ForVic opened a new pull request, #3873:
URL: https://github.com/apache/celeborn/pull/3873
### 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:
```scala
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 memory is roughly `partitions x
threshold`.
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 memory budget. 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?
- [x] 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.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
