[CELEBORN-2480] Push columnar shuffle batches by estimated size in bytes

[ForVic]

What changes were proposed in this pull request?

Add an optional byte limit to the columnar shuffle writer. When celeborn.columnarShuffle.batch.sizeInBytes.enabled is true, ColumnarHashBasedShuffleWriter pushes a partition's batch once it reaches either celeborn.columnarShuffle.batch.size rows or celeborn.columnarShuffle.batch.sizeInBytes.threshold estimated bytes, whichever comes first. The estimate is the sum of UnsafeRow.getSizeInBytes() over the rows in the batch, so the column builders are unchanged.

Config	Default
celeborn.columnarShuffle.batch.sizeInBytes.enabled	false
celeborn.columnarShuffle.batch.sizeInBytes.threshold	64k, same as celeborn.client.push.buffer.max.size

The same change is made to the spark-3, spark-3.5 and spark-4 columnar writers.

Why are the changes needed?

The columnar writer keeps one in-memory batch per reduce partition and pushes it only after 10,000 rows. Row count is a poor proxy for memory: 10,000 rows can be a few hundred KB or hundreds of MB, depending on string widths. Spark's memory manager does not track this memory, so it cannot spill.

If a map task has fewer than 10,000 rows per partition, nothing is pushed until the task ends. That is common: 1M rows over 200 partitions is 5,000 rows each. The task's whole output then stays on the heap. The row-based hash writer instead caps each partition's buffer at celeborn.client.push.buffer.max.size, which defaults to 64 KiB.

Example: repartition 4M rows with a 2 KiB string into 200 partitions, using 4 map tasks:

spark.range(0, 4000000, 1, 4)
  .selectExpr("id", "cast(id % 1000 as int) as category",
    "repeat(md5(cast(id as string)), 64) as payload")
  .repartition(200, $"id")
  .write.format("noop").mode("overwrite").save()

Smallest driver heap that completes the job (tried 1g, 2g, 4g, 8g, 16g):

Writer	Min heap
Row-based hash writer	1g
Columnar, row limit only (current behavior)	OOM even at 16g
Columnar + byte limit, 64k (default)	1g
Columnar + byte limit, 256k	2g
Columnar + byte limit, 1m	4g
Columnar + byte limit, 4m	8g

As with the row-based writer, buffered row data is roughly partitions x threshold. This is not a hard memory cap: as before this change, the column builders pre-size their buffers for celeborn.columnarShuffle.batch.size rows (up to 80 KB per column per partition). Sizing those buffers from the threshold could be a follow-up.

Same job at a 24g heap, where every configuration fits (mean of 2 runs):

Columnar	Wall	Task time	Shuffle bytes
Row limit only	12.3s	40.7s	213.9 MB
64k	9.8s	31.4s	230.7 MB (+7.9%)
256k	8.0s	25.3s	221.8 MB (+3.7%)
1m	8.3s	26.3s	218.3 MB (+2.1%)

Smaller batches compress slightly worse, but they avoid most of the GC work. Narrow rows are barely affected: 40M rows of 64 bytes, 50,000 rows per partition per task, 4g heap:

Columnar	Wall	Shuffle bytes
Row limit only	11.4s	1655.9 MB
64k	11.5s	1594.7 MB
256k	10.8s	1663.0 MB
1m	11.0s	1655.9 MB (byte limit never reached)

The default of 64k matches the row-based writer's per-partition buffer size. A larger threshold trades memory for slightly smaller shuffle output.

Does this PR resolve a correctness bug?

• Yes

Does this PR introduce any user-facing change?

• Yes

Adds the two configs above. Behavior is unchanged unless celeborn.columnarShuffle.batch.sizeInBytes.enabled is set.

How was this patch tested?

• Added ColumnarHashBasedShuffleWriterSuiteJ#testPushBatchWhenSizeInBytesThresholdReached to the spark-3, spark-3.5 and spark-4 columnar modules. It checks the rows in each pushed batch, with the byte limit (8, 8, 8, 6) and without it (30).
• Ran build/mvn -P<profile> -pl client-spark/<module> -am test -Dtest=ColumnarHashBasedShuffleWriterSuiteJ -DwildcardSuites=none for spark-3.3, spark-3.5 and spark-4.0. All 14 tests pass in each. Maven CI doesn't run the columnar modules, so these were run locally.
• Regenerated docs/configuration/columnar-shuffle.md with ConfigurationSuite.
• Ran the jobs above on Spark 3.5.8 in local[4] with one local worker, JDK 17, on a 16 vCPU arm64 host. Spark had assets/spark-patch/Celeborn_Columnar_Shuffle_spark3.patch applied. On 3.5 the patch needs DataTypeUtils.fromAttributes in place of StructType.fromAttributes to compile. Timings are from a shared machine, so treat small differences as noise.

[ForVic]

maybe @SteNicholas or @RexXiong for review when you get a chance