Skip to content

fix: Native shuffle fails with a 2GB task serialization OOM on jobs with many partitions - #5392

Draft
parthchandra wants to merge 2 commits into
apache:mainfrom
parthchandra:shuffle-oom
Draft

fix: Native shuffle fails with a 2GB task serialization OOM on jobs with many partitions#5392
parthchandra wants to merge 2 commits into
apache:mainfrom
parthchandra:shuffle-oom

Conversation

@parthchandra

Copy link
Copy Markdown
Contributor

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Native shuffle fails with a 2GB task serialization OOM on jobs with very many partitions

1 participant