Can we make the return type of this just be `RDD[Seq[T]]`? It will be more general if we want to optimize the data structures later