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..1d15fde5f12 --- /dev/null +++ b/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))) +} 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..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,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. 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..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 @@ -72,6 +72,7 @@ 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"), @@ -79,6 +80,7 @@ object CometExecUtils { 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 af78e16e825..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 @@ -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) 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/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..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 616a9716c8c..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 @@ -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())) } /** 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/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index d6d54f4c8bc..07a6a769eb9 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 @@ -810,6 +810,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], @@ -817,7 +818,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 6b26b093486..b76f52330ac 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -1454,10 +1454,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 06660cef2a6..932d4a8bdc6 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} @@ -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) + } + } + } } 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) } 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..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 @@ -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 @@ -176,7 +176,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { inputRDD.iterator(inputRDD.partitions.head, taskContext) new CometNativeShuffleWriter[Int, Any]( - NativeShuffleSpec(null, childMetrics, null), + NativeShuffleSpec(null, childMetrics, null, CapturedSQLConf.capture()), null, Nil, writerMetrics, @@ -232,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),