Add batch number to DataFrame based on moving sum in spark
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.

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

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.

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

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




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.
I have a dataset that I need to process batchwise (due to API restrictions).
The sum of the column text_lenth of a batch cannot exceed 1000. And the maximum number of rows in a batch cannot be greater than 5.
For that I would like to add batch numbers to the single batches in order to process the data based on batch_numbers later.
How can I achieve that in pyspark (in Databricks). I am pretty new to all of this that I don't even know what to look for online.
I really appreciate your help.
The tables below illustrate what I am trying to achieve:
Original table




id
text_length




1
500


2
400


3
200


4
300


5
100


6
100


7
100


8
100


9
100


10
300




Resulting table




id
text_length
batch_number




1
500
1


2
400
1


3
200
2


4
300
2


5
100
2


6
100
2


7
100
2


8
100
3


9
100
3


10
300
3