From 042ab7d3f8fbf9ccee54de5370d1ba36f9ef2e17 Mon Sep 17 00:00:00 2001 From: msaf Date: Fri, 16 Oct 2026 09:54:52 +0100 Subject: [PATCH] fix: let a Comet plan's output be read on a thread other than the task's CometExecIterator looked up TaskContext.get() on every batch. A consumer that reads a partition's iterator on a thread of its own, such as the ClickHouse RowBinary writer that streams rows from its client's thread, saw a null context there and failed with a NullPointerException on the first batch. The ClickHouse client dropped that failure, so every insert of a Comet plan's output wrote no rows. Capture the task's context when the iterator is built, on the task thread, and use it for every batch. Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../org/apache/comet/CometExecIterator.scala | 13 +- .../comet/exec/CometOffTaskThreadSuite.scala | 155 ++++++++++++++++++ 4 files changed, 164 insertions(+), 6 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometOffTaskThreadSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 473006aa378..1b460c06dff 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -530,6 +530,7 @@ jobs: org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite org.apache.comet.exec.CometFilterShortCircuitSuite + org.apache.comet.exec.CometOffTaskThreadSuite org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index d00c7f6d2e2..f1f02a26878 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -178,6 +178,7 @@ jobs: org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite org.apache.comet.exec.CometFilterShortCircuitSuite + org.apache.comet.exec.CometOffTaskThreadSuite org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index b071de9e3f4..4a5e418c1ae 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -96,8 +96,9 @@ class CometExecIterator( private val memoryMXBean = ManagementFactory.getMemoryMXBean private val nativeLib = new Native() private val nativeUtil = new NativeUtil() - private val taskAttemptId = TaskContext.get().taskAttemptId() - private val taskCPUs = TaskContext.get().cpus() + private val taskContext = TaskContext.get() + private val taskAttemptId = taskContext.taskAttemptId() + private val taskCPUs = taskContext.cpus() private val cometTaskMemoryManager = new CometTaskMemoryManager(id, taskAttemptId) private val plan = { @@ -143,7 +144,7 @@ class CometExecIterator( // TaskContext and context ClassLoader. Read here because this class is only ever // constructed on a Spark task thread (see `taskAttemptId` above); a JNI-attached Tokio // worker has neither. See CometUdfBridge.evaluate. - TaskContext.get(), + taskContext, Thread.currentThread().getContextClassLoader) // Bind task-owned callbacks separately to preserve the existing createPlan JNI signature. @@ -220,7 +221,7 @@ class CometExecIterator( // Register a task completion listener to ensure native resources are released // when the task is done. - TaskContext.get().addTaskCompletionListener[Unit] { _ => + taskContext.addTaskCompletionListener[Unit] { _ => this.close() } @@ -228,7 +229,7 @@ class CometExecIterator( private val planLock = new Object private var planReleased = false - private val killWatch = CometExecIterator.watchForKill(TaskContext.get(), () => cancelPlan()) + private val killWatch = CometExecIterator.watchForKill(taskContext, () => cancelPlan()) private def cancelPlan(): Unit = planLock.synchronized { if (!planReleased) { @@ -241,7 +242,7 @@ class CometExecIterator( private def getNextBatch: Option[ColumnarBatch] = { assert(partitionIndex >= 0 && partitionIndex < numParts) - val ctx = TaskContext.get() + val ctx = taskContext CometTaskContextShim.killTaskIfInterrupted(ctx) try { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometOffTaskThreadSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometOffTaskThreadSuite.scala new file mode 100644 index 00000000000..4279f207598 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometOffTaskThreadSuite.scala @@ -0,0 +1,155 @@ +/* + * 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.comet.exec + +import java.util.concurrent.Executors + +import scala.concurrent.{Await, ExecutionContext, Future} +import scala.concurrent.duration.DurationInt + +import org.apache.spark.SparkContext +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.comet.CometExec +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.functions.{broadcast, col, count, lit, rand, round} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +class CometOffTaskThreadSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + private val rows = 60000 + private val days = 15 + + private def withTable(f: DataFrame => Unit): Unit = { + withTempPath { dir => + spark + .range(rows) + .selectExpr( + s"date_add(DATE'2026-10-01', cast(id % $days AS int)) AS date", + "concat('team', cast(id % 7 AS string)) AS team", + "concat('app', cast(id AS string)) AS app_id", + "id AS apps", + "cast(id * 1.5 AS double) AS cost") + .repartition(12) + .write + .parquet(dir.getCanonicalPath) + f(spark.read.parquet(dir.getCanonicalPath)) + } + } + + private def splitByDate(df: DataFrame): DataFrame = { + val counts = df + .groupBy(col("date")) + .agg(count(lit(1)) as "cnt") + .withColumn("maxRand", round(col("cnt") / 5e6)) + df.join(broadcast(counts), Seq("date"), "left_outer") + .repartitionByRange(col("date"), round(rand() * col("maxRand")).as("rand")) + .drop("rand", "maxRand", "cnt") + } + + private def partitionSizes(rdd: RDD[_]): Seq[(Int, Int)] = + rdd + .mapPartitionsWithIndex { case (i, it) => Iterator((i, it.size)) } + .collect() + .toSeq + + private def rereadOffTaskThread(rdd: RDD[_], partitions: Seq[Int]): Seq[(Int, Int)] = { + val pool = Executors.newFixedThreadPool(4) + implicit val ec: ExecutionContext = ExecutionContext.fromExecutorService(pool) + try { + val reads = partitions.map { i => + Future( + (i, CometOffTaskThreadSuite.countOffTaskThread(spark.sparkContext, rdd, Seq(i)).head)) + } + Await.result(Future.sequence(reads), 5.minutes).sortBy(_._1) + } finally { + pool.shutdown() + } + } + + private def hasCometExec(df: DataFrame): Boolean = + find(df.queryExecution.executedPlan)(_.isInstanceOf[CometExec]).isDefined + + Seq("native", "jvm").foreach { mode => + Seq("true", "false").foreach { aqe => + test(s"a reused $mode shuffle read on a thread other than the task's, aqe=$aqe") { + withTable { df => + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> mode, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + "spark.sql.adaptive.coalescePartitions.parallelismFirst" -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "50", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "1k", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_SIZE.key -> "1k") { + val split = splitByDate(df) + val rdd = split.rdd + assert(hasCometExec(split)) + val sizes = partitionSizes(rdd).filter(_._2 > 0) + assert(sizes.size >= days) + assert(sizes.map(_._2).sum == rows) + assert(rereadOffTaskThread(rdd, sizes.map(_._1)) == sizes) + } + } + } + } + } + + test("a native scan read on a thread other than the task's") { + withTable { df => + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "true") { + val filtered = df.filter(col("apps") % 3 === 0).select(col("app_id"), col("cost")) + assert(hasCometExec(filtered)) + val rdd = filtered.rdd + val counts = CometOffTaskThreadSuite.countOffTaskThread( + spark.sparkContext, + rdd, + rdd.partitions.indices) + assert(counts.sum == rows / 3) + } + } + } +} + +object CometOffTaskThreadSuite { + + def countOffTaskThread(sc: SparkContext, rdd: RDD[_], partitions: Seq[Int]): Seq[Int] = + sc.runJob(rdd, countOnNewThread _, partitions).toSeq + + private def countOnNewThread(it: Iterator[_]): Int = { + var size = -1 + var error: Throwable = null + val reader = new Thread(new Runnable { + override def run(): Unit = + try { + size = it.size + } catch { + case e: Throwable => error = e + } + }) + reader.start() + reader.join() + if (error != null) throw error + size + } +}