scala - Apache Spark's RDD splitting according to the particular size

Question

I am trying to read strings from a text file, but I want to limit each line according to a particular size. For example;

Here is my representing the file.

aaaaa bbb ccccc

When trying to read this file by sc.textFile, RDD would appear this one.

scala> val rdd = sc.textFile("textFile")
scala> rdd.collect
res1: Array[String] = Array(aaaaa, bbb, ccccc)

But I want to limit the size of this RDD. For example, if the limit is 3, then I should get like this one.

Array[String] = Array(aaa, aab, bbc, ccc, c)

What is the best performance way to do that?

Answer 1

Not a particularly efficient solution (not terrible either) but you can do something like this:

val pairs = rdd
  .flatMap(x => x)  // Flatten
  .zipWithIndex  // Add indices
  .keyBy(_._2 / 3)  // Key by index / n

// We'll use a range partitioner to minimize the shuffle 
val partitioner = new RangePartitioner(pairs.partitions.size, pairs)

pairs
  .groupByKey(partitioner)  // group
  // Sort, drop index, concat
  .mapValues(_.toSeq.sortBy(_._2).map(_._1).mkString("")) 
  .sortByKey()
  .values

It is possible to avoid the shuffle by passing data required to fill the partitions explicitly but it takes some effort to code. See my answer to Partition RDD into tuples of length n.

If you can accept some misaligned records on partitions boundaries then simple mapPartitions with grouped should do the trick at much lower cost:

rdd.mapPartitions(_.flatMap(x => x).grouped(3).map(_.mkString("")))

It is also possible to use sliding RDD:

rdd.flatMap(x => x).sliding(3, 3).map(_.mkString(""))