Add batch number to DataFrame based on moving sum in spark

If you are not looking for the optimal solution but for a way to solve the problem in Spark without too much complexity, we can divide the problem into two steps:

1. split the data in blocks with 5 rows each, ignoring the text length
2. if the sum of the text lenghts in one block is too large, split this block into further blocks

This solution is not optimal as it produces too many batches.

Step 1 can be implemented using zipWithIndex. When creating the batch ids, we leave enough "space" for dividing the batches later. At the end of this step all rows in one block are grouped into a single list as input for step 2:

df = ...

r = df.rdd.zipWithIndex().toDF() \
    .select("_1.id", "_1.text_length", "_2") \
    .withColumn("batch", F.expr("cast(_2 / 5 as long)*5")) \
    .withColumn("data", F.struct("id", "text_length", "batch")) \
    .groupBy("batch") \
    .agg(F.collect_list("data").alias("data"))

Part 2 consists mainly of an udf that checks if within one batch the maximum text length is exceeded. If so, the batch id for the following elements is increased by one. As we have skipped enough batch ids in part 1, we do not get any collisions.

def splitBatchIfNecessary(data):
    text_length = 0
    batch = -1
    for d in data:
        text_length = text_length + d.text_length
        if text_length > 1000:
          if batch == -1:
            text_length = 0
            batch = d.batch + 1
            yield (d.id, d.text_length, d.batch)
          else:
            text_length = d.text_length
            batch = batch + 1
            yield (d.id, d.text_length, batch)          
        else:
          if batch == -1:
            batch = d.batch
          yield (d.id, d.text_length, batch)

schema=r.schema["data"].dataType
split_udf = F.udf(splitBatchIfNecessary, schema)

r = r.withColumn("data",split_udf(F.col("data")) ) \
      .selectExpr("explode(data)") \
      .select("col.*") 

Output:

+---+-----------+-----+                                                         
| id|text_length|batch|
+---+-----------+-----+
|  1|        500|    0|
|  2|        400|    0|
|  3|        200|    1|
|  4|        300|    1|
|  5|        100|    1|
|  6|        100|    5|
|  7|        100|    5|
|  8|        100|    5|
|  9|        100|    5|
| 10|        300|    5|
+---+-----------+-----+

Possible optimizations would be to replace zipWithIndex by zipWithUniqueIds (but getting slightly more "incomplete" batches) or using a vectorized udf.

Unlike mck states, this is not the "partition problem".

What the issues are, are that 1) Spark works with partitions - not just 1 such partition in order to be effective and 2) there is not a grouping attribute to ensure that 'batches' can form naturally or be distilled naturally within a partition only. Moreover, can we have negative numbers or fractions? - this is not stated. Max 5 entries it is stated, however.

1. This means that processing would need to be just based on one partition, but it may not be big enough aka OOM.

2. Trying to process per partition is pointless as all work would need to be done per partition N, N+1, and so on die to offset effects in the partition N-1. I have worked out a solution here on SO that took partition boundaries into account, but this is against the principle of Spark and the use case was more simplistic.

3. Actually not a Spark use case. It is a sequential algorithm as opposed to parallel algorithm, use PL/SQL, Scala, JAVA, C++.

4. The only way would be:

   • looping over a fixed size partition that has had zipWithIndex applied globally (for safety)
     • process with Scala into batches - temp result
     • take all items from last created batch and union with next partition
     • remove last batch from temp results
     • repeat cycle

5. NB: Approximations to get around partitioning boundary aspects of data seem not to work --> the other answer proves that in fact. You get a compromise result, not the actually answer. And to correct it is not on fact that easy as batches have gaps and may be in other partitions as result of grouping.