fix: apply the session's Comet configs in map stages of jobs outside a SQL execution

spark/src/main/scala/org/apache/spark/sql/comet/CapturedSQLConf.scala

@@ -0,0 +1,61 @@
+/*
+ * 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.internal.{ReadOnlySQLConf, SQLConf}
+
+import org.apache.comet.CometConf
+
+/**
+ * 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.
+ */
+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`, 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 (configs.isEmpty || context == null || !SQLConf.get.isInstanceOf[ReadOnlySQLConf] ||
+      configs.keys.exists(context.getLocalProperty(_) != null)) {
+      f
+    } else {
+      SQLConf.withExistingConf(executorSideConf)(f)
+    }
+  }
+}
+
+object CapturedSQLConf {
+
+  /** 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)))
+}

spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala

@@ -93,6 +93,8 @@ 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(", "))
 
+  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 +105,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.

spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala

@@ -72,13 +72,15 @@ 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,
+        sqlConf,
         outputAttribute.map(_.dataType))
     }
   }

spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala

@@ -155,6 +155,7 @@ case class CometTakeOrderedAndProjectExec(
     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 {
@@ -178,6 +179,7 @@ case class CometTakeOrderedAndProjectExec(
               serializedTopK,
               numParts,
               idx,
+              sqlConf,
               child.output.map(_.dataType))
           }
         }
@@ -207,6 +209,7 @@ case class CometTakeOrderedAndProjectExec(
           serializedFinalPlan,
           1,
           0,
+          sqlConf,
           finalOutputTypes)
         setSubqueries(it.id, this)
 

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,

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 =>

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

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
@@ -235,6 +235,7 @@ case class CometShuffleExchangeExec(
               nativeChild.nativeOp,
               nativeChildMetricNode,
               ctx,
+              CapturedSQLConf.capture(),
               positionalRoundRobin))
         case None =>
           child match {
@@ -968,7 +969,7 @@ object CometShuffleExchangeExec
       outputPartitioning,
       serializer,
       metrics,
-      NativeShuffleSpec(scanOp, childMetricNode, ctx))
+      NativeShuffleSpec(scanOp, childMetricNode, ctx, CapturedSQLConf.capture()))
   }
 
   /**

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()

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

@@ -798,14 +798,16 @@ 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],
       numOutputCols: Int,
       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,

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")
     }

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}
@@ -1983,4 +1984,99 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
       }
     }
   }
+
+  // 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, 4)
+        .selectExpr("id % 10 as k", "id as v")
+        .write
+        .parquet(dir.toString)
+      withSQLConf(
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+        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")
+    appender.setThreshold(Level.DEBUG)
+    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)
+      }
+    }
+  }
 }