This is an automated email from the ASF dual-hosted git repository.
parthchandra pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 7e0e5d2a26 fix: Native shuffle fails with a 2GB task serialization OOM
on jobs with many partitions (#5392)
7e0e5d2a26 is described below
commit 7e0e5d2a26928b05b936c3f2a9abd375d6efc41f
Author: Parth Chandra <[email protected]>
AuthorDate: Thu Aug 20 09:30:11 2026 -0700
fix: Native shuffle fails with a 2GB task serialization OOM on jobs with
many partitions (#5392)
* fix: Native shuffle fails with a 2GB task serialization OOM on jobs with
very many partitions
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
.../shuffle/CometNativeShuffleInputRDD.scala | 31 +++++--
.../shuffle/CometNativeShuffleWriter.scala | 12 ++-
.../shuffle/CometShuffleExchangeExec.scala | 16 +++-
.../org/apache/spark/sql/comet/operators.scala | 7 +-
.../comet/exec/CometNativeShuffleSuite.scala | 28 ++++++
.../shuffle/CometNativeShuffleInputRDDSuite.scala | 101 +++++++++++++++++++++
8 files changed, 181 insertions(+), 16 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml
b/.github/workflows/pr_build_linux.yml
index dff1ec58fd..aaa9f0a255 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -327,6 +327,7 @@ jobs:
org.apache.comet.exec.CometShuffle4_0Suite
org.apache.comet.exec.CometNativeColumnarToRowSuite
org.apache.comet.exec.CometNativeShuffleSuite
+
org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
org.apache.comet.exec.CometShuffleEncryptionSuite
org.apache.comet.exec.CometShuffleManagerSuite
org.apache.comet.exec.CometAsyncShuffleSuite
diff --git a/.github/workflows/pr_build_macos.yml
b/.github/workflows/pr_build_macos.yml
index aca9da1f26..c6afabec31 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -143,6 +143,7 @@ jobs:
org.apache.comet.exec.CometShuffle4_0Suite
org.apache.comet.exec.CometNativeColumnarToRowSuite
org.apache.comet.exec.CometNativeShuffleSuite
+
org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
org.apache.comet.exec.CometShuffleEncryptionSuite
org.apache.comet.exec.CometShuffleManagerSuite
org.apache.comet.exec.CometAsyncShuffleSuite
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
index 0579e57bce..b79fb9458c 100644
---
a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
+++
b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
@@ -38,7 +38,8 @@ private[shuffle] class CometNativeShuffleInputRDD(
sc: SparkContext,
var inputRDDs: Seq[RDD[_]],
numPartitionsParam: Int,
- shuffleScanIndices: Set[Int])
+ shuffleScanIndices: Set[Int],
+ @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty)
extends RDD[Product2[Int, ColumnarBatch]](
sc,
inputRDDs.map(rdd => new OneToOneDependency(rdd))) {
@@ -50,7 +51,13 @@ private[shuffle] class CometNativeShuffleInputRDD(
// `leafRdd.partitions` on the executor, which would otherwise trigger
getPartitions and
// hit the @transient-null trap (e.g. CometExecRDD.perPartitionByKey).
val inputParts = inputRDDs.map(_.partitions(i)).toArray
- new CometNativeShuffleInputPartition(i, inputParts)
+ // Slice this partition's plan data off the @transient full map here on
the driver. Carrying
+ // only the per-partition slice on the Partition object (serialized per
task) keeps the full
+ // O(numPartitions) map out of the broadcast task binary, which
otherwise blows the 2GB
+ // ByteArrayOutputStream limit on jobs with tens of millions of
partitions. Mirrors
+ // CometExecRDD.getPartitions.
+ val planDataByKey = perPartitionByKey.map { case (key, arr) => key ->
arr(i) }
+ new CometNativeShuffleInputPartition(i, inputParts, planDataByKey)
}.toArray
override def compute(
@@ -63,7 +70,11 @@ private[shuffle] class CometNativeShuffleInputRDD(
partition.inputPartitions,
shuffleScanIndices,
context)
- new CometNativeShuffleInputIterator(partition.index, inputObjects,
shuffleBlockIters)
+ new CometNativeShuffleInputIterator(
+ partition.index,
+ inputObjects,
+ shuffleBlockIters,
+ partition.planDataByKey)
}
override def getPreferredLocations(split: Partition): Seq[String] = {
@@ -84,19 +95,23 @@ private[shuffle] class CometNativeShuffleInputRDD(
private[shuffle] class CometNativeShuffleInputPartition(
override val index: Int,
- val inputPartitions: Array[Partition])
+ val inputPartitions: Array[Partition],
+ val planDataByKey: Map[String, Array[Byte]])
extends Partition
/**
* Iterator handed to [[CometNativeShuffleWriter.write]] via Spark's
ShuffleMapTask. Reports no
- * elements; the writer downcasts and reads `partitionIndex`, `inputObjects`,
and
- * `shuffleBlockIterators` directly to drive the unified native plan.
`inputObjects` are the
- * already-resolved native input slots (see
[[CometExecRDD.resolveInputObjects]]).
+ * elements; the writer downcasts and reads `partitionIndex`, `inputObjects`,
+ * `shuffleBlockIterators`, and `planDataByKey` directly to drive the unified
native plan.
+ * `inputObjects` are the already-resolved native input slots (see
+ * [[CometExecRDD.resolveInputObjects]]). `planDataByKey` is this partition's
slice of the scan
+ * plan data (one entry per scan key); the writer injects it into the native
plan.
*/
private[shuffle] class CometNativeShuffleInputIterator(
val partitionIndex: Int,
val inputObjects: Array[Object],
- val shuffleBlockIterators: Map[Int, CometShuffleBlockIterator])
+ val shuffleBlockIterators: Map[Int, CometShuffleBlockIterator],
+ val planDataByKey: Map[String, Array[Byte]])
extends Iterator[Product2[Int, ColumnarBatch]] {
override def hasNext: Boolean = false
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
index 9eca4681b9..3245cc3429 100644
---
a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
+++
b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
@@ -97,10 +97,14 @@ class CometNativeShuffleWriter[K, V](
val unifiedPlan = buildUnifiedPlan(tempDataFilename, tempIndexFilename)
val ctx = spec.execContext
val finalNativePlan = if (ctx.commonByKey.nonEmpty) {
- val partitionDataByKey = ctx.perPartitionByKey.map { case (k, arr) =>
- k -> arr(partitionIdx)
- }
- PlanDataInjector.injectPlanData(unifiedPlan, ctx.commonByKey,
partitionDataByKey)
+ // This partition's plan-data slice rides on the input iterator's
Partition object (populated
+ // in CometNativeShuffleInputRDD.getPartitions on the driver), not on
the spec. The spec's
+ // execContext.perPartitionByKey is emptied in
prepareNativeShuffleDependency so the full
+ // O(numPartitions) map stays out of the broadcast task binary.
+ PlanDataInjector.injectPlanData(
+ unifiedPlan,
+ ctx.commonByKey,
+ shuffleInputIter.planDataByKey)
} else {
unifiedPlan
}
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index 9505443fb3..84f313ad37 100644
---
a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++
b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -122,7 +122,8 @@ case class CometShuffleExchangeExec(
sparkContext,
ctx.inputs,
ctx.numPartitions,
- ctx.shuffleScanIndices)
+ ctx.shuffleScanIndices,
+ ctx.perPartitionByKey)
case None =>
// Non-native child (e.g. CometSparkToColumnarExec): no subtree to
inline. The dep gets
// built via the convenience overload below; we just need a real RDD
of batches.
@@ -786,8 +787,17 @@ object CometShuffleExchangeExec
case e: Expression => e.collect { case s: ScalarSubquery => s }
case _ => Nil
}
- val augmentedSpec = spec.copy(execContext =
- spec.execContext.copy(subqueries = spec.execContext.subqueries ++
partitioningSubqueries))
+ // Drop the per-partition plan-data map off the spec that lands on the
(non-transient)
+ // CometShuffleDependency.nativeShuffleSpec. Each partition's slice now
rides on the thin RDD's
+ // Partition objects (see CometNativeShuffleInputRDD.getPartitions), so
the full
+ // O(numPartitions) map is dead weight here and would blow the 2GB
ByteArrayOutputStream limit
+ // at stage submission on very-high-partition-count jobs.
NativeExecContext.perPartitionByKey is
+ // also @transient (the structural guard against any build path), but we
empty it explicitly
+ // here too so the map isn't retained on the driver via this dependency.
commonByKey stays
+ // (O(#scans), not O(#partitions), and the writer still needs it).
+ val augmentedSpec = spec.copy(execContext = spec.execContext.copy(
+ subqueries = spec.execContext.subqueries ++ partitioningSubqueries,
+ perPartitionByKey = Map.empty))
// The code block below is mostly brought over from
// ShuffleExchangeExec::prepareShuffleDependency
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index 025f68b111..d91ad02a6c 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
@@ -539,7 +539,12 @@ private[comet] case class NativeExecContext(
broadcastedHadoopConfForEncryption:
Option[Broadcast[SerializableConfiguration]],
encryptedFilePaths: Seq[String],
commonByKey: Map[String, Array[Byte]],
- perPartitionByKey: Map[String, Array[Array[Byte]]],
+ // @transient: this holds one serialized scan-plan-data blob per
partition, so at high partition
+ // counts it is huge. It is only read on the driver (to slice per
partition onto each task's
+ // Partition object - see CometNativeShuffleInputRDD / CometExecRDD); the
executor reads its own
+ // slice, never this map. Keeping it off the wire stops it from bloating
the broadcast task
+ // binary when this context rides on the non-transient
CometShuffleDependency.nativeShuffleSpec.
+ @transient perPartitionByKey: Map[String, Array[Array[Byte]]],
shuffleScanIndices: Set[Int],
hasScanInput: Boolean) {
// Catch shape divergence (e.g. broadcast scans with different partition
counts after DPP
diff --git
a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
index ced6531cd8..cd52f57955 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
@@ -108,6 +108,34 @@ class CometNativeShuffleSuite extends CometTestBase with
AdaptiveSparkPlanHelper
}
}
+ test("native shuffle over a multi-partition native scan re-threads
per-partition plan data") {
+ // End-to-end companion to CometNativeShuffleInputRDDSuite: that suite
proves the per-partition
+ // scan plan data no longer rides the broadcast task binary; this one
proves each task still
+ // gets its OWN slice at write time. Reading real Parquet files gives a
CometNativeScanExec with
+ // several map partitions, so perPartitionByKey holds one file-list slice
per partition and a
+ // correct result depends on task i seeing slice i (not partition 0's).
+ withTempDir { dir =>
+ val path = new Path(dir.toURI.toString, "multi.parquet")
+ // Spread rows across 8 files so the native scan can yield multiple map
partitions.
+ spark
+ .range(0, 10000, 1, numPartitions = 8)
+ .selectExpr("id AS _1", "CAST(id AS STRING) AS _2")
+ .write
+ .parquet(path.toString)
+
+ // Force one scan partition per file split so the per-partition array
has multiple distinct
+ // entries; otherwise Spark coalesces the tiny files into a single
partition.
+ withSQLConf(
+ "spark.sql.files.maxPartitionBytes" -> "1024",
+ "spark.sql.files.openCostInBytes" -> "0") {
+ readParquetFile(path.toString) { df =>
+ val shuffled = df.repartition(17, col("_1"))
+ checkShuffleAnswer(shuffled, 1, checkNativeOperators = true)
+ }
+ }
+ }
+ }
+
test("hash-based native shuffle") {
withParquetTable((0 until 5).map(i => (i, (i + 1).toLong)), "tbl") {
val df = sql("SELECT * FROM tbl").sortWithinPartitions($"_1".desc)
diff --git
a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
new file mode 100644
index 0000000000..55194e80a2
--- /dev/null
+++
b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.execution.shuffle
+
+import org.apache.spark.HashPartitioner
+import org.apache.spark.serializer.JavaSerializer
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext}
+import org.apache.spark.sql.execution.metric.SQLMetrics
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import org.apache.comet.serde.OperatorOuterClass.Operator
+
+/**
+ * Ensure the serialized shuffle-map-stage task binary does not grow with the
number of
+ * partitions.
+ *
+ * `DAGScheduler.submitMissingTasks` broadcasts the serialized `(stage.rdd,
stage.shuffleDep)`
+ * pair, which must fit in ~2GB. On the native-shuffle path the scan plan data
is one serialized
+ * blob per map partition; the leak lived under
+ * `CometShuffleDependency.nativeShuffleSpec.execContext`, not on the thin
RDD. So this builds
+ * both carriers and serializes the pair, which catches a regression in either
the RDD or the
+ * dependency guards without needing a real 2GB allocation.
+ *
+ * Lives in the `execution.shuffle` package so it can construct the
`private[shuffle]`
+ * [[CometNativeShuffleInputRDD]] and the `private[comet]`
[[NativeExecContext]] directly.
+ */
+class CometNativeShuffleInputRDDSuite extends CometTestBase {
+
+ test("serialized (rdd, dep) task binary size is independent of partition
count") {
+ val sc = spark.sparkContext
+ val ser = new JavaSerializer(sc.getConf).newInstance()
+
+ // One scan key with a 1KB blob per map partition -- the per-partition
plan data. Build both
+ // objects the DAGScheduler serializes for a shuffle-map stage: the thin
RDD, and the
+ // CometShuffleDependency whose NativeShuffleSpec holds a
NativeExecContext with this map.
+ def build(numPartitions: Int): (
+ CometNativeShuffleInputRDD,
+ CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]) = {
+ val perPartitionByKey =
+ Map("scan-0" -> Array.fill(numPartitions)(new Array[Byte](1024)))
+ val rdd = new CometNativeShuffleInputRDD(
+ sc,
+ inputRDDs = Seq.empty,
+ numPartitionsParam = numPartitions,
+ shuffleScanIndices = Set.empty,
+ perPartitionByKey = perPartitionByKey)
+ val execContext = NativeExecContext(
+ inputs = Seq.empty,
+ numPartitions = numPartitions,
+ subqueries = Seq.empty,
+ broadcastedHadoopConfForEncryption = None,
+ encryptedFilePaths = Seq.empty,
+ commonByKey = Map.empty,
+ perPartitionByKey = perPartitionByKey,
+ shuffleScanIndices = Set.empty,
+ hasScanInput = false)
+ val spec =
+ NativeShuffleSpec(Operator.getDefaultInstance,
CometMetricNode(Map.empty), execContext)
+ val dep = new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch](
+ _rdd = rdd,
+ partitioner = new HashPartitioner(numPartitions),
+ decodeTime = SQLMetrics.createMetric(sc, "decode time"),
+ nativeShuffleSpec = Some(spec))
+ (rdd, dep)
+ }
+
+ // Pre-fix this pair grew from ~13KB to ~10MB between 10 and 10000
partitions. With the map held
+ // @transient on both the RDD and the NativeExecContext, the pair stays
roughly constant.
+ val (_, smallDep) = build(10)
+ val (largeRdd, largeDep) = build(10000)
+ val smallPair = ser.serialize((smallDep.rdd, smallDep)).limit()
+ val largePair = ser.serialize((largeDep.rdd, largeDep)).limit()
+ assert(
+ math.abs(largePair - smallPair) < 100 * 1024,
+ s"serialized (rdd, dep) grew with partition count (small=$smallPair,
large=$largePair); " +
+ "the per-partition plan-data map is leaking into the broadcast task
binary")
+
+ // Each task's Partition object still carries its own slice so the writer
can inject plan data.
+ val part =
largeRdd.partitions(7).asInstanceOf[CometNativeShuffleInputPartition]
+ assert(part.planDataByKey.keySet == Set("scan-0"))
+ assert(part.planDataByKey("scan-0").length == 1024)
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
