perf: key the codegen dispatcher's kernel cache on a digest of the expression

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());
+    }
+}

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)

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

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,38 +118,38 @@ 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
     }
     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) {

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