External Group By to handle huge keys

[mccheah]

This takes the Python implementation of group-by-key and brings Scala up to parity, making Scala's group-by-key fault-tolerant to a single large group. It does so via using an ExternalList data structure to combine the groups, where ExternalList can spill if it gets too large.

First, the performance testing:

I tried a few versions of this. In my first implementation, I wrote the ExternalList class, and simply replaced the CompactBuffer usage in PairRDDFunctions.groupByKey with external lists. I then ran the ExternalListSuite that did the group by key operation. With 7200000 unique integers, bucketed into 5 buckets, the trials yielded:

89.333 ms
89.992 ms
104.026 ms

I then switched to the current implementation. It matches exactly what Python does with an ExternalSorter. So in essence, this is a sort-based group by. It wasn't clear to me why sort-based group by is better... until I saw the numbers, for the same RDD and buckets:

49632 ms
53615 ms
54340 ms

Therefore I went with the current implementation.

Some caveats that I am concerned about:

1. ExternalList never cleans up the files it creates for spilling. This is noticeably different from ExternalAppendOnlyMap, which expects to only be iterated over once and then the map is discarded. We can't make that assumption for ExternalList, since the user could wish to iterate over the list a few times in any given RDD operation, meaning we have to keep the files around. I think that Spark itself will clean up its shuffle directories so we should be ok with not leaking files, but I'm still a bit paranoid.
2. Serialization is a bit hacky in ExternalList - it took a bit of work to get ExternalList to be both serializable and yet still be able to be spillable. For example, I use an extra companion object to hold references to "constants" that I don't want to be wiped out upon serialization. Due to the way Java serialization works, the code blocks in a Scala class will not be invoked again, leaving all vals as null. The time where the val makes sense to be instantiated only upon creating the class, I either have to make it a var and instantiate it in the serialization process, or mark the val as lazy so it is re-initialized upon use post-serialization.
3. ExternalList's spilling behavior is tied to shuffle parameters, which is a bit unintuitive. This is what Python does as well. I guess the alternative would be other Spark configurations, or parameters to group-by-key that have reasonable defaults. Pretty ugly either way.

cc @mingyukim , @zls88 , and @muyouzhi for review.

[JoshRosen]

ExternalList never cleans up the files it creates for spilling. This is noticeably different from ExternalAppendOnlyMap, which expects to only be iterated over once and then the map is discarded.

You might be able to use TaskContext's task completion listener mechanism to automatically perform cleanup at the end of a stage.

ExternalList's spilling behavior is tied to shuffle parameters, which is a bit unintuitive. This is what Python does as well.

It's quite likely that Tungsten will end up unifying / cleaning up our memory management interfaces such that we no longer have this weird scenario of anything that's spillable going through a ShuffleMemoryManager even in non-shuffle contexts.

[mccheah]

@JoshRosen , what will happen if the RDD with the external list is cached and then the task completion listener decides to clean it? Actually I didn't think about caching nearly enough but I think these ExternalList objects as it currently stands should be able to be cached as the backing files are still around.

[mccheah]

After having done some more reading on lazy val, it's become more clear that this should really be a var and instantiated on serialization because lazy val is definitely not "free"

[JoshRosen]

@mccheah Ah, right: caching definitely is a concern here. Creating the files within the shuffle spill directory should be sufficient to ensure that they're cleaned up once the application exits, but I think the concern with file leaks is that super-long-lived SparkContexts that perform tons and tons of queries might run out of disk space. I guess that we can't safely delete these files until the grouped RDD becomes unreachable and is garbage collected. Maybe there's a way to have the ContextCleaner handle cleanup in those cases, similar to how it cleans up shuffle files once ShuffleDependencies are garbage-collected.

[mccheah]

I think Python has cleanup logic in the list's "destructor" because I believe destructors are a thing in Python. Of course, no analogous thing exists in Scala / Java.

I can't think of a great way to make ContextCleaner be able to handle this situation. The hard part is getting ExternalList to hook itself into ContextCleaner and for the ExternalList to associate itself with a cached RDD. And also if the RDD isn't cached then the cleanup should happen on the task completion listener side. The interaction is a bit complex.

[mccheah]

I also thought a bit more about if this could be ignored for a first iteration but it's a pretty serious problem. If the list is spilled to disk, that means it couldn't fit in an executor's RAM, which means the list will probably be pretty big. So I could see the disk space being a legitimate concern.

[mccheah]

@JoshRosen the problem with using ContextCleaner is that it's a Spark driver side construct, right? However it is the executors that know explicitly when the ExternalLists are being constructed and spilled. It would seem more intuitive to tie the cleanup to the Executor side.

It might be necessary, but a little hacky, to create an equivalent executor-side context cleaner. It's either going to be that or the executors informing the driver to record all the external lists that are written, which is even less than ideal.

I'm thinking of creating something generic like "ExecutorCleaner" and put it in the SparkEnv for the executor side and start it with an API to clean up ExternalLists. What do you think?

[mccheah]

Ok I created an ExecutorCleaner class (not the best named class...). ExternalList objects register themselves for cleanup, but only after the task finishes because the task needs to build up the entire ExternalList first before registering the files that are to be purged. Probably still ugly parts of the code style-wise but feel free to look over the logic.

[justinuang]

@mccheah, where is the code for the python destructor for the list? From a quick scan of the python code, it looks like they just create an iterable that returns a generator iterator that reads from the spill dirs and removes them right afterwards. I'm probably missing something, because this implies that the iterable actually can only be iterated over once.

[justinuang]

Generally this is frowned upon, because foreach could possibly implemented in a parallel way, leading to concurrency issues if c1 isn't thread safe. I think in this case, we should be fine because c2 is not a parallel collection but just flagging the pattern.

I think you can do this instead

c1 ++= c2

http://www.scala-lang.org/api/current/index.html#scala.collection.generic.Growable@++=(xs:scala.collection.TraversableOnce[A]):Growable.this.type

[mccheah]

I have to define ++= on ExternalList to do this. Which is doable, just it wasn't done.

[justinuang]

Looks good overall, I'm mainly worried about the cleanup logic, lets discuss after you show me where the python cleanup logic is?

[mccheah]

k I'm going to open the PR against master now.