You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/06 20:47:54 UTC
(datafusion-comet) 01/01: fix: initialize dispatched kernels with the partition the native plan computes (#6700)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6700-6cd5ea0ed84e10087fac3066baa30ebd65c870ce
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 8d0bb01a5a99924024e07c1722e6eddd9260e167
Author: Oleks V <co...@users.noreply.github.com>
AuthorDate: Tue Oct 6 20:47:30 2026 +0000
fix: initialize dispatched kernels with the partition the native plan computes (#6700)
* fix: initialize dispatched kernels with the partition the native plan computes
The JVM codegen dispatcher seeded kernels from `TaskContext.partitionId()`, which differs from
the partition being computed under a union, a coalesce and a cartesian product. It also kept one
kernel per task, so the parent partitions of a coalesce shared random state.
The native planner now passes its partition index and plan id through the UDF bridge.
Nondeterministic kernels are cached per plan and initialized with that index. Deterministic
kernels never read it and stay shared across the plans of a task.
* fix: drop a closed plan's dispatched kernels and check every node for determinism
A coalesce kept the nondeterministic kernel of every parent partition until the task ended.
`CometExecIterator.close` now calls `CometUdfBridge.releasePlan`, which hands the closed plan to a
new `CometUDF.releasePlan` hook, and the dispatcher drops that plan's kernels there. Kernels
shared by every plan stay until the task ends.
The dispatcher shared a kernel whenever the root expression reported deterministic, but `Invoke`
leaves its target object out of `deterministic`. So `make_valid_utf8` over `spark_partition_id()`
still shared one kernel across the plans of a coalesce. It now checks every node.
---
native/core/src/execution/planner.rs | 2 +
native/core/src/parquet/parquet_exec.rs | 2 +
native/jni-bridge/src/comet_udf_bridge.rs | 2 +-
native/spark-expr/src/jvm_udf/mod.rs | 13 ++++
.../java/org/apache/comet/udf/CometUdfBridge.java | 33 ++++++++-
.../scala/org/apache/comet/CometExecIterator.scala | 5 ++
.../main/scala/org/apache/comet/udf/CometUDF.scala | 19 +++++
.../comet/udf/codegen/CometScalaUDFCodegen.scala | 84 +++++++++++++++-------
.../org/apache/comet/CometCodegenSourceSuite.scala | 4 +-
.../scala/org/apache/comet/CometCodegenSuite.scala | 63 ++++++++++++++--
.../org/apache/comet/udf/CometUdfBridgeSuite.scala | 4 ++
11 files changed, 197 insertions(+), 34 deletions(-)
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 99a52fb562..10f3eddb36 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -1045,6 +1045,8 @@ impl PhysicalPlanner {
udf.return_nullable,
self.task_context.clone(),
self.class_loader.clone(),
+ self.partition,
+ self.exec_context_id,
)))
}
ExprStruct::NativeScalarUdf(call) => {
diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs
index 8149c3656b..6f9803d994 100644
--- a/native/core/src/parquet/parquet_exec.rs
+++ b/native/core/src/parquet/parquet_exec.rs
@@ -417,6 +417,8 @@ mod tests {
true,
None,
None,
+ 0,
+ 0,
));
let literal: Arc<dyn PhysicalExpr> = Arc::new(Literal::new(ScalarValue::Int32(Some(1))));
let nested: Arc<dyn PhysicalExpr> =
diff --git a/native/jni-bridge/src/comet_udf_bridge.rs b/native/jni-bridge/src/comet_udf_bridge.rs
index bbd7465d00..6532fe0403 100644
--- a/native/jni-bridge/src/comet_udf_bridge.rs
+++ b/native/jni-bridge/src/comet_udf_bridge.rs
@@ -42,7 +42,7 @@ impl<'a> CometUdfBridge<'a> {
JNIString::new(Self::JVM_CLASS),
jni::jni_str!("evaluate"),
jni::jni_sig!(
- "(Ljava/lang/String;[J[JJJILorg/apache/spark/TaskContext;Ljava/lang/ClassLoader;)V"
+ "(Ljava/lang/String;[J[JJJIIJLorg/apache/spark/TaskContext;Ljava/lang/ClassLoader;)V"
),
)?,
method_evaluate_ret: ReturnType::Primitive(Primitive::Void),
diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs
index 232926744d..0a37330097 100644
--- a/native/spark-expr/src/jvm_udf/mod.rs
+++ b/native/spark-expr/src/jvm_udf/mod.rs
@@ -54,9 +54,14 @@ pub struct JvmScalarUdfExpr {
/// duration of the call. `None` when no driving Spark task is available (unit tests, direct
/// native driver runs); the bridge then installs nothing.
class_loader: Option<Arc<Global<JObject<'static>>>>,
+ /// Index of the partition this native plan computes. See `CometUDF.evaluate`.
+ partition: i32,
+ /// Id of this native plan. See `CometUDF.evaluate`.
+ exec_context_id: i64,
}
impl JvmScalarUdfExpr {
+ #[allow(clippy::too_many_arguments)]
pub fn new(
class_name: String,
args: Vec<Arc<dyn PhysicalExpr>>,
@@ -64,6 +69,8 @@ impl JvmScalarUdfExpr {
return_nullable: bool,
task_context: Option<Arc<Global<JObject<'static>>>>,
class_loader: Option<Arc<Global<JObject<'static>>>>,
+ partition: i32,
+ exec_context_id: i64,
) -> Self {
debug_assert!(
!class_name.is_empty(),
@@ -76,6 +83,8 @@ impl JvmScalarUdfExpr {
return_nullable,
task_context,
class_loader,
+ partition,
+ exec_context_id,
}
}
}
@@ -234,6 +243,8 @@ impl PhysicalExpr for JvmScalarUdfExpr {
JValue::Long(out_arr_ptr).as_jni(),
JValue::Long(out_sch_ptr).as_jni(),
JValue::Int(batch.num_rows() as i32).as_jni(),
+ JValue::Int(self.partition).as_jni(),
+ JValue::Long(self.exec_context_id).as_jni(),
JValue::Object(task_context_ref).as_jni(),
JValue::Object(class_loader_ref).as_jni(),
],
@@ -276,6 +287,8 @@ impl PhysicalExpr for JvmScalarUdfExpr {
self.return_nullable,
self.task_context.clone(),
self.class_loader.clone(),
+ self.partition,
+ self.exec_context_id,
)))
}
}
diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
index d2881bfe44..3ab34fe049 100644
--- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
+++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
@@ -32,6 +32,8 @@ import org.apache.spark.TaskContext;
import org.apache.spark.comet.CometTaskContextShim;
import org.apache.spark.util.TaskCompletionListener;
+import com.google.common.annotations.VisibleForTesting;
+
import org.apache.comet.util.ClassLoaders;
/**
@@ -86,6 +88,8 @@ public class CometUdfBridge {
* @param numRows row count of the current batch. Mirrors DataFusion's {@code
* ScalarFunctionArgs.number_rows}; the only batch-size signal a zero-input UDF (e.g. a
* zero-arg non-deterministic ScalaUDF) ever sees.
+ * @param partitionIndex index of the partition the calling native plan computes
+ * @param planId id of the calling native plan
* @param taskContext propagated Spark {@link TaskContext} from the driving Spark task thread, or
* {@code null} outside a Spark task. Treated as ground truth for the call: installed as the
* thread-local on entry, with the prior value (if any) saved and restored in {@code finally}.
@@ -113,6 +117,8 @@ public class CometUdfBridge {
long outArrayPtr,
long outSchemaPtr,
int numRows,
+ int partitionIndex,
+ long planId,
TaskContext taskContext,
ClassLoader classLoader) {
assert udfClassName != null && !udfClassName.isEmpty() : "udfClassName must be non-empty";
@@ -148,6 +154,8 @@ public class CometUdfBridge {
outArrayPtr,
outSchemaPtr,
numRows,
+ partitionIndex,
+ planId,
taskContext);
} finally {
// Unconditional: a no-op when nothing was installed, and it also undoes any change the
@@ -170,6 +178,8 @@ public class CometUdfBridge {
long outArrayPtr,
long outSchemaPtr,
int numRows,
+ int partitionIndex,
+ long planId,
TaskContext taskContext) {
long taskAttemptId = (taskContext != null) ? taskContext.taskAttemptId() : NO_TASK_ID;
@@ -228,7 +238,7 @@ public class CometUdfBridge {
inputs[i] = Data.importVector(importAllocator, inArr, inSch, null);
}
- result = udf.evaluate(inputs, numRows);
+ result = udf.evaluate(inputs, numRows, partitionIndex, planId);
// Recorded before the checks below, so the invariant holds however this exits.
resultIsInput = isOneOf(result, inputs);
if (!(result instanceof FieldVector)) {
@@ -290,6 +300,27 @@ public class CometUdfBridge {
}
}
+ /**
+ * Called by {@code CometExecIterator} once native plan {@code planId} of task {@code
+ * taskAttemptId} has closed. Passes the plan to every {@link CometUDF} instance of the task, see
+ * {@code CometUDF.releasePlan}.
+ */
+ public static void releasePlan(long taskAttemptId, long planId) {
+ ConcurrentHashMap<String, CometUDF> perTask = INSTANCES.get(taskAttemptId);
+ if (perTask != null) {
+ for (CometUDF udf : perTask.values()) {
+ udf.releasePlan(planId);
+ }
+ }
+ }
+
+ /** The instance of {@code udfClassName} that task {@code taskAttemptId} holds, or null. */
+ @VisibleForTesting
+ public static CometUDF instanceFor(long taskAttemptId, String udfClassName) {
+ ConcurrentHashMap<String, CometUDF> perTask = INSTANCES.get(taskAttemptId);
+ return perTask == null ? null : perTask.get(udfClassName);
+ }
+
/** Whether the UDF handed back one of the vectors it was given, rather than a new one. */
private static boolean isOneOf(ValueVector result, ValueVector[] inputs) {
for (ValueVector input : inputs) {
diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
index 0f03a6244e..27f880f913 100644
--- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
+++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
@@ -45,6 +45,7 @@ import org.apache.comet.exceptions.CometQueryExecutionException
import org.apache.comet.parquet.CometFileKeyUnwrapper
import org.apache.comet.serde.Config.ConfigMap
import org.apache.comet.shuffle.ShufflePartitionPusher
+import org.apache.comet.udf.CometUdfBridge
import org.apache.comet.vector.NativeUtil
/**
@@ -356,6 +357,10 @@ class CometExecIterator(
// global refs.
attempt(nativeLib.releasePlan(plan))
+ // Drops what the task's JVM UDFs keep for this plan, such as its dispatched kernels. `id` is
+ // the plan id the native plan passed them.
+ attempt(CometUdfBridge.releasePlan(taskAttemptId, id))
+
// Run the diagnostics even when teardown failed: a failed teardown is exactly when the
// non-zero memory usage warning below is most informative.
attempt {
diff --git a/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala
index 6b435c4064..db9178fc32 100644
--- a/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala
+++ b/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala
@@ -45,4 +45,23 @@ import org.apache.arrow.vector.ValueVector
*/
trait CometUDF {
def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector
+
+ /**
+ * The overload the bridge calls. `partitionIndex` is the index of the partition the calling
+ * native plan computes, which differs from `TaskContext.partitionId()` under a union, a
+ * coalesce or a cartesian product. `planId` tells apart the plans one task runs, such as the
+ * parent partitions of a coalesce. The default ignores both.
+ */
+ def evaluate(
+ inputs: Array[ValueVector],
+ numRows: Int,
+ partitionIndex: Int,
+ planId: Long): ValueVector =
+ evaluate(inputs, numRows)
+
+ /**
+ * Called once native plan `planId` has closed, so an implementation that keeps state per plan
+ * can drop it before the task ends. The default does nothing.
+ */
+ def releasePlan(planId: Long): Unit = ()
}
diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala
index 1cab882df9..c589fb6fa9 100644
--- a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala
+++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala
@@ -29,7 +29,7 @@ import scala.util.control.NonFatal
import org.apache.arrow.vector._
import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
import org.apache.arrow.vector.types.pojo.Field
-import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.SparkEnv
import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.expressions.Expression
import org.apache.spark.sql.comet.util.Utils
@@ -41,8 +41,8 @@ import org.apache.comet.udf.CometUDF
/**
* Arrow-direct codegen dispatcher. For each `(bound expression, input Arrow schema)` pair,
- * compiles a specialized [[CometBatchKernel]] on first encounter, initializes it with the task's
- * partition index, and caches the live instance.
+ * compiles a specialized [[CometBatchKernel]] on first encounter, initializes it with the index
+ * of the partition the calling native plan computes, and caches the live instance.
*
* Arg 0 is a `VarBinaryVector` scalar carrying the closure-serialized bound `Expression` bytes;
* args 1..N are the data columns the `BoundReference`s read in ordinal order.
@@ -56,16 +56,20 @@ import org.apache.comet.udf.CometUDF
* +----------------------------+ +----------------------------+ +----------------------------+
* | Key: generated Java | | Key: task + UDF class | | Key: bound expression + |
* | source | | | | input column shapes |
- * | Value: compiled Java class | | Value: dispatcher object | | Value: ready-to-run kernel |
- * | Scope: JVM, all queries | | Scope: one Spark task | | with state primed |
- * | share it | | | | Scope: one Spark task |
- * | Owner: Spark | | Owner: Comet | | (lives inside 2) |
+ * | Value: compiled Java class | | Value: dispatcher object | | (+ native plan if |
+ * | Scope: JVM, all queries | | Scope: one Spark task | | nondeterministic) |
+ * | share it | | | | Value: ready-to-run kernel |
+ * | Owner: Spark | | Owner: Comet | | with state primed |
+ * | | | | | Scope: one Spark task |
+ * | | | | | (lives inside 2), |
+ * | | | | | or one native plan |
+ * | | | | | if nondeterministic |
* | | | | | Owner: Comet |
* +----------------------------+ +----------------------------+ +----------------------------+
* }}}
*
- * Stateful expressions (`Rand`, `MonotonicallyIncreasingID`) advance inside the per-task kernel
- * across batches.
+ * Stateful expressions (`Rand`, `MonotonicallyIncreasingID`) advance inside the per-plan kernel
+ * across batches. `CometExecIterator.close` drops a plan's kernels through `releasePlan`.
*
* `evaluate` runs under `this.synchronized` because DataFusion operators like `HashJoinExec`
* pipeline build/probe via `OnceAsync` (`tokio::spawn`), so multiple Tokio worker threads can
@@ -81,13 +85,34 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
* Per-task cache keyed on serialized expression bytes plus per-column specs. The deserialized
* `boundExpr` carries mutable state (`NamedLambdaVariable.value` for HOFs, `Rand`'s
* `XORShiftRandom`) that must not be shared across concurrent tasks running the same query;
- * keeping the cache per-task gives each task its own copy. Guarded by `this.synchronized`.
+ * keeping the cache per-task gives each task its own copy. A nondeterministic kernel is seeded
+ * from the partition its plan computes, so its key also holds the plan: each plan a task runs,
+ * such as each parent partition of a coalesce, gets its own, dropped by `releasePlan` when the
+ * plan closes. Guarded by `this.synchronized`.
*/
private val kernelCache
: mutable.Map[CometScalaUDFCodegen.CacheKey, CometScalaUDFCodegen.CacheEntry] =
mutable.HashMap.empty
- override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector = {
+ // Kernels shared by every plan stay until the task ends.
+ override def releasePlan(planId: Long): Unit = this.synchronized {
+ kernelCache.keys.filter(_.planId == planId).toList.foreach(kernelCache.remove)
+ }
+
+ /** Plan ids of the cached kernels, `NoPlan` for a shared one. */
+ private[comet] def cachedPlanIds: List[Long] = this.synchronized {
+ kernelCache.keysIterator.map(_.planId).toList.sorted
+ }
+
+ // Callers that bypass the bridge (unit tests, benchmarks) have no native plan.
+ override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector =
+ evaluate(inputs, numRows, partitionIndex = 0, planId = CometScalaUDFCodegen.NoPlan)
+
+ override def evaluate(
+ inputs: Array[ValueVector],
+ numRows: Int,
+ partitionIndex: Int,
+ planId: Long): ValueVector = {
require(
inputs.length >= 1,
"CometScalaUDFCodegen requires at least 1 input (serialized expression), " +
@@ -115,12 +140,12 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
val n = numRows
val specsSeq = specs.toIndexedSeq
- val key = CometScalaUDFCodegen.CacheKey(ByteBuffer.wrap(bytes), specsSeq)
+ val key = CometScalaUDFCodegen.CacheKey(planId, ByteBuffer.wrap(bytes), specsSeq)
// Cache lookup and `process` run under one lock to serialize concurrent Tokio callers that
// would otherwise race on the kernel's per-batch instance fields.
this.synchronized {
- val entry = lookupOrCompile(key, bytes, specsSeq)
+ val entry = lookupOrCompile(key, bytes, specsSeq, partitionIndex)
val out = CometBatchKernelCodegen.allocateOutput(
entry.outputField,
@@ -144,9 +169,14 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
private def lookupOrCompile(
key: CometScalaUDFCodegen.CacheKey,
bytes: Array[Byte],
- specs: IndexedSeq[ArrowColumnSpec]): CometScalaUDFCodegen.CacheEntry = {
+ specs: IndexedSeq[ArrowColumnSpec],
+ partitionIndex: Int): CometScalaUDFCodegen.CacheEntry = {
assert(Thread.holdsLock(this), "lookupOrCompile must run under this.synchronized")
- kernelCache.get(key) match {
+ // A deterministic kernel never reads the partition index, so one instance under the planless
+ // key serves every plan in the task. Only a kernel with a nondeterministic node is stored per
+ // plan.
+ val sharedKey = key.copy(planId = CometScalaUDFCodegen.NoPlan)
+ kernelCache.get(sharedKey).orElse(kernelCache.get(key)) match {
case Some(entry) =>
CometScalaUDFCodegen.cacheHitCount.incrementAndGet()
entry
@@ -168,14 +198,17 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
}
val compiled = CometBatchKernelCodegen.compile(boundExpr, specs)
val kernel = compiled.newInstance()
- kernel.init(CometScalaUDFCodegen.currentPartitionIndex())
+ kernel.init(partitionIndex)
val outputField = CometBatchKernelCodegen.toFfiArrowField(
"codegen_result",
boundExpr.dataType,
boundExpr.nullable)
val entry =
CometScalaUDFCodegen.CacheEntry(compiled, kernel, boundExpr.dataType, outputField)
- kernelCache.put(key, entry)
+ // Walks the tree, because `deterministic` is not transitive everywhere. `Invoke` skips its
+ // `targetObject`, which is where Spark 4's `make_valid_utf8` puts its input.
+ val perPlan = boundExpr.exists(!_.deterministic)
+ kernelCache.put(if (perPlan) key else sharedKey, entry)
CometScalaUDFCodegen.compileCount.incrementAndGet()
CometScalaUDFCodegen.recordCompiledSignature(specs, boundExpr.dataType)
entry
@@ -300,21 +333,22 @@ object CometScalaUDFCodegen {
}
/**
- * Partition index for the kernel's `init`. Expressions whose `doGenCode` calls
- * `addPartitionInitializationStatement` (`Rand`, `Randn`, `Uuid`) reseed mutable state from
- * this. Falls back to 0 when the dispatcher is exercised outside a Spark task (unit tests).
+ * Plan id for kernels shared by every plan in a task, and for callers that bypass the bridge.
*/
- private def currentPartitionIndex(): Int =
- Option(TaskContext.get()).map(_.partitionId()).getOrElse(0)
+ private[comet] val NoPlan = -1L
/**
- * Cache key: serialized expression bytes plus per-column compile-time invariants. `hashCode`
- * walks `bytesKey` per lookup, so for large ScalaUDF closures it scales with closure size.
+ * Cache key: calling native plan (`NoPlan` for a deterministic expression), serialized
+ * expression bytes plus per-column compile-time invariants. `hashCode` walks `bytesKey` per
+ * lookup, so for large ScalaUDF closures it scales with closure size.
*
* TODO(perf-cache-key): if hot, options are a driver-precomputed hash piggybacked through the
* proto, per-instance last-key memoization, or a two-tier cache keyed on the generated source.
*/
- final case class CacheKey(bytesKey: ByteBuffer, specs: IndexedSeq[ArrowColumnSpec])
+ final case class CacheKey(
+ planId: Long,
+ bytesKey: ByteBuffer,
+ specs: IndexedSeq[ArrowColumnSpec])
/** Snapshot of dispatcher cache counters and current size. */
final case class DispatcherStats(compileCount: Long, cacheHitCount: Long, cacheSize: Int) {
diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
index 400ec6fcd2..f1272e151e 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
@@ -1535,8 +1535,8 @@ class CometCodegenSourceSuite extends AnyFunSuite {
IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = true))
val nonNullable =
IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = false))
- val k1 = CometScalaUDFCodegen.CacheKey(bytes, nullable)
- val k2 = CometScalaUDFCodegen.CacheKey(bytes, nonNullable)
+ val k1 = CometScalaUDFCodegen.CacheKey(0L, bytes, nullable)
+ val k2 = CometScalaUDFCodegen.CacheKey(0L, bytes, nonNullable)
assert(
k1 != k2,
"expected nullable=true and nullable=false specs to produce distinct cache keys")
diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
index 8f6d7c4a8d..5a80865da9 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
@@ -45,6 +45,7 @@ import org.apache.comet.codegen.CometBatchKernelCodegen
import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec
import org.apache.comet.serde.{CometInvokeTargets, CometScalaUDF, QueryPlanSerde}
import org.apache.comet.serde.ExprOuterClass.Expr.ExprStructCase
+import org.apache.comet.udf.CometUdfBridge
import org.apache.comet.udf.codegen.CometScalaUDFCodegen
import org.apache.comet.vector.CometVector
@@ -1621,11 +1622,10 @@ class CometCodegenSuite
// `init(int partitionIndex)`, called once per kernel allocation. Spark seeds
// `XORShiftRandom(seed + partitionIndex)` per partition, so different partitions produce
// different sequences for the same seed. Matching Spark across partitions requires the
- // kernel to see the real partition index, which the dispatcher derives from
- // `TaskContext.get().partitionId()`, live on this path thanks to the bridge-level
- // TaskContext propagation. Composing with a ScalaUDF (identity on Double here) forces the
- // tree through codegen dispatch so the Rand evaluation runs inside our kernel's init
- // rather than via Spark's normal codegen.
+ // kernel to see the real partition index, which the native plan passes through the bridge.
+ // Composing with a ScalaUDF (identity on Double here) forces the tree through codegen
+ // dispatch so the Rand evaluation runs inside our kernel's init rather than via Spark's
+ // normal codegen.
spark.udf.register("dblId", (d: Double) => d)
val df = spark
.range(0, 1024, 1, numPartitions = 4)
@@ -1645,6 +1645,59 @@ class CometCodegenSuite
}
}
+ test("dispatched kernels see the index of the partition the native plan computes") {
+ // Under a union, a coalesce and a cartesian product the native plan computes a partition
+ // whose index differs from `TaskContext.partitionId()`, and the coalesce runs two plans in
+ // one task. See https://github.com/apache/datafusion-comet/issues/6570.
+ val df = spark
+ .range(0, 8, 1, numPartitions = 2)
+ .selectExpr("id", "map(1, spark_partition_id()) AS m", "round(rand(42), 6) AS r")
+ Seq(df.union(df), df.coalesce(1)).foreach { q =>
+ assertCodegenRan(checkSparkAnswerAndOperator(q))
+ }
+ withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ // `CartesianProductExec` stays in Spark, over native children.
+ assertCodegenRan(checkSparkAnswer(df.crossJoin(spark.range(0, 4, 1, 2))))
+ }
+ if (isSpark40Plus) {
+ // Spark 4 lowers `make_valid_utf8` to an `Invoke` whose `deterministic` skips its target
+ // object, so the partition-seeded input sits under a root that reports deterministic.
+ val invoke = spark
+ .range(0, 8, 1, numPartitions = 2)
+ .selectExpr(
+ "id",
+ "make_valid_utf8(cast(spark_partition_id() AS STRING)) AS p",
+ "make_valid_utf8(cast(round(rand(42), 6) AS STRING)) AS r")
+ assertCodegenRan(checkSparkAnswerAndOperator(invoke.coalesce(1)))
+ }
+ }
+
+ test("plans of a coalesce share deterministic kernels and drop their own when they close") {
+ // The coalesce runs four native plans in one task. The deterministic `round` compiles once
+ // for all of them and the nondeterministic one once per plan. Each plan's kernel is dropped
+ // when the plan closes, so once the rows are read only the shared one is left.
+ val df = spark
+ .range(0, 16, 1, numPartitions = 4)
+ .selectExpr("round(id / 3, 2) AS d", "round(rand(42), 6) AS r")
+ .coalesce(1)
+ CometScalaUDFCodegen.resetStats()
+ val cached = df.queryExecution.toRdd
+ .mapPartitions { rows =>
+ rows.foreach(_ => ())
+ val dispatcher = CometUdfBridge.instanceFor(
+ TaskContext.get().taskAttemptId(),
+ classOf[CometScalaUDFCodegen].getName)
+ Iterator(Option(dispatcher).map(_.asInstanceOf[CometScalaUDFCodegen].cachedPlanIds))
+ }
+ .collect()
+ .toSeq
+ assert(
+ cached == Seq(Some(List(CometScalaUDFCodegen.NoPlan))),
+ s"expected only the shared kernel to outlive the plans, got plan ids $cached")
+ val stats = CometScalaUDFCodegen.stats()
+ assert(stats.compileCount == 5, s"expected 1 shared and 4 per-plan compiles, got $stats")
+ }
+
test("ScalaUDF composed with reused scalar subquery across projection and filter") {
// The same scalar subquery appears in two sites: the projection (which the dispatcher
// compiles into a fused kernel) and the filter (a separate operator). Each site holds its
diff --git a/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala b/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala
index 7efc5f12cb..71ec1f856e 100644
--- a/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala
@@ -102,6 +102,8 @@ class CometUdfBridgeSuite extends AnyFunSuite {
outArrays(0).memoryAddress(),
outSchemas(0).memoryAddress(),
numRows,
+ 0,
+ 0L,
null,
null)
@@ -172,6 +174,8 @@ class CometUdfBridgeSuite extends AnyFunSuite {
outArrays(0).memoryAddress(),
outSchemas(0).memoryAddress(),
numRows,
+ 0,
+ 0L,
null,
null)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org