Skip to content

fix: have native shuffle read a converted child's Arrow stream #6689

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
andygrove merged 2 commits into from
Oct 6, 2026
Merged

fix: have native shuffle read a converted child's Arrow stream #6689

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3a731a9
fix: have native shuffle read a converted child's Arrow stream
andygrove Oct 4, 2026
3652546
Merge apache/main into fix/native-shuffle-reads-arrow-stream
andygrove Oct 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 68 additions & 24 deletions ...rc/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import java.util.function.Supplier
import scala.concurrent.Future
import scala.jdk.CollectionConverters._

import org.apache.arrow.c.ArrowArrayStream
import org.apache.spark._
import org.apache.spark.internal.config
import org.apache.spark.rdd.RDD
Expand All @@ -35,7 +36,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.LazilyGeneratedOrdering
import org.apache.spark.sql.catalyst.plans.logical.Statistics
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.comet.{CometFilterExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometPlan, CometProjectExec, CometScanWrapper, CometSinkPlaceHolder, CometSparkToColumnarExec, NativeExecContext}
import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNativeArrowSource}
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec
import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, Exchange, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin}
Expand Down Expand Up @@ -144,9 +145,16 @@ case class CometShuffleExchangeExec(
ctx.perPartitionByKey,
positionalRoundRobin.isDefined)
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.
child.executeColumnar()
child match {
// Native reads the source's Arrow stream, as it does for a native operator's input. The
// batches from the source's `executeColumnar` share vectors that it reuses for the next
// batch, and the convenience overload below closes each batch once native has it.
// Closing a struct vector drops its children, so the next batch would lose them.
case source: CometNativeArrowSource => source.doExecuteAsArrowStream()
// Other non-native child: no subtree to inline. The dep gets built via the convenience
// overload below; we just need a real RDD of batches.
case _ => child.executeColumnar()
}
}
} else if (shuffleType == CometColumnarShuffle) {
// Row-based shuffle. CometNativeExec.doExecute wraps columnar output with
Expand Down Expand Up @@ -227,12 +235,27 @@ case class CometShuffleExchangeExec(
ctx,
positionalRoundRobin))
case None =>
CometShuffleExchangeExec.prepareShuffleDependency(
inputRDD.asInstanceOf[RDD[ColumnarBatch]],
child.output,
outputPartitioning,
serializer,
metrics)
child match {
case _: CometNativeArrowSource =>
CometShuffleExchangeExec.prepareArrowStreamShuffleDependency(
inputRDD.asInstanceOf[RDD[ArrowArrayStream]],
// The range partitioner samples rows, so it needs them as batches.
outputPartitioning match {
case _: RangePartitioning => Some(child.executeColumnar())
case _ => None
},
child.output,
outputPartitioning,
serializer,
metrics)
case _ =>
CometShuffleExchangeExec.prepareShuffleDependency(
inputRDD.asInstanceOf[RDD[ColumnarBatch]],
child.output,
outputPartitioning,
serializer,
metrics)
}
}
metrics("numPartitions").set(dep.partitioner.numPartitions.toLong)
val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
Expand Down Expand Up @@ -857,6 +880,36 @@ object CometShuffleExchangeExec
outputPartitioning: Partitioning,
serializer: Serializer,
metrics: Map[String, SQLMetric]): ShuffleDependency[Int, ColumnarBatch, ColumnarBatch] = {
// Wrap the raw batches as an RDD[ArrowArrayStream] so the leaf reaches native via the Arrow C
// Stream Interface, matching how CometNativeExec.buildNativeContext feeds the native-child
// path.
val streamRDD = CometArrowStream.wrapColumnarBatchRDD(
rdd,
StructType(
outputAttributes.map(a => StructField(a.name, a.dataType, a.nullable, a.metadata))),
CometArrowStream.NATIVE_TIMEZONE,
"ShuffleWriterInput")
prepareArrowStreamShuffleDependency(
streamRDD,
Some(rdd),
outputAttributes,
outputPartitioning,
serializer,
metrics)
}

