Batching within an Apache Spark RDD map - classification

Question

I have a situation where an underlying function operates significantly more efficiently when given batches to work on. I have existing code like this:
// subjects: RDD[Subject]
val subjects = Subject.load(job, sparkContext, config)
val classifications = subjects.flatMap(subject => classify(subject)).reduceByKey(_ + _)
classifications.saveAsTextFile(config.output)
The classify method works on single elements but would be more efficient operating on groups of elements. I considered using coalesce to split the RDD into chunks and acting on each chunk as a group, however there are two problems with this:
I'm not sure how to return the mapped RDD.
classify doesn't know in advance how big the groups should be and it varies based on the contents of the input.
Sample code on how classify could be called in an ideal situation (the output is kludgey since it can't spill for very large inputs):
def classifyRdd (subjects: RDD[Subject]): RDD[(String, Long)] = {
val classifier = new Classifier
subjects.foreach(subject => classifier.classifyInBatches(subject))
classifier.classifyRemaining
classifier.results
}
This way classifyInBatches can have code like this internally:
def classifyInBatches(subject: Subject) {
if (!internals.canAdd(subject)) {
partialResults.add(internals.processExisting)
}
internals.add(subject) // Assumption: at least one will fit.
}
What can I do in Apache Spark that will allow behavior somewhat like this?

Answer 1

Try using the mapPartitions method, which allows your map function to consume a partition as an iterator and produce an iterator of output.
You should be able to write something like this:
subjectsRDD.mapPartitions { subjects =>
val classifier = new Classifier
subjects.foreach(subject => classifier.classifyInBatches(subject))
classifier.classifyRemaining
classifier.results
}