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/07 15:58:01 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6709-1e8cf985fd3d9d3b1f8450ff73b749c3a625a527 created (now 9b54a0c7d8)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6709-1e8cf985fd3d9d3b1f8450ff73b749c3a625a527
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 9b54a0c7d8 perf: key the codegen dispatcher's kernel cache on a digest of the expression (#6709)
This branch includes the following new commits:
new 9b54a0c7d8 perf: key the codegen dispatcher's kernel cache on a digest of the expression (#6709)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: perf: key the codegen dispatcher's kernel cache on a digest of the expression (#6709)
Posted by gi...@apache.org.
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-6709-1e8cf985fd3d9d3b1f8450ff73b749c3a625a527
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 9b54a0c7d8cb7724edee9eba303561de0a51481d
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Oct 7 15:57:39 2026 +0000
perf: key the codegen dispatcher's kernel cache on a digest of the expression (#6709)
* perf: key the codegen dispatcher's kernel cache on a digest of the expression
Every batch copied the closure-serialized expression out of the
dispatcher's first argument, about 7 KB for a Scala UDF, and hashed the
copy byte by byte to find the kernel. The serde now computes a SHA-256
digest of the bytes once on the driver and sends it ahead of them. The
dispatcher looks the kernel up by the 32-byte digest and reads the
serialized expression only to compile on a cache miss.
Closes #6705.
* perf: build a JVM UDF's literal arguments once instead of for every batch
JvmScalarUdfExpr evaluated each literal argument for every batch, which
cloned its value and copied the clone into a one-row array. For the codegen
dispatcher that is two copies of the serialized expression per batch. The
arrays are now built when the expression is created and shared by every
batch.
* test: add a benchmark of a Scala UDF's per-batch cost on the codegen dispatcher
Runs SELECT max(f(c)) end to end at three batch sizes, and calls the
dispatcher directly next to the generated kernel it runs, so a cost the
dispatcher pays per batch shows up apart from the cost of the UDF itself.
* docs: say that a CometUDF must not modify its inputs
A literal argument is now the same vector on every batch of a task, so a
CometUDF that wrote into one would change what later batches see. A
column argument can already share its buffers with the rest of the plan.
CometUDF's scaladoc now says inputs are read-only.
* test: check that new children rebuild a JVM UDF's literal arrays
JvmScalarUdfExpr derives its literal arrays from its children. If
with_new_children kept the old arrays, a position that no longer held a
literal would go on sending the old literal's value.
* docs: say that a CometUDF's literal argument shares buffers, not a vector
CometUdfBridge.evaluate imports a new vector for every input on every
call and closes it when the call returns. What stays the same is the
native array behind a literal argument, which JvmScalarUdfExpr builds
once and sends with every batch, so its buffers last as long as the
expression rather than the task.
* test: check that the serde ships a digest of the serialized expression
---
native/spark-expr/src/jvm_udf/mod.rs | 70 ++++++-
.../org/apache/comet/serde/CometScalaUDF.scala | 21 ++-
.../main/scala/org/apache/comet/udf/CometUDF.scala | 3 +
.../comet/udf/codegen/CometScalaUDFCodegen.scala | 89 +++++----
.../org/apache/comet/CometCodegenSourceSuite.scala | 12 +-
.../scala/org/apache/comet/CometCodegenSuite.scala | 80 +++++++-
.../benchmark/CometScalaUDFDispatchBenchmark.scala | 204 +++++++++++++++++++++
.../sql/benchmark/CometTimeExtractBenchmark.scala | 7 +-
8 files changed, 432 insertions(+), 54 deletions(-)
diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs
index 0a37330097..6355b1a9da 100644
--- a/native/spark-expr/src/jvm_udf/mod.rs
+++ b/native/spark-expr/src/jvm_udf/mod.rs
@@ -26,6 +26,7 @@ use arrow::record_batch::RecordBatch;
use datafusion::common::Result as DFResult;
use datafusion::logical_expr::ColumnarValue;
+use datafusion::physical_expr::expressions::Literal;
use datafusion::physical_expr::PhysicalExpr;
use datafusion_comet_common::{decode_string_arrays, zero_offsets};
@@ -39,6 +40,12 @@ use jni::objects::{Global, JObject, JValue};
pub struct JvmScalarUdfExpr {
class_name: String,
args: Vec<Arc<dyn PhysicalExpr>>,
+ /// The length-1 array sent for each argument that is a literal. Built once here, because
+ /// building it for every batch copies the value twice: `Literal::evaluate` clones it and
+ /// `to_array_of_size` copies the clone. The codegen dispatcher's serialized expression is
+ /// such a literal, several KB for a Scala UDF. `None` for any other argument, and for a
+ /// literal whose array cannot be built, which `evaluate` then reports.
+ literal_arrays: Vec<Option<ArrayRef>>,
return_type: DataType,
return_nullable: bool,
/// Captured at `createPlan` time and threaded here by the planner. Passed through the
@@ -76,9 +83,17 @@ impl JvmScalarUdfExpr {
!class_name.is_empty(),
"JvmScalarUdfExpr requires a non-empty class name"
);
+ let literal_arrays = args
+ .iter()
+ .map(|arg| {
+ arg.downcast_ref::<Literal>()
+ .and_then(|literal| literal.value().to_array_of_size(1).ok())
+ })
+ .collect();
Self {
class_name,
args,
+ literal_arrays,
return_type,
return_nullable,
task_context,
@@ -143,9 +158,13 @@ impl PhysicalExpr for JvmScalarUdfExpr {
let arrays: Vec<ArrayRef> = self
.args
.iter()
- .map(|e| match e.evaluate(batch)? {
- ColumnarValue::Array(a) => Ok(a),
- ColumnarValue::Scalar(s) => s.to_array_of_size(1),
+ .zip(&self.literal_arrays)
+ .map(|(e, literal_array)| match literal_array {
+ Some(a) => Ok(Arc::clone(a)),
+ None => match e.evaluate(batch)? {
+ ColumnarValue::Array(a) => Ok(a),
+ ColumnarValue::Scalar(s) => s.to_array_of_size(1),
+ },
})
.collect::<DFResult<_>>()?;
@@ -292,3 +311,48 @@ impl PhysicalExpr for JvmScalarUdfExpr {
)))
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{Array, BinaryArray};
+ use datafusion::common::ScalarValue;
+ use datafusion::physical_expr::expressions::Column;
+
+ #[test]
+ fn builds_literal_arguments_once() {
+ let bytes = vec![7u8; 4096];
+ let udf = JvmScalarUdfExpr::new(
+ "org.example.Udf".to_string(),
+ vec![
+ Arc::new(Literal::new(ScalarValue::Binary(Some(bytes.clone())))),
+ Arc::new(Column::new("a", 0)),
+ ],
+ DataType::Int64,
+ true,
+ None,
+ None,
+ 0,
+ 0,
+ );
+ let literal = udf.literal_arrays[0]
+ .as_ref()
+ .expect("a literal argument's array is built up front");
+ let literal = literal.as_any().downcast_ref::<BinaryArray>().unwrap();
+ assert_eq!(literal.len(), 1);
+ assert_eq!(literal.value(0), bytes.as_slice());
+ assert!(udf.literal_arrays[1].is_none());
+
+ // New children rebuild the arrays, so a position that no longer holds a literal does not
+ // keep sending the old literal's array.
+ let swapped = Arc::new(udf)
+ .with_new_children(vec![
+ Arc::new(Column::new("a", 0)),
+ Arc::new(Literal::new(ScalarValue::Binary(Some(bytes)))),
+ ])
+ .unwrap();
+ let swapped = swapped.downcast_ref::<JvmScalarUdfExpr>().unwrap();
+ assert!(swapped.literal_arrays[0].is_none());
+ assert!(swapped.literal_arrays[1].is_some());
+ }
+}
diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala
index 58caedcff1..c79b4839bb 100644
--- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala
@@ -117,7 +117,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
}
// Serialize via Spark's closure serializer: respects the task context classloader (so user
- // UDF jars are visible) and matches Spark's wire format. The bytes become arg 0 of the
+ // UDF jars are visible) and matches Spark's wire format. The bytes become arg 1 of the
// JvmScalarUdf proto and self-describe the expression so this works in cluster mode without
// executor-side driver registry state.
//
@@ -145,12 +145,17 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
s"(${e.getClass.getSimpleName}: ${e.getMessage})")
return None
}
- val exprArg = exprToProtoInternal(Literal(bytes, BinaryType), inputs, binding).getOrElse {
- withFallbackReason(
- expr,
- s"$exprName: codegen dispatch: could not serialize closure-serialized bound " +
- "expression payload")
- return None
+ // Arg 0 is a digest of the bytes. The dispatcher finds each batch's kernel by it, so it reads
+ // the bytes only to compile on a cache miss instead of copying and hashing them for every
+ // batch.
+ val exprArgs = Seq(CometScalaUDFCodegen.digest(bytes), bytes).map { payload =>
+ exprToProtoInternal(Literal(payload, BinaryType), inputs, binding).getOrElse {
+ withFallbackReason(
+ expr,
+ s"$exprName: codegen dispatch: could not serialize closure-serialized bound " +
+ "expression payload")
+ return None
+ }
}
val dataArgs = attrs.map { a =>
@@ -169,7 +174,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
val udfBuilder = ExprOuterClass.JvmScalarUdf
.newBuilder()
.setClassName(classOf[CometScalaUDFCodegen].getName)
- .addArgs(exprArg)
+ exprArgs.foreach(udfBuilder.addArgs)
dataArgs.foreach(udfBuilder.addArgs)
udfBuilder
.setReturnType(returnTypeProto)
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 db9178fc32..a1d7293991 100644
--- a/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala
+++ b/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala
@@ -28,6 +28,9 @@ import org.apache.arrow.vector.ValueVector
* - Vector arguments arrive at the row count of the current batch.
* - Scalar (literal-folded) arguments arrive as length-1 vectors and must be read at index 0.
* - The returned vector's length must match `numRows`.
+ * - Inputs belong to native execution and must not be modified. Each call receives new vectors,
+ * but a scalar argument's buffers are the same on every call for the life of the expression,
+ * and a vector argument can share its buffers with the rest of the plan.
*
* `numRows` mirrors DataFusion's `ScalarFunctionArgs.number_rows` and is the batch row count.
* UDFs that always have at least one batch-length input can read length from it and ignore
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 c589fb6fa9..11b2594fdd 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
@@ -20,6 +20,7 @@
package org.apache.comet.udf.codegen
import java.nio.ByteBuffer
+import java.security.MessageDigest
import java.util.Collections
import java.util.concurrent.atomic.AtomicLong
@@ -44,8 +45,11 @@ import org.apache.comet.udf.CometUDF
* 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.
+ * Arg 0 is a `VarBinaryVector` scalar carrying the [[CometScalaUDFCodegen.digest]] of the
+ * closure-serialized bound `Expression`; arg 1 is a `VarBinaryVector` scalar carrying the
+ * serialized bytes themselves; args 2..N are the data columns the `BoundReference`s read in
+ * ordinal order. A batch finds its kernel by the digest, so the serialized expression, several KB
+ * for a Scala UDF's closure, is read only to compile a kernel on a cache miss.
*
* Caching hierarchy, broadest scope on the left:
* {{{
@@ -82,13 +86,13 @@ import org.apache.comet.udf.CometUDF
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. 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`.
+ * Per-task cache keyed on the serialized expression's digest 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. 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] =
@@ -114,25 +118,25 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
partitionIndex: Int,
planId: Long): ValueVector = {
require(
- inputs.length >= 1,
- "CometScalaUDFCodegen requires at least 1 input (serialized expression), " +
- s"got ${inputs.length}")
- val exprVec = inputs(0).asInstanceOf[VarBinaryVector]
+ inputs.length >= 2,
+ "CometScalaUDFCodegen requires at least 2 inputs (expression digest and serialized " +
+ s"expression), got ${inputs.length}")
+ val digest = binaryScalar(inputs(0), "expression digest at arg 0")
require(
- exprVec.getValueCount >= 1 && !exprVec.isNull(0),
- "CometScalaUDFCodegen requires non-null serialized expression bytes at arg 0")
- val bytes = exprVec.get(0)
+ digest.length == CometScalaUDFCodegen.DigestLength,
+ s"CometScalaUDFCodegen requires a ${CometScalaUDFCodegen.DigestLength}-byte expression " +
+ s"digest at arg 0, got ${digest.length} bytes")
// TODO(dict-encoded): kernels assume materialized inputs. Dict-encoded vectors would fail the
// cast in `specFor` below. Fix is to materialize at the dispatcher (via
// `CDataDictionaryProvider`) or widen `emitTypedGetters` with a dict-index + lookup path.
- val numDataCols = inputs.length - 1
+ val numDataCols = inputs.length - 2
val dataCols = new Array[ValueVector](numDataCols)
val specs = new Array[ArrowColumnSpec](numDataCols)
var di = 0
while (di < numDataCols) {
- val v = inputs(di + 1)
+ val v = inputs(di + 2)
dataCols(di) = v
specs(di) = specFor(v)
di += 1
@@ -140,12 +144,12 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
val n = numRows
val specsSeq = specs.toIndexedSeq
- val key = CometScalaUDFCodegen.CacheKey(planId, ByteBuffer.wrap(bytes), specsSeq)
+ val key = CometScalaUDFCodegen.CacheKey(planId, ByteBuffer.wrap(digest), 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, partitionIndex)
+ val entry = lookupOrCompile(key, inputs(1), specsSeq, partitionIndex)
val out = CometBatchKernelCodegen.allocateOutput(
entry.outputField,
@@ -168,7 +172,7 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
private def lookupOrCompile(
key: CometScalaUDFCodegen.CacheKey,
- bytes: Array[Byte],
+ exprVec: ValueVector,
specs: IndexedSeq[ArrowColumnSpec],
partitionIndex: Int): CometScalaUDFCodegen.CacheEntry = {
assert(Thread.holdsLock(this), "lookupOrCompile must run under this.synchronized")
@@ -181,6 +185,7 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
CometScalaUDFCodegen.cacheHitCount.incrementAndGet()
entry
case None =>
+ val bytes = binaryScalar(exprVec, "serialized expression at arg 1")
val loader = Option(Thread.currentThread().getContextClassLoader)
.getOrElse(classOf[Expression].getClassLoader)
val boundExpr =
@@ -215,6 +220,15 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
}
}
+ /** The value of a binary scalar argument, which arrives as a length-1 vector. */
+ private def binaryScalar(v: ValueVector, what: String): Array[Byte] = {
+ val vec = v.asInstanceOf[VarBinaryVector]
+ require(
+ vec.getValueCount >= 1 && !vec.isNull(0),
+ s"CometScalaUDFCodegen requires a non-null $what")
+ vec.get(0)
+ }
+
/**
* Build the compile-time spec for one input Arrow vector. Recurses on complex types.
*
@@ -225,9 +239,10 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
*
* Top-level `nullable=true` is hardcoded: the cache key does not specialize on per-batch null
* density. Schema-declared nullability still reaches the kernel via `BoundReference.nullable`
- * embedded in `bytesKey`, so `BoundReference.doGenCode` elides its own `isNullAt` probe on
- * non-null columns. `StructFieldSpec.nullable` reads `field.isNullable` from Arrow metadata,
- * which is a schema property and therefore stable across batches.
+ * embedded in the serialized expression, which the key's digest covers, so
+ * `BoundReference.doGenCode` elides its own `isNullAt` probe on non-null columns.
+ * `StructFieldSpec.nullable` reads `field.isNullable` from Arrow metadata, which is a schema
+ * property and therefore stable across batches.
*/
private def specFor(v: ValueVector): ArrowColumnSpec = v match {
case map: MapVector =>
@@ -337,18 +352,24 @@ object CometScalaUDFCodegen {
*/
private[comet] val NoPlan = -1L
+ /** Length in bytes of a [[digest]]. */
+ val DigestLength = 32
+
/**
- * 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.
+ * Identifies a closure-serialized bound expression in the kernel cache. The serde computes it
+ * once on the driver and ships it as arg 0, next to the bytes, so the per-batch lookup hashes
+ * 32 bytes instead of the whole serialized expression. SHA-256 rather than a cheaper hash,
+ * because a hit is trusted without comparing the bytes: two expressions whose digests collided
+ * would share a kernel.
+ */
+ def digest(serialized: Array[Byte]): Array[Byte] =
+ MessageDigest.getInstance("SHA-256").digest(serialized)
+
+ /**
+ * Cache key: calling native plan (`NoPlan` for a deterministic expression), the bound
+ * expression's [[digest]] plus per-column compile-time invariants.
*/
- final case class CacheKey(
- planId: Long,
- bytesKey: ByteBuffer,
- specs: IndexedSeq[ArrowColumnSpec])
+ final case class CacheKey(planId: Long, digest: 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 f1272e151e..54b59dec9d 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala
@@ -1436,8 +1436,8 @@ class CometCodegenSourceSuite extends AnyFunSuite {
}
test("closure-serialized bytes diverge for failOnError / roundOff / timeZoneId variants") {
- // The dispatcher caches kernels by the closure-serialized bytes of the bound expression
- // (`CacheKey.bytesKey`). For expression classes that carry a runtime-dependent boolean
+ // The dispatcher caches kernels by a digest of the closure-serialized bytes of the bound
+ // expression (`CacheKey.digest`). For expression classes that carry a runtime-dependent boolean
// (failOnError, roundOff) or string (timeZoneId), two plan instances differing only in that
// field must serialize to distinct byte sequences so they receive distinct cache entries.
// A collision would let a kernel compiled for one variant (e.g. ANSI throw site) silently
@@ -1523,20 +1523,20 @@ class CometCodegenSourceSuite extends AnyFunSuite {
}
test("CacheKey discriminates on ArrowColumnSpec.nullable") {
- // Structural regression: same expression bytes and same Arrow vector class with different
+ // Structural regression: same expression digest and same Arrow vector class with different
// `nullable` must produce non-equal cache keys. The dispatcher today hardcodes `nullable=true`
// for top-level specs, so the two variants don't both arise from runtime data, but the case
// class equality contract still has to discriminate so that any future tiered cache or test
// construction can rely on it. The non-nullable variant's generated source emits a literal
// `false` from `isNullAt`, distinct codegen output that we never want to silently share with
// the nullable variant.
- val bytes = java.nio.ByteBuffer.wrap(Array[Byte](1, 2, 3))
+ val digest = java.nio.ByteBuffer.wrap(Array[Byte](1, 2, 3))
val nullable =
IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = true))
val nonNullable =
IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = false))
- val k1 = CometScalaUDFCodegen.CacheKey(0L, bytes, nullable)
- val k2 = CometScalaUDFCodegen.CacheKey(0L, bytes, nonNullable)
+ val k1 = CometScalaUDFCodegen.CacheKey(0L, digest, nullable)
+ val k2 = CometScalaUDFCodegen.CacheKey(0L, digest, 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 5a80865da9..7af83a5494 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
@@ -624,6 +624,30 @@ class CometCodegenSuite
.contains(Set("hypot", "cast", "checkoverflow", "add")))
}
+ test("the serde ships a digest of the serialized expression at arg 0 (#6705)") {
+ // The dispatcher trusts a cache hit on the digest without comparing the bytes, so the digest
+ // the serde ships must be the digest of the bytes it ships next to it. Calling the dispatcher
+ // directly cannot catch a mismatch, because those tests build their own digest.
+ val x = AttributeReference("x", DoubleType, nullable = false)()
+ def payload(e: Expression): (Array[Byte], Array[Byte]) = {
+ val proto = QueryPlanSerde.exprToProto(e, Seq(x)).get
+ assert(proto.hasJvmScalarUdf)
+ val args = proto.getJvmScalarUdf.getArgsList
+ (
+ args.get(0).getLiteral.getBytesVal.toByteArray,
+ args.get(1).getLiteral.getBytesVal.toByteArray)
+ }
+
+ val (hypotDigest, hypotBytes) = payload(Hypot(x, Literal(4.0d)))
+ val (otherDigest, otherBytes) = payload(Hypot(x, Literal(5.0d)))
+ assert(hypotDigest.sameElements(CometScalaUDFCodegen.digest(hypotBytes)))
+ assert(otherDigest.sameElements(CometScalaUDFCodegen.digest(otherBytes)))
+
+ // Expressions that differ only in a literal must not share a kernel.
+ assert(!hypotBytes.sameElements(otherBytes))
+ assert(!hypotDigest.sameElements(otherDigest))
+ }
+
test("tags copied onto the shared TrueLiteral do not leak into unrelated plans") {
// Catalyst copies a rewritten node's tags onto its replacement, so a tagged expression that an
// earlier query rewrote into `Literal.TrueLiteral` brands that process-wide singleton for the
@@ -866,7 +890,7 @@ class CometCodegenSuite
"same UDF over nullable and non-nullable columns gets distinct kernels with independent state") {
// Two columns, same type, different schema-declared nullability. Same UDF applied to each
// alongside a per-projection MonotonicallyIncreasingID. Each projection has its own MII
- // child (different bytesKey), so each kernel must have its own counter advancing 0..N-1.
+ // child (a different digest), so each kernel must have its own counter advancing 0..N-1.
// If the dispatcher collapses them onto one kernel or shares state somehow, the counters
// would interleave and the output would diverge from Spark.
spark.udf.register("withId", (s: String, id: Long) => s"${s}_${id}")
@@ -2329,6 +2353,7 @@ class CometCodegenSuite
// because Spark 4.1 still rejects TIME columns in file-based data sources, so no SQL query
// can produce a TIME input today.
val timeVec = new TimeNanoVector("tm", CometArrowAllocator)
+ val digestVec = new VarBinaryVector("digest", CometArrowAllocator)
val exprVec = new VarBinaryVector("expr", CometArrowAllocator)
var out: ValueVector = null
try {
@@ -2342,21 +2367,72 @@ class CometCodegenSuite
val serialized = SparkEnv.get.closureSerializer.newInstance().serialize(expr)
val bytes = new Array[Byte](serialized.remaining())
serialized.get(bytes)
+ digestVec.allocateNew()
+ digestVec.setSafe(0, CometScalaUDFCodegen.digest(bytes))
+ digestVec.setValueCount(1)
exprVec.allocateNew()
exprVec.setSafe(0, bytes)
exprVec.setValueCount(1)
- out = new CometScalaUDFCodegen().evaluate(Array(exprVec, timeVec), 2)
+ out = new CometScalaUDFCodegen().evaluate(Array(digestVec, exprVec, timeVec), 2)
val comet = CometVector.getVector(out.asInstanceOf[FieldVector], null)
assert(comet.getLong(0) === 45296000000000L)
assert(comet.isNullAt(1))
} finally {
if (out != null) out.close()
+ digestVec.close()
exprVec.close()
timeVec.close()
}
}
+ test("dispatcher finds a compiled kernel by the expression digest alone (#6705)") {
+ // The serde ships a digest of the serialized expression at arg 0 and the bytes at arg 1, and
+ // the dispatcher reads the bytes only to compile on a cache miss. The second call passes a
+ // null at arg 1, so it succeeds only if the digest finds the kernel the first call compiled.
+ val expr = Add(BoundReference(0, LongType, nullable = true), Literal(1L))
+ val serialized = SparkEnv.get.closureSerializer.newInstance().serialize(expr)
+ val bytes = new Array[Byte](serialized.remaining())
+ serialized.get(bytes)
+
+ def binaryScalar(name: String, value: Array[Byte]): VarBinaryVector = {
+ val v = new VarBinaryVector(name, CometArrowAllocator)
+ v.allocateNew()
+ if (value == null) v.setNull(0) else v.setSafe(0, value)
+ v.setValueCount(1)
+ v
+ }
+ val digestVec = binaryScalar("digest", CometScalaUDFCodegen.digest(bytes))
+ val exprVec = binaryScalar("expr", bytes)
+ val nullExprVec = binaryScalar("expr", null)
+ val input = new BigIntVector("x", CometArrowAllocator)
+ try {
+ input.allocateNew(2)
+ input.set(0, 41L)
+ input.setNull(1)
+ input.setValueCount(2)
+
+ val dispatcher = new CometScalaUDFCodegen()
+ Seq(exprVec, nullExprVec).foreach { arg1 =>
+ val out =
+ dispatcher.evaluate(Array(digestVec, arg1, input), 2).asInstanceOf[BigIntVector]
+ try {
+ assert(out.get(0) === 42L)
+ assert(out.isNull(1))
+ } finally out.close()
+ }
+
+ // The layout before the digest, with the serialized expression at arg 0, is refused
+ // rather than taken for a digest.
+ val e = intercept[IllegalArgumentException] {
+ dispatcher.evaluate(Array(exprVec, input), 2)
+ }
+ assert(e.getMessage.contains("expression digest"), e.getMessage)
+ } finally {
+ Seq(digestVec, exprVec, nullExprVec, input).foreach(_.close())
+ }
+ }
+
// Runtime coverage for nullable nested `getStruct` / `getArray` / `getMap` element reads is
// exercised through HOFs in `CometCodegenHOFSuite`. Static emitter assertions live in
// `CometCodegenSourceSuite`.
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala
new file mode 100644
index 0000000000..91a9337f26
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala
@@ -0,0 +1,204 @@
+/*
+ * 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.benchmark
+
+import org.apache.arrow.vector.{BigIntVector, ValueVector, VarBinaryVector}
+import org.apache.spark.SparkEnv
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSeq, BindReferences, Expression, ScalaUDF}
+import org.apache.spark.sql.execution.adaptive.QueryStageExec
+
+import org.apache.comet.{CometArrowAllocator, CometConf}
+import org.apache.comet.codegen.CometBatchKernelCodegen
+import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec
+import org.apache.comet.udf.codegen.CometScalaUDFCodegen
+
+/**
+ * What a row-level Scala UDF costs on the codegen dispatcher, split into the cost paid per batch
+ * and the cost paid per row.
+ *
+ * Part 1 runs `SELECT max(f(c))` over a Parquet `bigint` column end to end at several batch
+ * sizes. The total of a cost paid once per batch grows as batches shrink, and the total of a cost
+ * paid per row does not. A tenth of the rows are null. Spark wraps a UDF with a primitive
+ * parameter in a null check, `if(isnull(c), null, f(knownnotnull(c)))`, which runs natively
+ * around the dispatched call; a UDF with a boxed parameter gets no such wrapper and converts its
+ * input through Spark's encoder instead.
+ *
+ * Part 2 calls the dispatcher directly on the same in-memory batch over and over, with no Spark
+ * query and no native code, and does the same with the generated kernel it runs. The difference
+ * between the two is the dispatcher's own work for each batch: reading its arguments and finding
+ * the kernel.
+ *
+ * To run this benchmark:
+ * {{{
+ * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometScalaUDFDispatchBenchmark
+ * }}}
+ * Results will be written to "spark/benchmarks/CometScalaUDFDispatchBenchmark-**results.txt".
+ */
+object CometScalaUDFDispatchBenchmark extends CometBenchmarkBase {
+
+ private val Rows = 4L * 1024 * 1024
+ private val BatchSizes = Seq(1024, 8192, 65536)
+ private val LayerBatchRows = Seq(1024, 8192)
+
+ override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+ spark.udf.register("prim_add_one", (x: Long) => x + 1)
+ spark.udf.register(
+ "boxed_add_one",
+ (x: java.lang.Long) => if (x == null) null else java.lang.Long.valueOf(x + 1))
+
+ withTempPath { dir =>
+ withTempTable(tbl, "parquetV1Table") {
+ spark.range(Rows).createOrReplaceTempView(tbl)
+ prepareTable(dir, spark.sql(s"SELECT IF(id % 10 = 0, NULL, id * 7919) AS c FROM $tbl"))
+
+ checkPlans()
+ BatchSizes.foreach(endToEnd)
+ LayerBatchRows.foreach(jvmLayers)
+ }
+ }
+ }
+
+ /** `dispatch` turns the dispatcher on; the cases without a UDF leave it off. */
+ private case class EndToEndCase(label: String, f: String, dispatch: Boolean)
+
+ private val endToEndCases = Seq(
+ EndToEndCase("dispatch, primitive param", "prim_add_one(c)", dispatch = true),
+ EndToEndCase("dispatch, boxed param", "boxed_add_one(c)", dispatch = true),
+ EndToEndCase("native c + 1", "c + 1", dispatch = false),
+ EndToEndCase("no function, max(c)", "c", dispatch = false))
+
+ private def configs(dispatch: Boolean, batchSize: Int): Seq[(String, String)] = Seq(
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_BATCH_SIZE.key -> batchSize.toString,
+ CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString)
+
+ private def query(f: String): String = s"SELECT max($f) FROM parquetV1Table"
+
+ /** Fails unless Comet runs every operator of every case. */
+ private def checkPlans(): Unit = endToEndCases.foreach {
+ case EndToEndCase(label, f, dispatch) =>
+ withSQLConf(configs(dispatch, BatchSizes.head): _*) {
+ val df = spark.sql(query(f))
+ df.noop()
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ // A query stage is a leaf, so check the plan inside each one separately.
+ val stagePlans = collect(plan) { case s: QueryStageExec => s.plan }
+ (plan +: stagePlans)
+ .flatMap(findFirstNonCometOperator(_, classOf[QueryStageExec]))
+ .headOption
+ .foreach { op =>
+ throw new IllegalStateException(s"$label ran ${op.nodeName} in Spark:\n$plan")
+ }
+ }
+ }
+
+ private def endToEnd(batchSize: Int): Unit = {
+ val benchmark =
+ new Benchmark(s"max(f(c)), batch size $batchSize", Rows, minNumIters = 5, output = output)
+ endToEndCases.foreach { case EndToEndCase(label, f, dispatch) =>
+ benchmark.addCase(label) { _ =>
+ withSQLConf(configs(dispatch, batchSize): _*) {
+ spark.sql(query(f)).noop()
+ }
+ }
+ }
+ benchmark.run()
+ }
+
+ /** The bound UDF subtree the serde would dispatch, and its closure-serialized bytes. */
+ private def boundUdf(f: String): (Expression, Array[Byte]) = {
+ val plan = spark.sql(s"SELECT $f FROM parquetV1Table").queryExecution.optimizedPlan
+ val udf = plan.expressions.flatMap(_.collect { case u: ScalaUDF => u }).head
+ val attrs = udf.collect { case a: AttributeReference => a }.distinct
+ val bound = BindReferences.bindReference(udf: Expression, AttributeSeq(attrs))
+ val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(bound)
+ val bytes = new Array[Byte](buffer.remaining())
+ buffer.get(bytes)
+ (bound, bytes)
+ }
+
+ private def binaryScalar(value: Array[Byte]): VarBinaryVector = {
+ val v = new VarBinaryVector("arg", CometArrowAllocator)
+ v.allocateNew()
+ v.setSafe(0, value)
+ v.setValueCount(1)
+ v
+ }
+
+ /** The arguments the serde ships ahead of the data columns. */
+ private def dispatcherArgs(bytes: Array[Byte]): Seq[VarBinaryVector] =
+ Seq(CometScalaUDFCodegen.digest(bytes), bytes).map(binaryScalar)
+
+ private def jvmLayers(batchRows: Int): Unit = {
+ val in = new BigIntVector("c", CometArrowAllocator)
+ in.allocateNew(batchRows)
+ var i = 0
+ while (i < batchRows) {
+ if (i % 10 == 0) in.setNull(i) else in.set(i, i * 7919L)
+ i += 1
+ }
+ in.setValueCount(batchRows)
+ val inputs: Array[ValueVector] = Array(in)
+ val batchesPerIter = (Rows / batchRows).toInt
+
+ val benchmark =
+ new Benchmark(s"JVM layers, $batchRows-row batches", Rows, minNumIters = 5, output = output)
+
+ def addBatchCase(label: String)(evaluate: () => ValueVector): Unit =
+ benchmark.addCase(label) { _ =>
+ var b = 0
+ while (b < batchesPerIter) {
+ evaluate().close()
+ b += 1
+ }
+ }
+
+ val toClose = Seq("prim_add_one(c)" -> "primitive param", "boxed_add_one(c)" -> "boxed param")
+ .flatMap { case (f, param) =>
+ val (bound, bytes) = boundUdf(f)
+ val specs = IndexedSeq(ArrowColumnSpec(classOf[BigIntVector], nullable = true))
+ val kernel = CometBatchKernelCodegen.compile(bound, specs).newInstance()
+ kernel.init(0)
+ val field = CometBatchKernelCodegen.toFfiArrowField("r", bound.dataType, bound.nullable)
+ addBatchCase(s"generated kernel only, $param") { () =>
+ val out = CometBatchKernelCodegen.allocateOutput(field, batchRows, 0)
+ kernel.process(inputs, out, batchRows)
+ out.setValueCount(batchRows)
+ out
+ }
+
+ val leading = dispatcherArgs(bytes)
+ val dispatcher = new CometScalaUDFCodegen
+ val dispatchInputs: Array[ValueVector] = (leading :+ in).toArray
+ addBatchCase(s"dispatcher evaluate, $param (${bytes.length}-byte expression)") { () =>
+ dispatcher.evaluate(dispatchInputs, batchRows)
+ }
+ leading
+ }
+
+ try benchmark.run()
+ finally {
+ toClose.foreach(_.close())
+ in.close()
+ }
+ }
+}
diff --git a/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala b/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala
index 025fe43d45..38d4ac8b54 100644
--- a/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala
+++ b/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala
@@ -85,6 +85,7 @@ object CometTimeExtractBenchmark extends CometBenchmarkBase {
private def runKernels(precision: Int, nulls: Boolean): Unit = {
val size = 10000
val time = new TimeNanoVector("t", CometArrowAllocator)
+ val digest = new VarBinaryVector("digest", CometArrowAllocator)
val serialized = new VarBinaryVector("expr", CometArrowAllocator)
try {
time.allocateNew(size)
@@ -114,11 +115,14 @@ object CometTimeExtractBenchmark extends CometBenchmarkBase {
val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(expr)
val bytes = new Array[Byte](buffer.remaining())
buffer.get(bytes)
+ digest.allocateNew()
+ digest.setSafe(0, CometScalaUDFCodegen.digest(bytes))
+ digest.setValueCount(1)
serialized.allocateNew()
serialized.setSafe(0, bytes)
serialized.setValueCount(1)
val dispatcher = new CometScalaUDFCodegen()
- val inputs = Array[org.apache.arrow.vector.ValueVector](serialized, time)
+ val inputs = Array[org.apache.arrow.vector.ValueVector](digest, serialized, time)
val result = dispatcher.evaluate(inputs, size)
try {
val vector =
@@ -170,6 +174,7 @@ object CometTimeExtractBenchmark extends CometBenchmarkBase {
}
}
} finally {
+ digest.close()
serialized.close()
time.close()
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org