/**
* [[prepareShuffleDependency]] for input that native reads as one Arrow stream per partition,
* such as a [[CometNativeArrowSource]]'s. `samplingRDD` gives the same rows as batches, and is
* only required for [[RangePartitioning]].
*/
private def prepareArrowStreamShuffleDependency(
streamRDD: RDD[ArrowArrayStream],
samplingRDD: Option[RDD[ColumnarBatch]],
outputAttributes: Seq[Attribute],
outputPartitioning: Partitioning,
serializer: Serializer,
metrics: Map[String, SQLMetric]): ShuffleDependency[Int, ColumnarBatch, ColumnarBatch] = {

val scanBuilder = OperatorOuterClass.Scan.newBuilder().setSource("ShuffleWriterInput")
val scanTypes = outputAttributes.flatMap { attr =>
Expand All @@ -867,29 +920,20 @@ object CometShuffleExchangeExec
s"$outputAttributes contains unsupported data types for CometShuffleExchangeExec.")
}
scanBuilder.addAllFields(scanTypes.asJava)
// The native consumer of the stream.
val scanOp = OperatorOuterClass.Operator.newBuilder().setScan(scanBuilder).build()

// Wrap the raw batches as an RDD[ArrowArrayStream] so the leaf reaches native via the Arrow C
// Stream Interface, matching how CometNativeExec.buildNativeContext feeds the native-child
// path. The synthetic Scan("ShuffleWriterInput") above is the native consumer.
val streamRDD = CometArrowStream.wrapColumnarBatchRDD(
rdd,
StructType(
outputAttributes.map(a => StructField(a.name, a.dataType, a.nullable, a.metadata))),
CometArrowStream.NATIVE_TIMEZONE,
"ShuffleWriterInput")

val childMetricNode = CometMetricNode(Map.empty)
val thinRDD = new CometNativeShuffleInputRDD(
rdd.sparkContext,
streamRDD.sparkContext,
Seq(streamRDD),
rdd.getNumPartitions,
streamRDD.getNumPartitions,
shuffleScanIndices = Set.empty,
spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode)))

val ctx = NativeExecContext(
inputs = Seq(streamRDD),
numPartitions = rdd.getNumPartitions,
numPartitions = streamRDD.getNumPartitions,
subqueries = Seq.empty,
broadcastedHadoopConfForEncryption = None,
encryptedFilePaths = Seq.empty,
Expand All @@ -902,7 +946,7 @@ object CometShuffleExchangeExec
// is `shuffleWriterMetrics` at the root with one empty leaf for the Scan child.
prepareNativeShuffleDependency(
thinRDD,
Some(rdd),
samplingRDD,
outputAttributes,
outputPartitioning,
serializer,
Expand Down
29 changes: 27 additions & 2 deletions spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,14 +42,14 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.Final
import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
import org.apache.spark.sql.comet.{CometExec, CometHashAggregateExec, CometLocalTableScanExec, CometMetricNode, CometNativeExec, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec}
import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNativeArrowSource}
import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec}
import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan}
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, ShuffleQueryStageExec}
import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
import org.apache.spark.sql.functions.{broadcast, col, count, countDistinct, sum}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ArrayType, DataType, LongType, MapType, StructField, StructType}
import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, LongType, MapType, StringType, StructField, StructType}
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}

import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native}
Expand Down Expand Up @@ -694,6 +694,31 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
}
}

test("native shuffle reads a struct column that a conversion writes over several batches") {
// A conversion of Spark rows writes every batch into the same vectors. Native shuffle reads
// its Arrow stream, as a native operator does, rather than batches that it would close once
// native has them: closing a struct vector drops its children before the next batch (#6685).
val schema = new StructType()
.add("k", IntegerType)
.add("payload", new StructType().add("v", LongType).add("s", StringType))
val data = (0 until 200).map(i => Row(i, Row(i.toLong, s"s$i")))
withSQLConf(
CometConf.COMET_BATCH_SIZE.key -> "7",
CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true",
CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") {
// A converted RDD scan, and a local table scan.
val rdd = spark.createDataFrame(spark.sparkContext.parallelize(data, 1), schema)
val local = spark.createDataFrame(data.asJava, schema)
for (df <- Seq(rdd, local);
shuffled <- Seq(df.repartition(3, col("k")), df.repartitionByRange(3, col("k")))) {
val (_, plan) = checkSparkAnswer(shuffled)
val shuffles = collect(plan) { case s: CometShuffleExchangeExec => s }
assert(shuffles.map(_.shuffleType) == Seq(CometNativeShuffle), plan)
assert(shuffles.head.child.isInstanceOf[CometNativeArrowSource], plan)
}
}
}

test("native shuffle declines a struct data column with duplicate field names") {
// Java Arrow keys a struct vector's children by name, so a struct with two same-named
// fields cannot be imported back across the C data interface after a native shuffle, and
Expand Down
Loading