From a9f061b22303552188f61b60a9576ad5b9cb5ced Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 28 Oct 2026 08:51:22 -0700 Subject: [PATCH 1/4] fix: apply the session's Comet configs in map stages of jobs outside a SQL execution On an executor, SQLConf.get reads the task's local properties, which Spark fills with the session's configs only for jobs that run inside a SQL execution. A job launched from a query's RDD (Dataset.rdd with AQE off, QueryExecution.toRdd) runs its map stages outside one, and Spark's SQLExecutionRDD re-applies the session's configs only in the final stage. Comet reads its settings on the executor when it builds a native plan, so those map stages ran with every spark.comet.* setting at its default, including the memory pool and the configs passed to native. Capture the session's configs on the driver where Comet builds the RDDs and native shuffle that run native plans, and make them what SQLConf.get returns on the task thread while the plan is built, as SQLExecutionRDD does. Tasks inside a SQL execution are unchanged. Closes #7301. --- .../spark/sql/comet/CapturedSQLConf.scala | 68 +++++++++++++++++++ .../apache/spark/sql/comet/CometExecRDD.scala | 9 ++- .../spark/sql/comet/CometExecUtils.scala | 17 +++-- .../CometTakeOrderedAndProjectExec.scala | 36 ++++++---- .../shuffle/CometNativeShuffleInputRDD.scala | 13 ++-- .../shuffle/CometNativeShuffleWriter.scala | 2 +- .../shuffle/CometShuffleDependency.scala | 9 ++- .../comet/exec/CometNativeShuffleSuite.scala | 26 +++++++ 8 files changed, 150 insertions(+), 30 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CapturedSQLConf.scala diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CapturedSQLConf.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CapturedSQLConf.scala new file mode 100644 index 00000000000..b994a8a60eb --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CapturedSQLConf.scala @@ -0,0 +1,68 @@ +/* + * 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 + +import org.apache.spark.TaskContext +import org.apache.spark.sql.execution.SQLExecution +import org.apache.spark.sql.internal.SQLConf + +/** + * The SQL configs of the session that planned a Comet plan, captured on the driver so that the + * plan's tasks read the same values however its job was launched. + * + * On an executor, `SQLConf.get` reads the task's local properties, which Spark fills with the + * session's configs only for a job that runs inside a SQL execution. A job launched from a + * query's RDD (`Dataset.rdd`, `QueryExecution.toRdd`) can run outside one, and Spark's + * `SQLExecutionRDD` re-applies the session's configs only in the stage that computes it, the last + * one. Without this, every earlier stage resolves each `spark.comet.*` setting to its default, + * including the memory pool and the configs passed to the native plan. + * + * Works like `SQLExecutionRDD`: outside a SQL execution, [[withConf]] makes the captured configs + * what `SQLConf.get` returns on the task thread while it runs. Inside one, the task's local + * properties already carry the session's configs, and nothing changes. + */ +private[spark] class CapturedSQLConf private (configs: Map[String, String]) extends Serializable { + + @transient private lazy val executorSideConf: SQLConf = { + val conf = new SQLConf() + configs.foreach { case (key, value) => conf.setConfString(key, value) } + conf + } + + /** + * Runs `f` with the captured configs as `SQLConf.get` on this thread, if it is a task thread + * outside a SQL execution. Only reads made while `f` runs see them, so `f` has to read the + * settings it needs, as the `CometExecIterator` constructor does. + */ + def withConf[T](f: => T): T = { + val context = TaskContext.get() + if (context == null || context.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) != null) { + f + } else { + SQLConf.withExistingConf(executorSideConf)(f) + } + } +} + +private[spark] object CapturedSQLConf { + + /** Captures the configs of the active session. Call it on the driver, while planning. */ + def capture(): CapturedSQLConf = new CapturedSQLConf(SQLConf.get.getAllConfs) +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index d20022f2b53..c255a4ee8a1 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -93,6 +93,10 @@ private[spark] class CometExecRDD( s"All per-partition arrays must have length $numPartitions, but found: " + perPartitionByKey.map { case (key, arr) => s"$key -> ${arr.length}" }.mkString(", ")) + // Captured here, on the driver, so that the native plan gets the session's settings in a job + // that runs outside a SQL execution. See CapturedSQLConf. + private val sqlConf = CapturedSQLConf.capture() + override protected def getPartitions: Array[Partition] = { (0 until numPartitions).map { idx => val inputParts = inputRDDs.map(_.partitions(idx)).toArray @@ -103,7 +107,10 @@ private[spark] class CometExecRDD( }.toArray } - override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { + override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = + sqlConf.withConf(computeNative(split, context)) + + private def computeNative(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { // Must precede resolveInputObjects and the CometExecIterator: completion listeners run in // reverse registration order, so registering first means this listener runs last, after // nested native blocks and the iterator have published their final metric values. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala index 5c97bebf271..4adc10c1bf2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala @@ -72,14 +72,17 @@ object CometExecUtils { val limitOp = CometExecUtils.getLimitNativePlan(outputAttribute, limit, offset).get val serializedPlan = CometExec.serializeNativePlan(limitOp) val inputSchema = Utils.fromAttributes(outputAttribute) + val sqlConf = CapturedSQLConf.capture() childPlan.mapPartitionsWithIndexInternal { case (idx, iter) => - CometExec.getCometIterator( - CometArrowStream.inputObjects(iter, inputSchema, "CometExecUtils-getNativeLimit"), - outputAttribute.length, - serializedPlan, - numParts, - idx, - outputAttribute.map(_.dataType)) + sqlConf.withConf { + CometExec.getCometIterator( + CometArrowStream.inputObjects(iter, inputSchema, "CometExecUtils-getNativeLimit"), + outputAttribute.length, + serializedPlan, + numParts, + idx, + outputAttribute.map(_.dataType)) + } } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala index af78e16e825..e713e6ff85f 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala @@ -152,6 +152,7 @@ case class CometTakeOrderedAndProjectExec( protected override def doExecuteColumnar(): RDD[ColumnarBatch] = { val childRDD = child.executeColumnar() + val sqlConf = CapturedSQLConf.capture() if (childRDD.getNumPartitions == 0 || limit == 0) { new ParallelCollectionRDD(sparkContext, Seq.empty[ColumnarBatch], 1, Map.empty) } else { @@ -172,13 +173,16 @@ case class CometTakeOrderedAndProjectExec( val numOutputCols = child.output.length val inputSchema = CometUtils.fromAttributes(child.output) childRDD.mapPartitionsWithIndexInternal { case (idx, iter) => - CometExec.getCometIterator( - CometArrowStream.inputObjects(iter, inputSchema, "CometTakeOrderedAndProject-topK"), - numOutputCols, - serializedTopK, - numParts, - idx, - child.output.map(_.dataType)) + sqlConf.withConf { + CometExec.getCometIterator( + CometArrowStream + .inputObjects(iter, inputSchema, "CometTakeOrderedAndProject-topK"), + numOutputCols, + serializedTopK, + numParts, + idx, + child.output.map(_.dataType)) + } } } @@ -200,14 +204,16 @@ case class CometTakeOrderedAndProjectExec( val finalOutputTypes = output.map(_.dataType) val finalInputSchema = CometUtils.fromAttributes(child.output) singlePartitionRDD.mapPartitionsInternal { iter => - val it = CometExec.getCometIterator( - CometArrowStream - .inputObjects(iter, finalInputSchema, "CometTakeOrderedAndProject-final"), - finalOutputTypes.length, - serializedFinalPlan, - 1, - 0, - finalOutputTypes) + val it = sqlConf.withConf { + CometExec.getCometIterator( + CometArrowStream + .inputObjects(iter, finalInputSchema, "CometTakeOrderedAndProject-final"), + finalOutputTypes.length, + serializedFinalPlan, + 1, + 0, + finalOutputTypes) + } setSubqueries(it.id, this) Option(TaskContext.get()).foreach { context => 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 8770496bff2..641fd36da64 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 @@ -21,7 +21,7 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark._ import org.apache.spark.rdd.{DeterministicLevel, RDD} -import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode} +import org.apache.spark.sql.comet.{CapturedSQLConf, CometExecRDD, CometMetricNode} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometShuffleBlockIterator @@ -37,6 +37,9 @@ import org.apache.comet.CometShuffleBlockIterator * @param positionalRoundRobin * whether the writer fed by this RDD places rows by position rather than by content; see * [[CometShuffleExchangeExec.positionalRoundRobinSpec]] and `getOutputDeterministicLevel`. + * @param sqlConf + * the session's configs, captured on the driver, for resolving the inputs in a job that runs + * outside a SQL execution. See [[CapturedSQLConf]]. */ private[shuffle] class CometNativeShuffleInputRDD( sc: SparkContext, @@ -45,7 +48,8 @@ private[shuffle] class CometNativeShuffleInputRDD( shuffleScanIndices: Set[Int], spillMetricNode: CometMetricNode, @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty, - positionalRoundRobin: Boolean = false) + positionalRoundRobin: Boolean = false, + sqlConf: CapturedSQLConf = CapturedSQLConf.capture()) extends RDD[Product2[Int, ColumnarBatch]]( sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { @@ -63,7 +67,8 @@ private[shuffle] class CometNativeShuffleInputRDD( shuffleScanIndices, spillMetricNode, perPartitionByKey, - positionalRoundRobin) + positionalRoundRobin, + sqlConf) /** * Spark's `isOrderSensitive` rule, applied to the RDD graph below the native plan. Spark only @@ -109,7 +114,7 @@ private[shuffle] class CometNativeShuffleInputRDD( override def compute( split: Partition, - context: TaskContext): Iterator[Product2[Int, ColumnarBatch]] = { + context: TaskContext): Iterator[Product2[Int, ColumnarBatch]] = sqlConf.withConf { spillMetricNode.reportSpillMetrics(context) val partition = split.asInstanceOf[CometNativeShuffleInputPartition] val (inputObjects, shuffleBlockIters) = 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 7aef724c495..cf8b6ffed9b 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 @@ -95,7 +95,7 @@ class CometNativeShuffleWriter[K, V]( override def write(inputs: Iterator[Product2[K, V]]): Unit = { try { - writeInternal(inputs) + spec.sqlConf.withConf(writeInternal(inputs)) } catch { case failure: Throwable => remoteDestination.foreach { destination => diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index be91caa0ec9..f50bda5a1ed 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala @@ -28,7 +28,7 @@ import org.apache.spark.shuffle.ShuffleWriteProcessor import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.physical.Partitioning -import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext} +import org.apache.spark.sql.comet.{CapturedSQLConf, CometMetricNode, NativeExecContext} import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types.StructType @@ -51,7 +51,12 @@ case class NativeShuffleSpec( * size so that a re-executed map task places rows exactly as the attempt it replaces. See * `CometShuffleExchangeExec.positionalRoundRobinSpec`. */ - positionalRoundRobin: Option[PositionalRoundRobin] = None) + positionalRoundRobin: Option[PositionalRoundRobin] = None, + /** + * The session's configs, captured on the driver, which the writer applies while it builds and + * runs the native plan in a job that runs outside a SQL execution. See [[CapturedSQLConf]]. + */ + sqlConf: CapturedSQLConf = CapturedSQLConf.capture()) /** * Parameters for positional round-robin placement, resolved on the driver. 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 06660cef2a6..1f30786a27d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -1983,4 +1983,30 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } } + + // https://github.com/apache/datafusion-comet/issues/7301 + test("map stage of a job outside a SQL execution uses the session's Comet configs") { + withTempPath { dir => + spark + .range(0, 1000, 1, 1) + .selectExpr("id % 10 as k", "id as v") + .write + .parquet(dir.toString) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_BATCH_SIZE.key -> "10") { + val df = spark.read.parquet(dir.toString).repartition(2, col("k")) + // A job launched from the query's RDD runs outside a SQL execution, so Spark does not + // pass the session's configs to the map stage's tasks. + df.queryExecution.toRdd.foreach(_ => ()) + val exchange = find(df.queryExecution.executedPlan) { + case _: CometShuffleExchangeExec => true + case _ => false + }.get + // The scan fused into the shuffle writer emits 10-row batches, not one of the default + // 8192 rows. + assert(exchange.metrics("input_batches").value == 100) + } + } + } } From 53b2216af6bcb313aab27b851b567417f64371f4 Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 28 Oct 2026 09:02:53 -0700 Subject: [PATCH 2/4] refactor: capture only Comet configs once per shuffle and let configs Spark passed win Capture only spark.comet.* keys, and apply them only when no conf is set on the thread and the task got none of the session's Comet configs from Spark, so SQLExecution.withSQLConfPropagated still takes effect. Carry the native shuffle's one captured conf from its input RDD to the writer, apply it inside the serialized-plan getCometIterator, and drop the test workarounds this change makes redundant. --- .../spark/sql/comet/CapturedSQLConf.scala | 39 ++++++++----------- .../apache/spark/sql/comet/CometExecRDD.scala | 2 - .../spark/sql/comet/CometExecUtils.scala | 17 ++++---- .../CometTakeOrderedAndProjectExec.scala | 39 +++++++++---------- .../shuffle/CometNativeShuffleInputRDD.scala | 9 +++-- .../shuffle/CometNativeShuffleWriter.scala | 29 +++++++------- .../shuffle/CometShuffleDependency.scala | 9 +---- .../apache/spark/sql/comet/operators.scala | 4 +- .../apache/comet/exec/CometJoinSuite.scala | 7 +--- .../comet/exec/CometNativeShuffleSuite.scala | 11 ++---- .../apache/comet/vector/NativeUtilSuite.scala | 6 ++- .../sql/comet/CometTaskMetricsSuite.scala | 26 ++++++------- 12 files changed, 89 insertions(+), 109 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CapturedSQLConf.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CapturedSQLConf.scala index b994a8a60eb..1d15fde5f12 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CapturedSQLConf.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CapturedSQLConf.scala @@ -20,25 +20,17 @@ package org.apache.spark.sql.comet import org.apache.spark.TaskContext -import org.apache.spark.sql.execution.SQLExecution -import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.{ReadOnlySQLConf, SQLConf} + +import org.apache.comet.CometConf /** - * The SQL configs of the session that planned a Comet plan, captured on the driver so that the - * plan's tasks read the same values however its job was launched. - * - * On an executor, `SQLConf.get` reads the task's local properties, which Spark fills with the - * session's configs only for a job that runs inside a SQL execution. A job launched from a - * query's RDD (`Dataset.rdd`, `QueryExecution.toRdd`) can run outside one, and Spark's - * `SQLExecutionRDD` re-applies the session's configs only in the stage that computes it, the last - * one. Without this, every earlier stage resolves each `spark.comet.*` setting to its default, - * including the memory pool and the configs passed to the native plan. - * - * Works like `SQLExecutionRDD`: outside a SQL execution, [[withConf]] makes the captured configs - * what `SQLConf.get` returns on the task thread while it runs. Inside one, the task's local - * properties already carry the session's configs, and nothing changes. + * The session's `spark.comet.*` configs, captured on the driver. Spark passes them to tasks as + * local properties only inside a SQL execution, so jobs launched from `Dataset.rdd` or + * `QueryExecution.toRdd` read Comet defaults in every stage before the last, which + * `SQLExecutionRDD` covers. [[withConf]] fills that gap the same way. */ -private[spark] class CapturedSQLConf private (configs: Map[String, String]) extends Serializable { +class CapturedSQLConf private (configs: Map[String, String]) extends Serializable { @transient private lazy val executorSideConf: SQLConf = { val conf = new SQLConf() @@ -47,13 +39,13 @@ private[spark] class CapturedSQLConf private (configs: Map[String, String]) exte } /** - * Runs `f` with the captured configs as `SQLConf.get` on this thread, if it is a task thread - * outside a SQL execution. Only reads made while `f` runs see them, so `f` has to read the - * settings it needs, as the `CometExecIterator` constructor does. + * Runs `f` with the captured configs as `SQLConf.get`, unless this is not a task thread, a conf + * is already set on it, or Spark passed the task the session's configs (it passes all or none). */ def withConf[T](f: => T): T = { val context = TaskContext.get() - if (context == null || context.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) != null) { + if (configs.isEmpty || context == null || !SQLConf.get.isInstanceOf[ReadOnlySQLConf] || + configs.keys.exists(context.getLocalProperty(_) != null)) { f } else { SQLConf.withExistingConf(executorSideConf)(f) @@ -61,8 +53,9 @@ private[spark] class CapturedSQLConf private (configs: Map[String, String]) exte } } -private[spark] object CapturedSQLConf { +object CapturedSQLConf { - /** Captures the configs of the active session. Call it on the driver, while planning. */ - def capture(): CapturedSQLConf = new CapturedSQLConf(SQLConf.get.getAllConfs) + /** Captures the active session's Comet configs. Call it on the driver. */ + def capture(): CapturedSQLConf = + new CapturedSQLConf(SQLConf.get.getAllConfs.filter(_._1.startsWith(CometConf.COMET_PREFIX))) } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index c255a4ee8a1..c903b7f70b0 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -93,8 +93,6 @@ private[spark] class CometExecRDD( s"All per-partition arrays must have length $numPartitions, but found: " + perPartitionByKey.map { case (key, arr) => s"$key -> ${arr.length}" }.mkString(", ")) - // Captured here, on the driver, so that the native plan gets the session's settings in a job - // that runs outside a SQL execution. See CapturedSQLConf. private val sqlConf = CapturedSQLConf.capture() override protected def getPartitions: Array[Partition] = { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala index 4adc10c1bf2..46fe131ece1 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala @@ -74,15 +74,14 @@ object CometExecUtils { val inputSchema = Utils.fromAttributes(outputAttribute) val sqlConf = CapturedSQLConf.capture() childPlan.mapPartitionsWithIndexInternal { case (idx, iter) => - sqlConf.withConf { - CometExec.getCometIterator( - CometArrowStream.inputObjects(iter, inputSchema, "CometExecUtils-getNativeLimit"), - outputAttribute.length, - serializedPlan, - numParts, - idx, - outputAttribute.map(_.dataType)) - } + CometExec.getCometIterator( + CometArrowStream.inputObjects(iter, inputSchema, "CometExecUtils-getNativeLimit"), + outputAttribute.length, + serializedPlan, + numParts, + idx, + sqlConf, + outputAttribute.map(_.dataType)) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala index e713e6ff85f..f9597ca046b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala @@ -152,10 +152,10 @@ case class CometTakeOrderedAndProjectExec( protected override def doExecuteColumnar(): RDD[ColumnarBatch] = { val childRDD = child.executeColumnar() - val sqlConf = CapturedSQLConf.capture() if (childRDD.getNumPartitions == 0 || limit == 0) { new ParallelCollectionRDD(sparkContext, Seq.empty[ColumnarBatch], 1, Map.empty) } else { + val sqlConf = CapturedSQLConf.capture() val singlePartitionRDD = if (childRDD.getNumPartitions == 1) { childRDD } else { @@ -173,16 +173,14 @@ case class CometTakeOrderedAndProjectExec( val numOutputCols = child.output.length val inputSchema = CometUtils.fromAttributes(child.output) childRDD.mapPartitionsWithIndexInternal { case (idx, iter) => - sqlConf.withConf { - CometExec.getCometIterator( - CometArrowStream - .inputObjects(iter, inputSchema, "CometTakeOrderedAndProject-topK"), - numOutputCols, - serializedTopK, - numParts, - idx, - child.output.map(_.dataType)) - } + CometExec.getCometIterator( + CometArrowStream.inputObjects(iter, inputSchema, "CometTakeOrderedAndProject-topK"), + numOutputCols, + serializedTopK, + numParts, + idx, + sqlConf, + child.output.map(_.dataType)) } } @@ -204,16 +202,15 @@ case class CometTakeOrderedAndProjectExec( val finalOutputTypes = output.map(_.dataType) val finalInputSchema = CometUtils.fromAttributes(child.output) singlePartitionRDD.mapPartitionsInternal { iter => - val it = sqlConf.withConf { - CometExec.getCometIterator( - CometArrowStream - .inputObjects(iter, finalInputSchema, "CometTakeOrderedAndProject-final"), - finalOutputTypes.length, - serializedFinalPlan, - 1, - 0, - finalOutputTypes) - } + val it = CometExec.getCometIterator( + CometArrowStream + .inputObjects(iter, finalInputSchema, "CometTakeOrderedAndProject-final"), + finalOutputTypes.length, + serializedFinalPlan, + 1, + 0, + sqlConf, + finalOutputTypes) setSubqueries(it.id, this) Option(TaskContext.get()).foreach { context => 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 641fd36da64..7a83843992a 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,8 +38,7 @@ import org.apache.comet.CometShuffleBlockIterator * whether the writer fed by this RDD places rows by position rather than by content; see * [[CometShuffleExchangeExec.positionalRoundRobinSpec]] and `getOutputDeterministicLevel`. * @param sqlConf - * the session's configs, captured on the driver, for resolving the inputs in a job that runs - * outside a SQL execution. See [[CapturedSQLConf]]. + * the session's Comet configs, which the writer also applies; see [[CapturedSQLConf]]. */ private[shuffle] class CometNativeShuffleInputRDD( sc: SparkContext, @@ -127,7 +126,8 @@ private[shuffle] class CometNativeShuffleInputRDD( partition.index, inputObjects, shuffleBlockIters, - partition.planDataByKey) + partition.planDataByKey, + sqlConf) } override def getPreferredLocations(split: Partition): Seq[String] = { @@ -164,7 +164,8 @@ private[shuffle] class CometNativeShuffleInputIterator( val partitionIndex: Int, val inputObjects: Array[Object], val shuffleBlockIterators: Map[Int, CometShuffleBlockIterator], - val planDataByKey: Map[String, Array[Byte]]) + val planDataByKey: Map[String, Array[Byte]], + val sqlConf: CapturedSQLConf) 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 cf8b6ffed9b..bdc8526b3ee 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 @@ -95,7 +95,8 @@ class CometNativeShuffleWriter[K, V]( override def write(inputs: Iterator[Product2[K, V]]): Unit = { try { - spec.sqlConf.withConf(writeInternal(inputs)) + val input = nativeShuffleInput(inputs) + input.sqlConf.withConf(writeInternal(input)) } catch { case failure: Throwable => remoteDestination.foreach { destination => @@ -111,7 +112,20 @@ class CometNativeShuffleWriter[K, V]( } } - private def writeInternal(inputs: Iterator[Product2[K, V]]): Unit = { + // The dep's _rdd is always a CometNativeShuffleInputRDD on this path. Pattern-match instead + // of asInstanceOf so a future RDD-layering change produces a clear error here rather than a + // bare ClassCastException deeper in the stack. + private def nativeShuffleInput( + inputs: Iterator[Product2[K, V]]): CometNativeShuffleInputIterator = inputs match { + case it: CometNativeShuffleInputIterator => it + case other => + throw new IllegalStateException( + "CometNativeShuffleWriter expects its input iterator to be a " + + "CometNativeShuffleInputIterator (produced by CometNativeShuffleInputRDD), got " + + s"${other.getClass.getName}") + } + + private def writeInternal(shuffleInputIter: CometNativeShuffleInputIterator): Unit = { val localOutput = if (remoteDestination.isEmpty) { val resolver = SparkEnv.get.shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver] @@ -123,17 +137,6 @@ class CometNativeShuffleWriter[K, V]( None } - // The dep's _rdd is always a CometNativeShuffleInputRDD on this path. Pattern-match instead - // of asInstanceOf so a future RDD-layering change produces a clear error here rather than a - // bare ClassCastException deeper in the stack. - val shuffleInputIter = inputs match { - case it: CometNativeShuffleInputIterator => it - case other => - throw new IllegalStateException( - "CometNativeShuffleWriter expects its input iterator to be a " + - "CometNativeShuffleInputIterator (produced by CometNativeShuffleInputRDD), got " + - s"${other.getClass.getName}") - } val partitionIdx = shuffleInputIter.partitionIndex val inputObjects = shuffleInputIter.inputObjects val shuffleBlockIters = shuffleInputIter.shuffleBlockIterators diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index f50bda5a1ed..be91caa0ec9 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala @@ -28,7 +28,7 @@ import org.apache.spark.shuffle.ShuffleWriteProcessor import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.physical.Partitioning -import org.apache.spark.sql.comet.{CapturedSQLConf, CometMetricNode, NativeExecContext} +import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext} import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types.StructType @@ -51,12 +51,7 @@ case class NativeShuffleSpec( * size so that a re-executed map task places rows exactly as the attempt it replaces. See * `CometShuffleExchangeExec.positionalRoundRobinSpec`. */ - positionalRoundRobin: Option[PositionalRoundRobin] = None, - /** - * The session's configs, captured on the driver, which the writer applies while it builds and - * runs the native plan in a job that runs outside a SQL execution. See [[CapturedSQLConf]]. - */ - sqlConf: CapturedSQLConf = CapturedSQLConf.capture()) + positionalRoundRobin: Option[PositionalRoundRobin] = None) /** * Parameters for positional round-robin placement, resolved on the driver. 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 3f331fc5b9c..6eb3e76eadb 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 @@ -798,6 +798,7 @@ object CometExec { /** * Create a CometExecIterator with a pre-serialized native plan. Use this overload when * executing the same plan across multiple partitions to avoid serializing the plan repeatedly. + * `sqlConf` is captured on the driver; see [[CapturedSQLConf]]. */ def getCometIterator( inputObjects: Array[Object], @@ -805,7 +806,8 @@ object CometExec { serializedPlan: Array[Byte], numParts: Int, partitionIdx: Int, - outputTypes: Seq[DataType] = Seq.empty): CometExecIterator = { + sqlConf: CapturedSQLConf, + outputTypes: Seq[DataType]): CometExecIterator = sqlConf.withConf { new CometExecIterator( newIterId, inputObjects, diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index b39d16f1235..d54ac56d5bc 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -36,7 +36,7 @@ import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.Join import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, SQLExecution, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.aggregate.SortAggregateExec import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException @@ -1450,10 +1450,7 @@ class CometJoinSuite extends CometTestBase { case _ => false }.isEmpty), plan) - // Propagates the session's confs, such as the batch size, to the tasks as collect() does. - val parts = SQLExecution.withSQLConfPropagated(spark) { - query.queryExecution.toRdd.map(_.getInt(0)).glom().collect() - } + val parts = query.queryExecution.toRdd.map(_.getInt(0)).glom().collect() parts.zipWithIndex.foreach { case (keys, i) => assert(keys.sameElements(keys.sorted), s"partition $i is out of order\n$plan") } 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 1f30786a27d..9a7b875e46b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -1996,15 +1996,10 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", CometConf.COMET_BATCH_SIZE.key -> "10") { val df = spark.read.parquet(dir.toString).repartition(2, col("k")) - // A job launched from the query's RDD runs outside a SQL execution, so Spark does not - // pass the session's configs to the map stage's tasks. + // Runs outside a SQL execution, so Spark passes no configs to the map tasks. df.queryExecution.toRdd.foreach(_ => ()) - val exchange = find(df.queryExecution.executedPlan) { - case _: CometShuffleExchangeExec => true - case _ => false - }.get - // The scan fused into the shuffle writer emits 10-row batches, not one of the default - // 8192 rows. + val exchange = checkCometExchange(df, 1, native = true).head + // 1000 rows in batches of 10. The default batch size would give 1. assert(exchange.metrics("input_batches").value == 100) } } diff --git a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala index a3046b7f99b..50eb2959cf1 100644 --- a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala +++ b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala @@ -34,7 +34,7 @@ import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvid import org.apache.arrow.vector.types.TimeUnit import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType} import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.comet.CometExec +import org.apache.spark.sql.comet.{CapturedSQLConf, CometExec} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.ConstantColumnVector import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, IntegerType, MapType, StringType, StructField, StructType, TimestampNTZType, TimestampType} @@ -507,10 +507,12 @@ class NativeUtilSuite extends CometTestBase { .build() .toByteArray + val sqlConf = CapturedSQLConf.capture() val actual = spark.sparkContext .parallelize(Seq(0), 1) .mapPartitions { _ => - val iterator = CometExec.getCometIterator(Array.empty[Object], 1, plan, 1, 0) + val iterator = + CometExec.getCometIterator(Array.empty[Object], 1, plan, 1, 0, sqlConf, Seq.empty) try { assert(iterator.hasNext) val column = iterator.next().column(0) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala index 9bf5e955954..4266f7f1605 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala @@ -36,7 +36,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometBypassMergeSortShuffleH import org.apache.spark.sql.comet.execution.shuffle.CometColumnarShuffle import org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffle import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec -import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} +import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.execution.metric.SQLMetric @@ -1224,19 +1224,17 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { // The task deserializes this metric together with the scan's metric node, so both // refer to the task-side copy that native execution publishes into. val outputRows = scan.metrics("output_rows") - val (midStream, atEnd) = SQLExecution.withSQLConfPropagated(spark) { - scan - .executeColumnar() - .mapPartitions { batches => - Thread.sleep(100) - batches.next() - val mid = outputRows.value - batches.foreach(_ => ()) - Iterator((mid, outputRows.value)) - } - .collect() - .head - } + val (midStream, atEnd) = scan + .executeColumnar() + .mapPartitions { batches => + Thread.sleep(100) + batches.next() + val mid = outputRows.value + batches.foreach(_ => ()) + Iterator((mid, outputRows.value)) + } + .collect() + .head check(midStream, atEnd) } From 04cdf52d63b8074b2dcfa6679278c121c3775b7e Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 28 Oct 2026 09:26:05 -0700 Subject: [PATCH 3/4] fix: apply the session's Comet configs to shuffle reads and test each native plan site A Comet shuffle read with no native plan above it in its stage, such as one feeding a Spark operator in an intermediate stage, still read Comet defaults. Capture the configs in CometShuffledBatchRDD and apply them in compute and computeAsShuffleBlockIterator, which makes the wrap in CometNativeShuffleInputRDD.compute redundant, so drop it. Make that RDD's sqlConf required so each call site captures on the driver explicitly. Replace the single test with one per site: the native shuffle writer, a native plan in a map stage, the native limit, the native top-K, each checking the memory pool every native plan resolves, and a shuffle read under spark.comet.shuffle.readBufferSize=0. --- .../shuffle/CometNativeShuffleInputRDD.scala | 13 +-- .../shuffle/CometShuffleExchangeExec.scala | 6 +- .../shuffle/CometShuffledRowRDD.scala | 16 +-- .../comet/exec/CometNativeShuffleSuite.scala | 98 ++++++++++++++++--- .../CometNativeShuffleInputRDDSuite.scala | 15 ++- 5 files changed, 117 insertions(+), 31 deletions(-) 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 7a83843992a..bc652d3092e 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 @@ import org.apache.comet.CometShuffleBlockIterator * whether the writer fed by this RDD places rows by position rather than by content; see * [[CometShuffleExchangeExec.positionalRoundRobinSpec]] and `getOutputDeterministicLevel`. * @param sqlConf - * the session's Comet configs, which the writer also applies; see [[CapturedSQLConf]]. + * the session's Comet configs, captured on the driver, which the writer applies; see + * [[CapturedSQLConf]]. */ private[shuffle] class CometNativeShuffleInputRDD( sc: SparkContext, @@ -47,8 +48,8 @@ private[shuffle] class CometNativeShuffleInputRDD( shuffleScanIndices: Set[Int], spillMetricNode: CometMetricNode, @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty, - positionalRoundRobin: Boolean = false, - sqlConf: CapturedSQLConf = CapturedSQLConf.capture()) + sqlConf: CapturedSQLConf, + positionalRoundRobin: Boolean = false) extends RDD[Product2[Int, ColumnarBatch]]( sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { @@ -66,8 +67,8 @@ private[shuffle] class CometNativeShuffleInputRDD( shuffleScanIndices, spillMetricNode, perPartitionByKey, - positionalRoundRobin, - sqlConf) + sqlConf, + positionalRoundRobin) /** * Spark's `isOrderSensitive` rule, applied to the RDD graph below the native plan. Spark only @@ -113,7 +114,7 @@ private[shuffle] class CometNativeShuffleInputRDD( override def compute( split: Partition, - context: TaskContext): Iterator[Product2[Int, ColumnarBatch]] = sqlConf.withConf { + context: TaskContext): Iterator[Product2[Int, ColumnarBatch]] = { spillMetricNode.reportSpillMetrics(context) val partition = split.asInstanceOf[CometNativeShuffleInputPartition] val (inputObjects, shuffleBlockIters) = 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 616a9716c8c..0ab25fcfb8b 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 @@ -36,7 +36,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Exp 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, CometSparkToColumnarExec, NativeExecContext} +import org.apache.spark.sql.comet.{CapturedSQLConf, CometFilterExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometPlan, CometProjectExec, CometSinkPlaceHolder, CometSparkToColumnarExec, NativeExecContext} import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNativeArrowSource} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec @@ -144,6 +144,7 @@ case class CometShuffleExchangeExec( ctx.shuffleScanIndices, CometMetricNode(metrics, Seq(nativeChildMetricNode)), ctx.perPartitionByKey, + CapturedSQLConf.capture(), positionalRoundRobin.isDefined) case None => child match { @@ -946,7 +947,8 @@ object CometShuffleExchangeExec Seq(streamRDD), streamRDD.getNumPartitions, shuffleScanIndices = Set.empty, - spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode))) + spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode)), + sqlConf = CapturedSQLConf.capture()) val ctx = NativeExecContext( inputs = Seq(streamRDD), diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala index dddb9767761..637d38b998d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala @@ -22,6 +22,7 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark._ import org.apache.spark.rdd.RDD import org.apache.spark.shuffle.sort.SortShuffleManager +import org.apache.spark.sql.comet.CapturedSQLConf import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.metric.{SQLMetric, SQLShuffleReadMetricsReporter} import org.apache.spark.sql.internal.SQLConf @@ -47,6 +48,8 @@ class CometShuffledBatchRDD( Array.tabulate(dependency.partitioner.numPartitions)(i => CoalescedPartitionSpec(i, i + 1))) } + private val sqlConf = CapturedSQLConf.capture() + dependency.rdd.context.setLocalProperty( SortShuffleManager.FETCH_SHUFFLE_BLOCKS_IN_BATCH_ENABLED_KEY, SQLConf.get.fetchShuffleBlocksInBatch.toString) @@ -166,7 +169,7 @@ class CometShuffledBatchRDD( */ def computeAsShuffleBlockIterator( split: Partition, - context: TaskContext): CometShuffleBlockIterator = { + context: TaskContext): CometShuffleBlockIterator = sqlConf.withConf { val readerAndMetrics: (CometShuffleReader[_, _], SQLShuffleReadMetricsReporter) = createReader(split, context) // Register before CometExecIterator so Spark's LIFO completion order closes the native stream @@ -180,11 +183,12 @@ class CometShuffledBatchRDD( CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt) } - override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { - val reader: CometShuffleReader[_, _] = createReader(split, context)._1 - // TODO: Reads IPC by native code - reader.read().asInstanceOf[Iterator[Product2[Int, ColumnarBatch]]].map(_._2) - } + override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = + sqlConf.withConf { + val reader: CometShuffleReader[_, _] = createReader(split, context)._1 + // TODO: Reads IPC by native code + reader.read().asInstanceOf[Iterator[Product2[Int, ColumnarBatch]]].map(_._2) + } override def clearDependencies(): Unit = { super.clearDependencies() 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 9a7b875e46b..19577ea9209 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -35,13 +35,14 @@ import org.apache.arrow.memory.ArrowBuf import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{Field, Schema} import org.apache.hadoop.fs.Path -import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.logging.log4j.Level +import org.apache.spark.{SparkEnv, SparkException, TaskContext} import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row} import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, SortOrder} 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, RangePartitioning} -import org.apache.spark.sql.comet.{CometExec, CometHashAggregateExec, CometLocalTableScanExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} +import org.apache.spark.sql.comet.{CometCollectLimitExec, CometExec, CometHashAggregateExec, CometLocalTableScanExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNativeArrowSource} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{FileSourceScanExec, LocalTableScanExec, SparkPlan} @@ -1984,23 +1985,96 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } - // https://github.com/apache/datafusion-comet/issues/7301 - test("map stage of a job outside a SQL execution uses the session's Comet configs") { + // https://github.com/apache/datafusion-comet/issues/7301: a job launched from a query's RDD + // runs outside a SQL execution, so Spark passes no configs to its map tasks. + private def withSessionConfsOutsideSqlExecution(f: String => Unit): Unit = { withTempPath { dir => spark - .range(0, 1000, 1, 1) + .range(0, 1000, 1, 4) .selectExpr("id % 10 as k", "id as v") .write .parquet(dir.toString) withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - CometConf.COMET_BATCH_SIZE.key -> "10") { - val df = spark.read.parquet(dir.toString).repartition(2, col("k")) - // Runs outside a SQL execution, so Spark passes no configs to the map tasks. - df.queryExecution.toRdd.foreach(_ => ()) - val exchange = checkCometExchange(df, 1, native = true).head - // 1000 rows in batches of 10. The default batch size would give 1. - assert(exchange.metrics("input_batches").value == 100) + CometConf.COMET_OFFHEAP_MEMORY_POOL_TYPE.key -> "greedy_unified") { + f(dir.toString) + } + } + } + + /** Runs `df` from its RDD and checks that every native plan got the session's memory pool. */ + private def checkNativePlansUseSessionConfs(df: DataFrame): Unit = { + val appender = new LogAppender("native plan memory pools") + withLogAppender(appender, Seq(classOf[CometExecIterator].getName), Some(Level.DEBUG)) { + df.queryExecution.toRdd.foreach(_ => ()) + } + val pools = appender.loggingEvents + .map(_.getMessage.getFormattedMessage) + .filter(_.startsWith("memoryPoolType=")) + .map(_.stripPrefix("memoryPoolType=").takeWhile(_ != ',')) + assert(pools.nonEmpty) + assert(pools.forall(_ == "greedy_unified"), pools) + } + + test("native shuffle writer outside a SQL execution uses the session's Comet configs") { + withSessionConfsOutsideSqlExecution { path => + val df = spark.read.parquet(path).repartition(2, col("k")) + checkCometExchange(df, 1, native = true) + checkNativePlansUseSessionConfs(df) + } + } + + test("native plan in a map stage outside a SQL execution uses the session's Comet configs") { + withSessionConfsOutsideSqlExecution { path => + withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + val df = spark.read.parquet(path).where(col("v") > 10).repartition(2, col("k")) + checkCometExchange(df, 1, native = false) + checkNativePlansUseSessionConfs(df) + } + } + } + + test("native limit outside a SQL execution uses the session's Comet configs") { + withSessionConfsOutsideSqlExecution { path => + val df = spark.read.parquet(path).limit(500) + assert( + find(df.queryExecution.executedPlan)(_.isInstanceOf[CometCollectLimitExec]).isDefined) + checkNativePlansUseSessionConfs(df) + } + } + + test("native top-K outside a SQL execution uses the session's Comet configs") { + withSessionConfsOutsideSqlExecution { path => + val df = spark.read.parquet(path).orderBy(col("v")).limit(50) + assert( + find(df.queryExecution.executedPlan)( + _.isInstanceOf[CometTakeOrderedAndProjectExec]).isDefined) + checkNativePlansUseSessionConfs(df) + } + } + + test("shuffle read outside a SQL execution uses the session's Comet configs") { + withSessionConfsOutsideSqlExecution { path => + // Only the shuffle reader reads this setting, and 0 fails its check, so the job fails only + // if the read sees the session's value. Spark's aggregate and project read the Comet + // shuffle in an intermediate stage, with no native plan above the read. + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.key -> "0") { + val df = spark.read + .parquet(path) + .groupBy("k") + .agg(sum("v").as("s")) + .repartition(3, col("s")) + .groupBy((col("s") % 7).as("b")) + .count() + assert(collect(df.queryExecution.executedPlan) { + case e: CometShuffleExchangeExec if e.shuffleType == CometColumnarShuffle => e + }.nonEmpty) + val e = intercept[SparkException](df.queryExecution.toRdd.foreach(_ => ())) + assert(e.getMessage.contains(CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.key), e) } } } 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 index 871255b1ecc..990233f20aa 100644 --- 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 @@ -23,7 +23,7 @@ import org.apache.spark.{HashPartitioner, Partition, TaskContext} import org.apache.spark.rdd.{DeterministicLevel, RDD} 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.comet.{CapturedSQLConf, CometMetricNode, NativeExecContext} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -61,7 +61,8 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { Seq(parent), 0, Set.empty, - CometMetricNode(Map.empty)) + CometMetricNode(Map.empty), + sqlConf = CapturedSQLConf.capture()) assert(input.outputDeterministicLevel == level) assert(input.copyForLocalShuffle().outputDeterministicLevel == level) } @@ -90,6 +91,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { 0, Set.empty, CometMetricNode(Map.empty), + sqlConf = CapturedSQLConf.capture(), positionalRoundRobin = true) assert(input.outputDeterministicLevel == expected, s"parent was $parentLevel") // The flag has to survive the copy, or a local-shuffle fallback silently drops the @@ -119,7 +121,8 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { 2, Set.empty, CometMetricNode(Map.empty), - planData) + planData, + CapturedSQLConf.capture()) val local = remote.copyForLocalShuffle() assert(local.id != remote.id) @@ -172,7 +175,8 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { Seq(nestedInput), 1, Set.empty, - CometMetricNode(writerMetrics, Seq(childMetrics))) + CometMetricNode(writerMetrics, Seq(childMetrics)), + sqlConf = CapturedSQLConf.capture()) inputRDD.iterator(inputRDD.partitions.head, taskContext) new CometNativeShuffleWriter[Int, Any]( @@ -218,7 +222,8 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { numPartitionsParam = numPartitions, shuffleScanIndices = Set.empty, spillMetricNode = CometMetricNode(writerMetrics, Seq(childMetricNode)), - perPartitionByKey = perPartitionByKey) + perPartitionByKey = perPartitionByKey, + sqlConf = CapturedSQLConf.capture()) val execContext = NativeExecContext( inputs = Seq.empty, numPartitions = numPartitions, From bcf48626d59c6eb29d9f0808de2055e761ad1e7d Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 28 Oct 2026 13:26:06 -0700 Subject: [PATCH 4/4] fix: apply the captured configs to the Celeborn frame size and capture DEBUG logs in tests The native Celeborn writer read spark.comet.shuffle.rss.maxFrameBytes in getWriter, before the writer applies the captured configs. Carry the configs on NativeShuffleSpec, which getWriter can reach through the dependency, and apply them to that read as well as in the writer, in place of CometNativeShuffleInputRDD and its iterator. Spark's LogAppender keeps INFO and above by default, so set its threshold to DEBUG for the memory pool line the tests read. --- .../shuffle/CometCelebornShuffleManager.scala | 9 +++--- .../shuffle/CometNativeShuffleInputRDD.scala | 13 ++------- .../shuffle/CometNativeShuffleWriter.scala | 29 +++++++++---------- .../shuffle/CometShuffleDependency.scala | 4 ++- .../shuffle/CometShuffleExchangeExec.scala | 7 ++--- .../comet/exec/CometNativeShuffleSuite.scala | 1 + .../CometNativeShuffleInputRDDSuite.scala | 21 +++++++------- 7 files changed, 38 insertions(+), 46 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala index b4c2a258cf9..921994b543c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala @@ -170,10 +170,11 @@ class CometCelebornShuffleManager private[shuffle] ( .computeIfAbsent(handle.shuffleId, _ => new ConcurrentHashMap[Int, AnyRef]()) .put(resolved.celebornShuffleId, resolved.client) + val spec = dependency.nativeShuffleSpec.getOrElse { + throw new IllegalStateException("Native Comet shuffle has no execution plan") + } new CometNativeShuffleWriter[K, V]( - dependency.nativeShuffleSpec.getOrElse { - throw new IllegalStateException("Native Comet shuffle has no execution plan") - }, + spec, dependency.outputPartitioning.getOrElse { throw new IllegalStateException("Native Comet shuffle has no output partitioning") }, @@ -189,7 +190,7 @@ class CometCelebornShuffleManager private[shuffle] ( CelebornNativeShuffleDestination( resolved.pusher, CometCelebornShuffleManager.maxNativeFrameBytes( - CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.get().toInt, + spec.sqlConf.withConf(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.get().toInt), resolved.pusher), dependency.partitioner.numPartitions, commitAuthorized = true, 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 bc652d3092e..8770496bff2 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 @@ -21,7 +21,7 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark._ import org.apache.spark.rdd.{DeterministicLevel, RDD} -import org.apache.spark.sql.comet.{CapturedSQLConf, CometExecRDD, CometMetricNode} +import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometShuffleBlockIterator @@ -37,9 +37,6 @@ import org.apache.comet.CometShuffleBlockIterator * @param positionalRoundRobin * whether the writer fed by this RDD places rows by position rather than by content; see * [[CometShuffleExchangeExec.positionalRoundRobinSpec]] and `getOutputDeterministicLevel`. - * @param sqlConf - * the session's Comet configs, captured on the driver, which the writer applies; see - * [[CapturedSQLConf]]. */ private[shuffle] class CometNativeShuffleInputRDD( sc: SparkContext, @@ -48,7 +45,6 @@ private[shuffle] class CometNativeShuffleInputRDD( shuffleScanIndices: Set[Int], spillMetricNode: CometMetricNode, @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty, - sqlConf: CapturedSQLConf, positionalRoundRobin: Boolean = false) extends RDD[Product2[Int, ColumnarBatch]]( sc, @@ -67,7 +63,6 @@ private[shuffle] class CometNativeShuffleInputRDD( shuffleScanIndices, spillMetricNode, perPartitionByKey, - sqlConf, positionalRoundRobin) /** @@ -127,8 +122,7 @@ private[shuffle] class CometNativeShuffleInputRDD( partition.index, inputObjects, shuffleBlockIters, - partition.planDataByKey, - sqlConf) + partition.planDataByKey) } override def getPreferredLocations(split: Partition): Seq[String] = { @@ -165,8 +159,7 @@ private[shuffle] class CometNativeShuffleInputIterator( val partitionIndex: Int, val inputObjects: Array[Object], val shuffleBlockIterators: Map[Int, CometShuffleBlockIterator], - val planDataByKey: Map[String, Array[Byte]], - val sqlConf: CapturedSQLConf) + 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 bdc8526b3ee..cf8b6ffed9b 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 @@ -95,8 +95,7 @@ class CometNativeShuffleWriter[K, V]( override def write(inputs: Iterator[Product2[K, V]]): Unit = { try { - val input = nativeShuffleInput(inputs) - input.sqlConf.withConf(writeInternal(input)) + spec.sqlConf.withConf(writeInternal(inputs)) } catch { case failure: Throwable => remoteDestination.foreach { destination => @@ -112,20 +111,7 @@ class CometNativeShuffleWriter[K, V]( } } - // The dep's _rdd is always a CometNativeShuffleInputRDD on this path. Pattern-match instead - // of asInstanceOf so a future RDD-layering change produces a clear error here rather than a - // bare ClassCastException deeper in the stack. - private def nativeShuffleInput( - inputs: Iterator[Product2[K, V]]): CometNativeShuffleInputIterator = inputs match { - case it: CometNativeShuffleInputIterator => it - case other => - throw new IllegalStateException( - "CometNativeShuffleWriter expects its input iterator to be a " + - "CometNativeShuffleInputIterator (produced by CometNativeShuffleInputRDD), got " + - s"${other.getClass.getName}") - } - - private def writeInternal(shuffleInputIter: CometNativeShuffleInputIterator): Unit = { + private def writeInternal(inputs: Iterator[Product2[K, V]]): Unit = { val localOutput = if (remoteDestination.isEmpty) { val resolver = SparkEnv.get.shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver] @@ -137,6 +123,17 @@ class CometNativeShuffleWriter[K, V]( None } + // The dep's _rdd is always a CometNativeShuffleInputRDD on this path. Pattern-match instead + // of asInstanceOf so a future RDD-layering change produces a clear error here rather than a + // bare ClassCastException deeper in the stack. + val shuffleInputIter = inputs match { + case it: CometNativeShuffleInputIterator => it + case other => + throw new IllegalStateException( + "CometNativeShuffleWriter expects its input iterator to be a " + + "CometNativeShuffleInputIterator (produced by CometNativeShuffleInputRDD), got " + + s"${other.getClass.getName}") + } val partitionIdx = shuffleInputIter.partitionIndex val inputObjects = shuffleInputIter.inputObjects val shuffleBlockIters = shuffleInputIter.shuffleBlockIterators diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index be91caa0ec9..ff813430ed9 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala @@ -28,7 +28,7 @@ import org.apache.spark.shuffle.ShuffleWriteProcessor import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.physical.Partitioning -import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext} +import org.apache.spark.sql.comet.{CapturedSQLConf, CometMetricNode, NativeExecContext} import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types.StructType @@ -44,6 +44,8 @@ case class NativeShuffleSpec( childNativeOp: OperatorOuterClass.Operator, childMetricNode: CometMetricNode, execContext: NativeExecContext, + /** The session's Comet configs, captured on the driver; see [[CapturedSQLConf]]. */ + sqlConf: CapturedSQLConf, /** * Set when round-robin placement is positional rather than content-hashed. Both the decision * and the group size are resolved once on the driver: the decision because it depends on the 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 0ab25fcfb8b..3a1a431cbd2 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 @@ -144,7 +144,6 @@ case class CometShuffleExchangeExec( ctx.shuffleScanIndices, CometMetricNode(metrics, Seq(nativeChildMetricNode)), ctx.perPartitionByKey, - CapturedSQLConf.capture(), positionalRoundRobin.isDefined) case None => child match { @@ -236,6 +235,7 @@ case class CometShuffleExchangeExec( nativeChild.nativeOp, nativeChildMetricNode, ctx, + CapturedSQLConf.capture(), positionalRoundRobin)) case None => child match { @@ -947,8 +947,7 @@ object CometShuffleExchangeExec Seq(streamRDD), streamRDD.getNumPartitions, shuffleScanIndices = Set.empty, - spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode)), - sqlConf = CapturedSQLConf.capture()) + spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode))) val ctx = NativeExecContext( inputs = Seq(streamRDD), @@ -970,7 +969,7 @@ object CometShuffleExchangeExec outputPartitioning, serializer, metrics, - NativeShuffleSpec(scanOp, childMetricNode, ctx)) + NativeShuffleSpec(scanOp, childMetricNode, ctx, CapturedSQLConf.capture())) } /** 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 19577ea9209..932d4a8bdc6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -2005,6 +2005,7 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper /** Runs `df` from its RDD and checks that every native plan got the session's memory pool. */ private def checkNativePlansUseSessionConfs(df: DataFrame): Unit = { val appender = new LogAppender("native plan memory pools") + appender.setThreshold(Level.DEBUG) withLogAppender(appender, Seq(classOf[CometExecIterator].getName), Some(Level.DEBUG)) { df.queryExecution.toRdd.foreach(_ => ()) } 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 index 990233f20aa..6ca87258b92 100644 --- 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 @@ -61,8 +61,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { Seq(parent), 0, Set.empty, - CometMetricNode(Map.empty), - sqlConf = CapturedSQLConf.capture()) + CometMetricNode(Map.empty)) assert(input.outputDeterministicLevel == level) assert(input.copyForLocalShuffle().outputDeterministicLevel == level) } @@ -91,7 +90,6 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { 0, Set.empty, CometMetricNode(Map.empty), - sqlConf = CapturedSQLConf.capture(), positionalRoundRobin = true) assert(input.outputDeterministicLevel == expected, s"parent was $parentLevel") // The flag has to survive the copy, or a local-shuffle fallback silently drops the @@ -121,8 +119,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { 2, Set.empty, CometMetricNode(Map.empty), - planData, - CapturedSQLConf.capture()) + planData) val local = remote.copyForLocalShuffle() assert(local.id != remote.id) @@ -175,12 +172,11 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { Seq(nestedInput), 1, Set.empty, - CometMetricNode(writerMetrics, Seq(childMetrics)), - sqlConf = CapturedSQLConf.capture()) + CometMetricNode(writerMetrics, Seq(childMetrics))) inputRDD.iterator(inputRDD.partitions.head, taskContext) new CometNativeShuffleWriter[Int, Any]( - NativeShuffleSpec(null, childMetrics, null), + NativeShuffleSpec(null, childMetrics, null, CapturedSQLConf.capture()), null, Nil, writerMetrics, @@ -222,8 +218,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { numPartitionsParam = numPartitions, shuffleScanIndices = Set.empty, spillMetricNode = CometMetricNode(writerMetrics, Seq(childMetricNode)), - perPartitionByKey = perPartitionByKey, - sqlConf = CapturedSQLConf.capture()) + perPartitionByKey = perPartitionByKey) val execContext = NativeExecContext( inputs = Seq.empty, numPartitions = numPartitions, @@ -237,7 +232,11 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { perPartitionFilePaths = Array.tabulate(numPartitions) { idx => Seq(s"file:/tmp/part-$idx.parquet") }) - val spec = NativeShuffleSpec(Operator.getDefaultInstance, childMetricNode, execContext) + val spec = NativeShuffleSpec( + Operator.getDefaultInstance, + childMetricNode, + execContext, + CapturedSQLConf.capture()) val dep = new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]( _rdd = rdd, partitioner = new HashPartitioner(numPartitions),