parthchandra opened a new pull request, #5392:
URL: https://github.com/apache/datafusion-comet/pull/5392
## Which issue does this PR close?
Closes #5391.
## Rationale for this change
A native-shuffle job with a very large number of partitions fails at stage
submission, before any task runs:
(From the issue:
Task serialization failed: java.lang.OutOfMemoryError: Required array
length 2147483639 + 794 is too large
at java.io.ByteArrayOutputStream.ensureCapacity(...)
at org.apache.spark.scheduler.DAGScheduler.submitMissingTasks(...)
`DAGScheduler` serializes the `(RDD, ShuffleDependency)` pair into a
single broadcast byte array that must fit in ~2GB. On the native-shuffle path,
`CometShuffleDependency.nativeShuffleSpec` is a non-transient field holding a
`perPartitionByKey` map with one scan-plan-data blob (the partition's file
list) per map partition. The failing job had ~38.7M
partitions, so ~38.7M protobufs got baked into that one blob and blew the
limit — even though each task only reads its own slice. Plain Spark avoids this
by keeping per-partition file lists `@transient` and shipping them per task;
Comet's own `CometExecRDD` does the same. The native shuffle path was the
exception.
## What changes are included in this PR?
Mirror `CometExecRDD`: keep the map off the serialized dependency and give
each task only its slice.
- `CometNativeShuffleInputRDD` takes `perPartitionByKey` as a `@transient`
arg and, in `getPartitions` (driver-side), slices out each partition's entry
onto its `Partition` object. That slice flows to the writer, which injects it
instead of indexing the full map.
- `NativeExecContext.perPartitionByKey` is now `@transient` so no build
path can serialize the full map, and we also empty it when building the
dependency.
- `commonByKey` is unchanged — it's sized by scan count, not partition
count, and the writer still needs it.
## How are these changes tested?
- A unit test serializes the input RDD at 10 vs 10,000 partitions and
checks the size stays roughly constant, and that each partition carries only
its own slice.
- An end-to-end test runs a native shuffle over a multi-partition native
scan and compares results against Spark with an all-native plan, so a wrong
slice would fail it.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
