From 3a731a903e5eb3431e8e23cbba86f6658aab6322 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 4 Oct 2026 14:20:06 -0600 Subject: [PATCH] fix: have native shuffle read a converted child's Arrow stream A native shuffle over a Spark-to-Arrow conversion, such as CometSparkToColumnarExec or CometLocalTableScanExec, wrapped the conversion's executeColumnar batches in a ColumnarBatchArrowReader, which closes each batch once native has it. The conversion writes every batch into the same vectors, and closing a struct vector drops its children, so the next batch failed to import. Native shuffle now reads such a child's Arrow stream, as a native operator does. --- .../shuffle/CometShuffleExchangeExec.scala | 92 ++++++++++++++----- .../comet/exec/CometNativeShuffleSuite.scala | 29 +++++- 2 files changed, 95 insertions(+), 26 deletions(-) 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 7a0d231a141..a992f090f72 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 @@ -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 @@ -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, CometSinkPlaceHolder, 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, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin} @@ -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 @@ -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) @@ -831,6 +854,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 => @@ -841,29 +894,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, @@ -876,7 +920,7 @@ object CometShuffleExchangeExec // is `shuffleWriterMetrics` at the root with one empty leaf for the Scan child. prepareNativeShuffleDependency( thinRDD, - Some(rdd), + samplingRDD, outputAttributes, outputPartitioning, serializer, 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 b817ed7a503..0c56e11b9ab 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -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} @@ -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