fix: Native shuffle fails with a 2GB task serialization OOM on jobs with many partitions - #5392
Draft
parthchandra wants to merge 2 commits into
Draft
fix: Native shuffle fails with a 2GB task serialization OOM on jobs with many partitions#5392parthchandra wants to merge 2 commits into
parthchandra wants to merge 2 commits into
Conversation
…ith very many partitions
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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(...)
DAGSchedulerserializes the(RDD, ShuffleDependency)pair into a single broadcast byte array that must fit in ~2GB. On the native-shuffle path,CometShuffleDependency.nativeShuffleSpecis a non-transient field holding aperPartitionByKeymap with one scan-plan-data blob (the partition's file list) per map partition. The failing job had ~38.7Mpartitions, 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
@transientand shipping them per task; Comet's ownCometExecRDDdoes 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.CometNativeShuffleInputRDDtakesperPartitionByKeyas a@transientarg and, ingetPartitions(driver-side), slices out each partition's entry onto itsPartitionobject. That slice flows to the writer, which injects it instead of indexing the full map.NativeExecContext.perPartitionByKeyis now@transientso no build path can serialize the full map, and we also empty it when building the dependency.commonByKeyis unchanged — it's sized by scan count, not partition count, and the writer still needs it.How are these changes tested?