diff --git a/.github/workflows/pr_benchmark_check.yml b/.github/workflows/pr_benchmark_check.yml
index b07cc03c34..a879493a7f 100644
--- a/.github/workflows/pr_benchmark_check.yml
+++ b/.github/workflows/pr_benchmark_check.yml
@@ -84,9 +84,7 @@ jobs:
${{ runner.os }}-benchmark-maven-
- name: Check Scala compilation and linting
- # Pin to spark-4.0 (Scala 2.13.16) because the default profile is now
- # spark-4.1 / Scala 2.13.17, and semanticdb-scalac_2.13.17 is not yet
- # published, which breaks `-Psemanticdb`. See pr_build_linux.yml for
- # the same exclusion in the main lint matrix.
+ # Pinned to spark-4.0 because semanticdb-scalac_2.13.17 (spark-4.1 default)
+ # is not yet published, which breaks the -Psemanticdb scalafix lint.
run: |
- ./mvnw -B compile test-compile scalafix:scalafix -Dscalafix.mode=CHECK -Psemanticdb -Pspark-4.0 -DskipTests
+ ./mvnw -B compile test-compile scalafix:scalafix -Dscalafix.mode=CHECK -Pspark-4.0 -Psemanticdb -DskipTests
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index 5c1ae2dc47..6e6a526f71 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -309,6 +309,7 @@ jobs:
org.apache.comet.CometFuzzAggregateSuite
org.apache.comet.CometFuzzIcebergSuite
org.apache.comet.CometFuzzMathSuite
+ org.apache.comet.CometCodegenDispatchFuzzSuite
org.apache.comet.DataGeneratorSuite
- name: "shuffle"
value: |
@@ -386,6 +387,8 @@ jobs:
org.apache.comet.expressions.conditional.CometIfSuite
org.apache.comet.expressions.conditional.CometCoalesceSuite
org.apache.comet.expressions.conditional.CometCaseWhenSuite
+ org.apache.comet.CometCodegenDispatchSmokeSuite
+ org.apache.comet.CometCodegenSourceSuite
- name: "sql"
value: |
org.apache.spark.sql.CometToPrettyStringSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 29eca594a2..901a3b7f5c 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -157,6 +157,7 @@ jobs:
org.apache.comet.CometFuzzAggregateSuite
org.apache.comet.CometFuzzIcebergSuite
org.apache.comet.CometFuzzMathSuite
+ org.apache.comet.CometCodegenDispatchFuzzSuite
org.apache.comet.DataGeneratorSuite
- name: "shuffle"
value: |
@@ -233,6 +234,8 @@ jobs:
org.apache.comet.expressions.conditional.CometIfSuite
org.apache.comet.expressions.conditional.CometCoalesceSuite
org.apache.comet.expressions.conditional.CometCaseWhenSuite
+ org.apache.comet.CometCodegenDispatchSmokeSuite
+ org.apache.comet.CometCodegenSourceSuite
- name: "sql"
value: |
org.apache.spark.sql.CometToPrettyStringSuite
diff --git a/common/src/main/java/org/apache/comet/udf/CometBatchKernel.java b/common/src/main/java/org/apache/comet/udf/CometBatchKernel.java
new file mode 100644
index 0000000000..cfa61c9715
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/udf/CometBatchKernel.java
@@ -0,0 +1,68 @@
+/*
+ * 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.udf;
+
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.ValueVector;
+
+/**
+ * Abstract base extended by the Janino-compiled batch kernel emitted by {@code
+ * CometBatchKernelCodegen}. The generated subclass extends {@code CometInternalRow} (so Spark's
+ * {@code BoundReference.genCode} can call {@code this.getUTF8String(ord)} directly) and carries
+ * typed input fields baked at codegen time, one per input column. Expression evaluation plus Arrow
+ * read/write fuse into one method per expression tree.
+ *
+ * Input scope: any {@code ValueVector[]}; the generated subclass casts each slot to the concrete
+ * Arrow type the compile-time schema specified. Output is a generic {@code FieldVector}; the
+ * generated subclass casts to the concrete type matching the bound expression's {@code dataType}.
+ * Widen input support by adding vector classes to the getter switch in {@code
+ * CometBatchKernelCodegen.emitTypedGetters}; widen output support by adding cases in {@code
+ * CometBatchKernelCodegen.allocateOutput} and {@code emitOutputWriter}.
+ */
+public abstract class CometBatchKernel extends CometInternalRow {
+
+ protected final Object[] references;
+
+ protected CometBatchKernel(Object[] references) {
+ this.references = references;
+ }
+
+ /**
+ * Process one batch.
+ *
+ * @param inputs Arrow input vectors; length and concrete classes must match the schema the kernel
+ * was compiled against
+ * @param output Arrow output vector; caller allocates to the expression's {@code dataType}
+ * @param numRows number of rows in this batch
+ */
+ public abstract void process(ValueVector[] inputs, FieldVector output, int numRows);
+
+ /**
+ * Run partition-dependent initialization. The generated subclass overrides this to execute
+ * statements collected via {@code CodegenContext.addPartitionInitializationStatement}, for
+ * example reseeding {@code Rand}'s {@code XORShiftRandom} from {@code seed + partitionIndex}.
+ * Deterministic expressions leave this as a no-op.
+ *
+ *
The caller must invoke this before the first {@code process} call of each partition. The
+ * generated subclass is not thread-safe across concurrent {@code process} calls, so kernels are
+ * allocated per dispatcher invocation and init is run once on the fresh instance.
+ */
+ public void init(int partitionIndex) {}
+}
diff --git a/common/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/common/src/main/java/org/apache/comet/udf/CometUdfBridge.java
index aed53c57df..4e8662829f 100644
--- a/common/src/main/java/org/apache/comet/udf/CometUdfBridge.java
+++ b/common/src/main/java/org/apache/comet/udf/CometUdfBridge.java
@@ -19,7 +19,8 @@
package org.apache.comet.udf;
-import java.util.concurrent.ConcurrentHashMap;
+import java.util.LinkedHashMap;
+import java.util.Map;
import org.apache.arrow.c.ArrowArray;
import org.apache.arrow.c.ArrowSchema;
@@ -27,6 +28,8 @@
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.ValueVector;
+import org.apache.spark.TaskContext;
+import org.apache.spark.comet.CometTaskContextShim;
/**
* JNI entry point for native execution to invoke a {@link CometUDF}. Matches the static-method
@@ -35,10 +38,23 @@
*/
public class CometUdfBridge {
- // Process-wide cache of UDF instances keyed by class name. CometUDF
- // implementations are required to be stateless (see CometUDF), so a
- // single shared instance per class is safe across native worker threads.
- private static final ConcurrentHashMap INSTANCES = new ConcurrentHashMap<>();
+ // Per-thread, bounded LRU of UDF instances keyed by class name. Comet
+ // native execution threads (Tokio/DataFusion worker pool) are reused
+ // across tasks within an executor, so the effective lifetime of cached
+ // entries is the worker thread (i.e. the executor JVM). Fine for
+ // stateless UDFs; future stateful UDFs would need explicit per-task
+ // isolation.
+ private static final int CACHE_CAPACITY = 64;
+
+ private static final ThreadLocal> INSTANCES =
+ ThreadLocal.withInitial(
+ () ->
+ new LinkedHashMap(CACHE_CAPACITY, 0.75f, true) {
+ @Override
+ protected boolean removeEldestEntry(Map.Entry eldest) {
+ return size() > CACHE_CAPACITY;
+ }
+ });
/**
* Called from native via JNI.
@@ -48,30 +64,67 @@ public class CometUdfBridge {
* @param inputSchemaPtrs addresses of pre-allocated FFI_ArrowSchema structs (one per input)
* @param outArrayPtr address of pre-allocated FFI_ArrowArray for the result
* @param outSchemaPtr address of pre-allocated FFI_ArrowSchema for the result
+ * @param numRows number of rows in the current batch. Mirrors DataFusion's {@code
+ * ScalarFunctionArgs.number_rows} and gives UDFs an explicit batch-size signal for cases
+ * where no input arg is a batch-length array (e.g. a zero-arg non-deterministic ScalaUDF).
+ * UDFs that already read size from their input vectors can ignore it.
+ * @param taskContext Spark {@link TaskContext} captured on the driving Spark task thread and
+ * passed through from native. May be {@code null} when the bridge is invoked outside a Spark
+ * task (unit tests, direct native driver runs). When non-null and the current thread has no
+ * {@code TaskContext} of its own, the bridge installs it as the thread-local for the duration
+ * of the UDF call so the UDF body (including partition-sensitive built-ins like {@code Rand}
+ * / {@code Uuid} / {@code MonotonicallyIncreasingID} that read the partition index via {@code
+ * TaskContext.get().partitionId()}) sees the real context rather than null. The thread-local
+ * is cleared in a {@code finally} so Tokio workers don't leak a stale TaskContext across
+ * invocations.
*/
public static void evaluate(
String udfClassName,
long[] inputArrayPtrs,
long[] inputSchemaPtrs,
long outArrayPtr,
- long outSchemaPtr) {
- CometUDF udf =
- INSTANCES.computeIfAbsent(
- udfClassName,
- name -> {
- try {
- // Resolve via the executor's context classloader so user-supplied UDF jars
- // (added via spark.jars / --jars) are visible.
- ClassLoader cl = Thread.currentThread().getContextClassLoader();
- if (cl == null) {
- cl = CometUdfBridge.class.getClassLoader();
- }
- return (CometUDF)
- Class.forName(name, true, cl).getDeclaredConstructor().newInstance();
- } catch (ReflectiveOperationException e) {
- throw new RuntimeException("Failed to instantiate CometUDF: " + name, e);
- }
- });
+ long outSchemaPtr,
+ int numRows,
+ TaskContext taskContext) {
+ boolean installedTaskContext = false;
+ if (taskContext != null && TaskContext.get() == null) {
+ CometTaskContextShim.set(taskContext);
+ installedTaskContext = true;
+ }
+ try {
+ evaluateInternal(
+ udfClassName, inputArrayPtrs, inputSchemaPtrs, outArrayPtr, outSchemaPtr, numRows);
+ } finally {
+ if (installedTaskContext) {
+ CometTaskContextShim.unset();
+ }
+ }
+ }
+
+ private static void evaluateInternal(
+ String udfClassName,
+ long[] inputArrayPtrs,
+ long[] inputSchemaPtrs,
+ long outArrayPtr,
+ long outSchemaPtr,
+ int numRows) {
+ LinkedHashMap cache = INSTANCES.get();
+ CometUDF udf = cache.get(udfClassName);
+ if (udf == null) {
+ try {
+ // Resolve via the executor's context classloader so user-supplied UDF jars
+ // (added via spark.jars / --jars) are visible.
+ ClassLoader cl = Thread.currentThread().getContextClassLoader();
+ if (cl == null) {
+ cl = CometUdfBridge.class.getClassLoader();
+ }
+ udf =
+ (CometUDF) Class.forName(udfClassName, true, cl).getDeclaredConstructor().newInstance();
+ } catch (ReflectiveOperationException e) {
+ throw new RuntimeException("Failed to instantiate CometUDF: " + udfClassName, e);
+ }
+ cache.put(udfClassName, udf);
+ }
BufferAllocator allocator = org.apache.comet.package$.MODULE$.CometArrowAllocator();
@@ -84,23 +137,17 @@ public static void evaluate(
inputs[i] = Data.importVector(allocator, inArr, inSch, null);
}
- result = udf.evaluate(inputs);
+ result = udf.evaluate(inputs, numRows);
if (!(result instanceof FieldVector)) {
throw new RuntimeException(
"CometUDF.evaluate() must return a FieldVector, got: " + result.getClass().getName());
}
- // Result length must match the longest input. Scalar (length-1) inputs
- // are allowed to be shorter, but a vector input bounds the output.
- int expectedLen = 0;
- for (ValueVector v : inputs) {
- expectedLen = Math.max(expectedLen, v.getValueCount());
- }
- if (result.getValueCount() != expectedLen) {
+ if (result.getValueCount() != numRows) {
throw new RuntimeException(
"CometUDF.evaluate() returned "
+ result.getValueCount()
+ " rows, expected "
- + expectedLen);
+ + numRows);
}
ArrowArray outArr = ArrowArray.wrap(outArrayPtr);
ArrowSchema outSch = ArrowSchema.wrap(outSchemaPtr);
diff --git a/common/src/main/scala/org/apache/comet/CometConf.scala b/common/src/main/scala/org/apache/comet/CometConf.scala
index 9b376837f7..dcc6359304 100644
--- a/common/src/main/scala/org/apache/comet/CometConf.scala
+++ b/common/src/main/scala/org/apache/comet/CometConf.scala
@@ -380,6 +380,46 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val REGEXP_ENGINE_RUST = "rust"
+ val REGEXP_ENGINE_JAVA = "java"
+
+ val COMET_REGEXP_ENGINE: ConfigEntry[String] =
+ conf("spark.comet.exec.regexp.engine")
+ .category(CATEGORY_EXEC)
+ .doc(
+ "Experimental. Selects the engine used to evaluate supported regular-expression " +
+ s"expressions. `$REGEXP_ENGINE_RUST` uses the native DataFusion regexp engine. " +
+ s"`$REGEXP_ENGINE_JAVA` routes through a JVM-side UDF (java.util.regex.Pattern) for " +
+ "Spark-compatible semantics, at the cost of JNI roundtrips per batch. Expressions " +
+ "routed when set to java: rlike, regexp_extract, regexp_extract_all, regexp_replace, " +
+ "regexp_instr, and split.")
+ .stringConf
+ .transform(_.toLowerCase(Locale.ROOT))
+ .checkValues(Set(REGEXP_ENGINE_RUST, REGEXP_ENGINE_JAVA))
+ .createWithDefault(REGEXP_ENGINE_JAVA)
+
+ val CODEGEN_DISPATCH_AUTO = "auto"
+ val CODEGEN_DISPATCH_DISABLED = "disabled"
+ val CODEGEN_DISPATCH_FORCE = "force"
+
+ val COMET_CODEGEN_DISPATCH_MODE: ConfigEntry[String] =
+ conf("spark.comet.exec.codegenDispatch.mode")
+ .category(CATEGORY_EXEC)
+ .doc("Controls whether Comet routes eligible scalar expressions through the Arrow-direct " +
+ "codegen dispatcher (`CometCodegenDispatchUDF`) rather than through a native " +
+ s"DataFusion implementation or falling back to Spark. `$CODEGEN_DISPATCH_AUTO` lets " +
+ "each expression's serde decide its preferred path based on measured evidence " +
+ "(e.g. for regex, codegen is preferred when " +
+ s"spark.comet.exec.regexp.engine=$REGEXP_ENGINE_JAVA). " +
+ s"`$CODEGEN_DISPATCH_DISABLED` never uses codegen dispatch. `$CODEGEN_DISPATCH_FORCE` " +
+ "inverts the chain: every serde tries codegen first and falls through to its next " +
+ "preferred path only when `canHandle` rejects the expression. Useful for debugging " +
+ "and benchmarking.")
+ .stringConf
+ .transform(_.toLowerCase(Locale.ROOT))
+ .checkValues(Set(CODEGEN_DISPATCH_AUTO, CODEGEN_DISPATCH_DISABLED, CODEGEN_DISPATCH_FORCE))
+ .createWithDefault(CODEGEN_DISPATCH_AUTO)
+
val COMET_EXEC_SHUFFLE_WITH_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] =
conf("spark.comet.native.shuffle.partitioning.hash.enabled")
.category(CATEGORY_SHUFFLE)
diff --git a/common/src/main/scala/org/apache/comet/udf/CometArrayData.scala b/common/src/main/scala/org/apache/comet/udf/CometArrayData.scala
new file mode 100644
index 0000000000..36e11546e7
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/udf/CometArrayData.scala
@@ -0,0 +1,110 @@
+/*
+ * 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.udf
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampNTZType, TimestampType}
+import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
+
+import org.apache.comet.shims.CometInternalRowShim
+
+/**
+ * Shim base for Comet-owned [[ArrayData]] views used by the Arrow-direct codegen kernel.
+ *
+ * Provides `UnsupportedOperationException` defaults for every abstract method on `ArrayData` and
+ * `SpecializedGetters`. Codegen emits a concrete subclass per complex-typed input column,
+ * overriding only the small set of getters the element type requires (e.g. `numElements`,
+ * `isNullAt`, and `getUTF8String` for an `ArrayType(StringType)` input).
+ *
+ * Pattern mirrors [[CometInternalRow]]: centralize the boilerplate throws so the codegen- emitted
+ * subclasses stay short, and absorb forward-compat breakage if Spark adds abstract methods to
+ * `ArrayData` in a future version.
+ *
+ * Mixes in [[CometInternalRowShim]] for the same reason `CometInternalRow` does: Spark 4.x adds
+ * new abstract getters (`getVariant`, `getGeography`, `getGeometry`) on `SpecializedGetters` that
+ * both `InternalRow` and `ArrayData` inherit. The shim is per-profile and provides throwing
+ * defaults only on the profiles that declare those methods abstract.
+ */
+abstract class CometArrayData extends ArrayData with CometInternalRowShim {
+
+ override def numElements(): Int = unsupported("numElements")
+ override def isNullAt(ordinal: Int): Boolean = unsupported("isNullAt")
+
+ override def getBoolean(ordinal: Int): Boolean = unsupported("getBoolean")
+ override def getByte(ordinal: Int): Byte = unsupported("getByte")
+ override def getShort(ordinal: Int): Short = unsupported("getShort")
+ override def getInt(ordinal: Int): Int = unsupported("getInt")
+ override def getLong(ordinal: Int): Long = unsupported("getLong")
+ override def getFloat(ordinal: Int): Float = unsupported("getFloat")
+ override def getDouble(ordinal: Int): Double = unsupported("getDouble")
+ override def getDecimal(ordinal: Int, precision: Int, scale: Int): Decimal =
+ unsupported("getDecimal")
+ override def getUTF8String(ordinal: Int): UTF8String = unsupported("getUTF8String")
+ override def getBinary(ordinal: Int): Array[Byte] = unsupported("getBinary")
+ override def getInterval(ordinal: Int): CalendarInterval = unsupported("getInterval")
+ override def getStruct(ordinal: Int, numFields: Int): InternalRow = unsupported("getStruct")
+ override def getArray(ordinal: Int): ArrayData = unsupported("getArray")
+ override def getMap(ordinal: Int): MapData = unsupported("getMap")
+
+ /**
+ * Generic `get(ordinal, dataType)` dispatcher. Spark codegen sometimes calls this rather than
+ * the typed getter (`SafeProjection` uses it when deserializing struct-valued ScalaUDF args,
+ * for example); leaving it as a throw leaks NPEs once callers catch the
+ * `UnsupportedOperationException` and propagate null. Dispatches to the typed getter matching
+ * `dataType`; a null entry returns `null` outright.
+ */
+ override def get(ordinal: Int, dataType: DataType): AnyRef = {
+ if (isNullAt(ordinal)) return null
+ dataType match {
+ case BooleanType => java.lang.Boolean.valueOf(getBoolean(ordinal))
+ case ByteType => java.lang.Byte.valueOf(getByte(ordinal))
+ case ShortType => java.lang.Short.valueOf(getShort(ordinal))
+ case IntegerType | DateType => java.lang.Integer.valueOf(getInt(ordinal))
+ case LongType | TimestampType | TimestampNTZType =>
+ java.lang.Long.valueOf(getLong(ordinal))
+ case FloatType => java.lang.Float.valueOf(getFloat(ordinal))
+ case DoubleType => java.lang.Double.valueOf(getDouble(ordinal))
+ case _: StringType => getUTF8String(ordinal)
+ case BinaryType => getBinary(ordinal)
+ case dt: DecimalType => getDecimal(ordinal, dt.precision, dt.scale)
+ case st: StructType => getStruct(ordinal, st.size)
+ case _: ArrayType => getArray(ordinal)
+ case _: MapType => getMap(ordinal)
+ case other => unsupported(s"get for dataType $other")
+ }
+ }
+
+ override def setNullAt(i: Int): Unit = unsupported("setNullAt")
+ override def update(i: Int, value: Any): Unit = unsupported("update")
+
+ override def copy(): ArrayData = unsupported("copy")
+ override def array: Array[Any] = unsupported("array")
+ override def toString(): String = {
+ val n =
+ try numElements().toString
+ catch { case _: Throwable => "?" }
+ s"${getClass.getSimpleName}(numElements=$n)"
+ }
+
+ protected def unsupported(method: String): Nothing =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: $method not implemented for this array shape")
+}
diff --git a/common/src/main/scala/org/apache/comet/udf/CometBatchKernelCodegen.scala b/common/src/main/scala/org/apache/comet/udf/CometBatchKernelCodegen.scala
new file mode 100644
index 0000000000..8d18f4297e
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/udf/CometBatchKernelCodegen.scala
@@ -0,0 +1,642 @@
+/*
+ * 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.udf
+
+import org.apache.arrow.vector.{BigIntVector, BitVector, DateDayVector, DecimalVector, FieldVector, Float4Vector, Float8Vector, IntVector, SmallIntVector, TimeStampMicroTZVector, TimeStampMicroVector, TinyIntVector, ValueVector, VarBinaryVector, VarCharVector}
+import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, RegExpReplace, Unevaluable}
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodeAndComment, CodeFormatter, CodegenContext, CodeGenerator, CodegenFallback, ExprCode, GeneratedClass}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{DataType, StringType}
+
+import org.apache.comet.shims.CometExprTraitShim
+
+/**
+ * Compiles a bound [[Expression]] plus an input schema into a specialized [[CometBatchKernel]]
+ * that fuses Arrow input reads, expression evaluation, and Arrow output writes into one
+ * Janino-compiled method per (expression, schema) pair.
+ *
+ * Input- and output-side emission live in [[CometBatchKernelCodegenInput]] and
+ * [[CometBatchKernelCodegenOutput]]. This file is the orchestrator: the [[ArrowColumnSpec]]
+ * vocabulary, [[canHandle]] / [[allocateOutput]] / [[compile]] / [[generateSource]] entry points,
+ * and the cross-cutting kernel-shape decisions (null-intolerant short-circuit, CSE variant,
+ * per-expression specialized emitters).
+ *
+ * The generated kernel '''is''' the `InternalRow` that Spark's `BoundReference.genCode` reads
+ * from. `ctx.INPUT_ROW = "row"` plus `InternalRow row = this;` inside `process` routes every
+ * `row.getUTF8String(ord)` to the kernel's own typed getter (final method, constant ordinal; JIT
+ * devirtualizes and folds the switch). `row` rather than `this` because Spark's
+ * `splitExpressions` uses INPUT_ROW as a helper-method parameter name and `this` is a reserved
+ * Java keyword.
+ *
+ * For the full feature list (type surface, optimizations, cache layers, specialized emitters,
+ * open work items), see `docs/source/contributor-guide/jvm_udf_dispatch.md`.
+ */
+object CometBatchKernelCodegen extends Logging with CometExprTraitShim {
+
+ /**
+ * Per-column compile-time invariants. The concrete Arrow vector class and whether the column is
+ * nullable are baked into the generated kernel's typed fields and branches. Part of the cache
+ * key: different vector classes or nullability produce different kernels.
+ *
+ * Sealed hierarchy so that complex types (array/map/struct) can carry their nested element
+ * shape recursively. Today scalar, array, and struct specs exist; map cases will land as an
+ * additional subclass when the emitter covers them. A companion `apply` / `unapply` preserves
+ * the original scalar-only construction and extractor shape so existing callers don't need to
+ * change.
+ */
+ sealed trait ArrowColumnSpec {
+ def vectorClass: Class[_ <: ValueVector]
+ def nullable: Boolean
+ }
+
+ object ArrowColumnSpec {
+
+ /** Convenience constructor producing a [[ScalarColumnSpec]]. */
+ def apply(vectorClass: Class[_ <: ValueVector], nullable: Boolean): ArrowColumnSpec =
+ ScalarColumnSpec(vectorClass, nullable)
+
+ /**
+ * Backward-compatible extractor for the common scalar case. Callers that want array / struct
+ * / future map specs should pattern match on the subclass directly.
+ */
+ def unapply(spec: ArrowColumnSpec): Option[(Class[_ <: ValueVector], Boolean)] = spec match {
+ case ScalarColumnSpec(c, n) => Some((c, n))
+ case _ => None
+ }
+ }
+
+ /** Scalar column: one Arrow vector class per row slot, no nested structure. */
+ final case class ScalarColumnSpec(vectorClass: Class[_ <: ValueVector], nullable: Boolean)
+ extends ArrowColumnSpec
+
+ /**
+ * Array column: an Arrow `ListVector` wrapping a child spec. `elementSparkType` is the Spark
+ * `DataType` of the element so the nested-class getter emitter can choose the right template
+ * (e.g. `getUTF8String` for `StringType`, `getInt` for `IntegerType`). The child spec carries
+ * the Arrow child vector class. Nested arrays (`Array>`) work by the child being
+ * itself an `ArrayColumnSpec`.
+ */
+ final case class ArrayColumnSpec(
+ nullable: Boolean,
+ elementSparkType: DataType,
+ element: ArrowColumnSpec)
+ extends ArrowColumnSpec {
+ override def vectorClass: Class[_ <: ValueVector] = classOf[ListVector]
+ }
+
+ /**
+ * Struct column: an Arrow `StructVector` wrapping N typed child specs. Each entry carries the
+ * Spark field name (for schema identification in the cache key), the Spark `DataType` of the
+ * field (so per-field emitters pick the right read/write template), the child `ArrowColumnSpec`
+ * (so nested shapes like `Struct>` compose by trait-level recursion), and the
+ * field's `nullable` bit (so non-nullable fields elide their per-row null check at source
+ * level). Nested structs (`Struct>`) work by the child being itself a
+ * `StructColumnSpec`.
+ */
+ final case class StructColumnSpec(nullable: Boolean, fields: Seq[StructFieldSpec])
+ extends ArrowColumnSpec {
+ override def vectorClass: Class[_ <: ValueVector] = classOf[StructVector]
+ }
+
+ /** One field entry on a [[StructColumnSpec]]. */
+ final case class StructFieldSpec(
+ name: String,
+ sparkType: DataType,
+ nullable: Boolean,
+ child: ArrowColumnSpec)
+
+ /**
+ * Map column: an Arrow `MapVector` (subclass of `ListVector`) whose data vector is a
+ * `StructVector` with a key field at ordinal 0 and a value field at ordinal 1. `key` and
+ * `value` are themselves `ArrowColumnSpec` so nested shapes (`Map, Array>`,
+ * `Map, ...>`) compose by trait-level recursion. Nullable map entries are controlled
+ * per-column by the outer map's validity; nullable keys and values are carried in the child
+ * specs' `nullable` bit.
+ */
+ final case class MapColumnSpec(
+ nullable: Boolean,
+ keySparkType: DataType,
+ valueSparkType: DataType,
+ key: ArrowColumnSpec,
+ value: ArrowColumnSpec)
+ extends ArrowColumnSpec {
+ override def vectorClass: Class[_ <: ValueVector] = classOf[MapVector]
+ }
+
+ /**
+ * Resolve an Arrow vector class by its simple name, using the same classloader the codegen uses
+ * internally. Intended for tests: the `common` module shades `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so `classOf[VarCharVector]` at a call site in an unshaded
+ * module refers to a different [[Class]] object than the one the codegen compares against.
+ * Callers pass a simple name and get back the class the production code actually uses.
+ */
+ def vectorClassBySimpleName(name: String): Class[_ <: ValueVector] = name match {
+ case "BitVector" => classOf[BitVector]
+ case "TinyIntVector" => classOf[TinyIntVector]
+ case "SmallIntVector" => classOf[SmallIntVector]
+ case "IntVector" => classOf[IntVector]
+ case "BigIntVector" => classOf[BigIntVector]
+ case "Float4Vector" => classOf[Float4Vector]
+ case "Float8Vector" => classOf[Float8Vector]
+ case "DecimalVector" => classOf[DecimalVector]
+ case "DateDayVector" => classOf[DateDayVector]
+ case "TimeStampMicroVector" => classOf[TimeStampMicroVector]
+ case "TimeStampMicroTZVector" => classOf[TimeStampMicroTZVector]
+ case "VarCharVector" => classOf[VarCharVector]
+ case "VarBinaryVector" => classOf[VarBinaryVector]
+ case other => throw new IllegalArgumentException(s"unknown Arrow vector class: $other")
+ }
+
+ /**
+ * Result of compiling a bound [[Expression]] into a Janino kernel. The `factory` is the Spark
+ * [[GeneratedClass]] produced by Janino and is safe to share across threads and partitions: it
+ * holds no mutable state. The `freshReferences` closure regenerates the references array each
+ * time a new kernel instance is allocated.
+ *
+ * Why not cache a single `references` array: some expressions (notably [[ScalaUDF]]) embed
+ * stateful Spark `ExpressionEncoder` serializers into `references` via `ctx.addReferenceObj`.
+ * Those serializers reuse an internal `UnsafeRow` / `byte[]` buffer per `.apply(...)` call and
+ * are not thread-safe. If two kernels on different partitions shared one serializer instance,
+ * they would race on that buffer and produce garbage. Re-running `genCode` per kernel
+ * allocation costs microseconds; Janino compile costs milliseconds. Cache the expensive piece,
+ * refresh the cheap one, stay correct.
+ *
+ * Mirrors Spark `WholeStageCodegenExec`: compile once per plan, instantiate per partition, call
+ * `init(partitionIndex)` once, iterate.
+ */
+ final case class CompiledKernel(factory: GeneratedClass, freshReferences: () => Array[Any]) {
+ def newInstance(): CometBatchKernel =
+ factory.generate(freshReferences()).asInstanceOf[CometBatchKernel]
+ }
+
+ /**
+ * Plan-time predicate: can the codegen dispatcher handle this bound expression end to end? If
+ * it returns `None`, the serde is free to emit the codegen proto. If it returns `Some(reason)`,
+ * the serde must fall back (usually via `withInfo(...) + None`) so Spark runs the expression
+ * rather than crashing in the Janino compile at execute time.
+ *
+ * Checks:
+ * - every `BoundReference`'s data type is in
+ * [[CometBatchKernelCodegenInput.isSupportedInputType]] (i.e. the kernel has a typed getter
+ * for it)
+ * - the overall `expr.dataType` is in [[CometBatchKernelCodegenOutput.isSupportedOutputType]]
+ * (i.e. `allocateOutput` and `emitWrite` know how to materialize it)
+ * - the expression is scalar (no `AggregateFunction`, no generators). These never reach a
+ * scalar serde, but we belt-and-suspenders anyway.
+ *
+ * Intermediate node types are '''not''' checked. Spark's `doGenCode` materializes intermediates
+ * in local variables; only the leaves (which read from the row) and the root (which writes to
+ * the output vector) touch Arrow.
+ */
+ def canHandle(boundExpr: Expression): Option[String] = {
+ if (!CometBatchKernelCodegenOutput.isSupportedOutputType(boundExpr.dataType)) {
+ return Some(s"codegen dispatch: unsupported output type ${boundExpr.dataType}")
+ }
+ // Reject expressions that can't be safely compiled or cached:
+ // - AggregateFunction / Generator: non-scalar bridge shape.
+ // - CodegenFallback: opts out of `doGenCode`, which our compile path assumes works.
+ // Passing one in would emit interpreted-eval glue that our kernel can't splice cleanly.
+ // - Unevaluable: unresolved plan markers. Shouldn't reach a serde, but cheap to guard.
+ // `isCodegenInertUnevaluable` lets the shim exclude version-specific leaves that are
+ // `Unevaluable` but never touched by codegen (e.g. Spark 4.0's `ResolvedCollation`, which
+ // lives in `Collate.collation` as a type marker; `Collate.genCode` delegates to its child).
+ //
+ // Nondeterministic and stateful expressions are accepted: the dispatcher allocates one
+ // kernel instance per partition (per `CometCodegenDispatchUDF.ensureKernel`) and calls
+ // `init(partitionIndex)` once on partition entry, so per-row state on `Rand`,
+ // `MonotonicallyIncreasingID`, etc. advances correctly across batches in the same
+ // partition and resets across partitions.
+ //
+ // `ExecSubqueryExpression` (e.g. `ScalarSubquery`, `InSubqueryExec`) is also accepted, and
+ // works correctly via a four-link invariant:
+ // 1. The surrounding Comet operator inherits `SparkPlan.waitForSubqueries`, which calls
+ // `updateResult()` on every `ExecSubqueryExpression` in its `expressions` before the
+ // operator's compute path ever reaches the JVM UDF bridge.
+ // 2. `ScalarSubquery.result` (and equivalents on other subquery expressions) is a plain
+ // mutable field on the case class. `@volatile` affects cross-thread visibility but
+ // not serializability: Java/Kryo serializers include it.
+ // 3. `SparkEnv.closureSerializer` captures the populated `result` value in the bytes
+ // that travel through `CometCodegenDispatchUDF`'s arg-0 transport.
+ // 4. The dispatcher's cache key is those exact bytes (see
+ // `CometCodegenDispatchUDF.CacheKey`). Different `result` values produce different
+ // bytes, hence different cache entries, hence a fresh compile per distinct subquery
+ // value. No cross-query staleness.
+ //
+ // If any of those four links breaks (a different cache-key derivation that drops `result`;
+ // a Comet operator that bypasses `waitForSubqueries`; a transport that strips `@volatile`
+ // fields), subquery correctness regresses. Keep this invariant intact when refactoring the
+ // cache-key or transport layers.
+ boundExpr.find {
+ case _: org.apache.spark.sql.catalyst.expressions.aggregate.AggregateFunction => true
+ case _: org.apache.spark.sql.catalyst.expressions.Generator => true
+ case _: CodegenFallback => true
+ case u: Unevaluable if isCodegenInertUnevaluable(u) => false
+ case _: Unevaluable => true
+ case _ => false
+ } match {
+ case Some(bad) =>
+ return Some(
+ s"codegen dispatch: expression ${bad.getClass.getSimpleName} not supported " +
+ "(aggregate, generator, codegen-fallback, or unevaluable)")
+ case None =>
+ }
+ val badRef = boundExpr.collectFirst {
+ case b: BoundReference if !CometBatchKernelCodegenInput.isSupportedInputType(b.dataType) =>
+ b
+ }
+ badRef.map(b =>
+ s"codegen dispatch: unsupported input type ${b.dataType} at ordinal ${b.ordinal}")
+ }
+
+ /**
+ * Allocate an Arrow output vector matching the expression's `dataType`. Thin forwarder to
+ * [[CometBatchKernelCodegenOutput.allocateOutput]]. Kept on this object as part of the public
+ * API so external callers (`CometCodegenDispatchUDF`) do not have to know about the internal
+ * split.
+ */
+ def allocateOutput(
+ dataType: DataType,
+ name: String,
+ numRows: Int,
+ estimatedBytes: Int = -1): FieldVector =
+ CometBatchKernelCodegenOutput.allocateOutput(dataType, name, numRows, estimatedBytes)
+
+ /**
+ * Output of [[generateSource]]. `body` is the raw Java source Janino will compile; `code` is
+ * the post-`stripOverlappingComments` wrapper Janino actually takes as input; `references` are
+ * the runtime objects the generated constructor pulls from via `ctx.addReferenceObj` (cached
+ * patterns, replacement strings, etc.). Tests inspect `body` to assert the shape of the
+ * generated source. See `CometCodegenSourceSuite` for examples.
+ */
+ final case class GeneratedSource(body: String, code: CodeAndComment, references: Array[Any])
+
+ /**
+ * Generate the Java source for a kernel without compiling it. Factored out of [[compile]] so
+ * tests can assert on the emitted source (null short-circuit present, non-nullable `isNullAt`
+ * returns literal `false`, specialized emitter engaged, etc.) without paying for Janino.
+ */
+ def generateSource(
+ boundExpr: Expression,
+ inputSchema: Seq[ArrowColumnSpec]): GeneratedSource = {
+ val ctx = new CodegenContext
+ // `BoundReference.genCode` emits `${ctx.INPUT_ROW}.getUTF8String(ord)`. We alias a local
+ // `row` to `this` at the top of `process` so those reads resolve to the kernel's own typed
+ // getters (virtual dispatch on a concrete final class, JIT devirtualizes + folds the
+ // switch). `row` rather than `this` because Spark's `splitExpressions` uses INPUT_ROW as the
+ // parameter name of any helper method it emits; `this` is a reserved keyword, so using it
+ // as a parameter name produces `private UTF8String helper(InternalRow this)` which Janino
+ // rejects.
+ ctx.INPUT_ROW = "row"
+
+ val baseClass = classOf[CometBatchKernel].getName
+ // Resolve shaded Arrow class names at compile time so generated source
+ // matches the abstract method signature after Maven relocation.
+ val valueVectorClass = classOf[ValueVector].getName
+ val fieldVectorClass = classOf[FieldVector].getName
+
+ // Pick the per-row body. Specialized emitters get priority; the default reuses
+ // Spark's doGenCode.
+ //
+ // `outputSetup` holds once-per-batch declarations (typed child-vector casts for complex
+ // output) that `emitOutputWriter` factors out of the per-row body so they do not repeat on
+ // every row. Scalar outputs return an empty string here. Specialized emitters (like
+ // RegExpReplace) do not need setup because they write directly to the root `output`.
+ //
+ // TODO(method-size): perRowBody is inlined inside process's for-loop and not split.
+ // Sufficiently deep trees can exceed Janino's 64KB method size; wrap in
+ // ctx.splitExpressionsWithCurrentInputs when hit. See
+ // docs/source/contributor-guide/jvm_udf_dispatch.md#open-items.
+ val (concreteOutClass, outputSetup, perRowBody) = boundExpr match {
+ case rr: RegExpReplace if canSpecializeRegExpReplace(rr) =>
+ (classOf[VarCharVector].getName, "", specializedRegExpReplaceBody(ctx, rr, inputSchema))
+ case _ =>
+ // Class-field CSE. `generateExpressions` runs `subexpressionElimination` under the
+ // hood, which populates `ctx.subexprFunctions` with per-row helper calls that write
+ // common subexpression results into `addMutableState`-allocated fields; the returned
+ // `ExprCode` then references those fields. `subexprFunctionsCode` is the concatenated
+ // helper invocation block, spliced into the per-row body by `defaultBody` (inside the
+ // NullIntolerant else-branch when that short-circuit fires, otherwise before
+ // `ev.code`). See the "Subexpression elimination" section of the object-level
+ // Scaladoc for why we use this variant rather than the WSCG one.
+ val ev = if (SQLConf.get.subexpressionEliminationEnabled) {
+ ctx.generateExpressions(Seq(boundExpr), doSubexpressionElimination = true).head
+ } else {
+ boundExpr.genCode(ctx)
+ }
+ val subExprsCode = ctx.subexprFunctionsCode
+ val (cls, setup, snippet) =
+ CometBatchKernelCodegenOutput.emitOutputWriter(boundExpr.dataType, ev.value, ctx)
+ (cls, setup, defaultBody(boundExpr, ev, snippet, subExprsCode))
+ }
+
+ val typedFieldDecls = CometBatchKernelCodegenInput.emitInputFieldDecls(inputSchema)
+ val typedInputCasts = CometBatchKernelCodegenInput.emitInputCasts(inputSchema)
+ val decimalTypeByOrdinal = CometBatchKernelCodegenInput.decimalPrecisionByOrdinal(boundExpr)
+ val getters =
+ CometBatchKernelCodegenInput.emitTypedGetters(inputSchema, decimalTypeByOrdinal)
+ val nested = CometBatchKernelCodegenInput.emitNestedClasses(inputSchema)
+ val getArrayMethod = CometBatchKernelCodegenInput.emitGetArrayMethod(inputSchema)
+ val getStructMethod = CometBatchKernelCodegenInput.emitGetStructMethod(inputSchema)
+ val getMapMethod = CometBatchKernelCodegenInput.emitGetMapMethod(inputSchema)
+
+ val codeBody =
+ s"""
+ |public java.lang.Object generate(Object[] references) {
+ | return new SpecificCometBatchKernel(references);
+ |}
+ |
+ |final class SpecificCometBatchKernel extends $baseClass {
+ |
+ | ${ctx.declareMutableStates()}
+ |
+ | $typedFieldDecls
+ | private int rowIdx;
+ |
+ | public SpecificCometBatchKernel(Object[] references) {
+ | super(references);
+ | ${ctx.initMutableStates()}
+ | }
+ |
+ | @Override
+ | public void init(int partitionIndex) {
+ | ${ctx.initPartition()}
+ | }
+ |
+ | $getters
+ | $getArrayMethod
+ | $getStructMethod
+ | $getMapMethod
+ |
+ | @Override
+ | public void process(
+ | $valueVectorClass[] inputs,
+ | $fieldVectorClass outRaw,
+ | int numRows) {
+ | $concreteOutClass output = ($concreteOutClass) outRaw;
+ | $typedInputCasts
+ | $outputSetup
+ | // Alias the kernel as `row` so Spark-generated `${ctx.INPUT_ROW}.method()` reads
+ | // resolve to the kernel's own typed getters. Helper methods that Spark splits off
+ | // via `splitExpressions` also take `InternalRow row` as a parameter; we pass `this`
+ | // implicitly since callers substitute INPUT_ROW which we've set to `row`.
+ | org.apache.spark.sql.catalyst.InternalRow row = this;
+ | for (int i = 0; i < numRows; i++) {
+ | this.rowIdx = i;
+ | $perRowBody
+ | }
+ | }
+ |
+ | ${ctx.declareAddedFunctions()}
+ |
+ |$nested
+ |}
+ """.stripMargin
+
+ val code = CodeFormatter.stripOverlappingComments(
+ new CodeAndComment(codeBody, ctx.getPlaceHolderToComments()))
+ GeneratedSource(code.body, code, ctx.references.toArray)
+ }
+
+ def compile(boundExpr: Expression, inputSchema: Seq[ArrowColumnSpec]): CompiledKernel = {
+ val src = generateSource(boundExpr, inputSchema)
+ val (clazz, _) =
+ try {
+ CodeGenerator.compile(src.code)
+ } catch {
+ case t: Throwable =>
+ logError(
+ s"CometBatchKernelCodegen: compile failed for ${boundExpr.getClass.getSimpleName}. " +
+ s"Generated source follows:\n${src.body}",
+ t)
+ throw t
+ }
+ // One log per unique (expr, schema) compile; the caller caches the result so subsequent
+ // batches with the same shape reuse this compile.
+ val specialized = boundExpr match {
+ case _: RegExpReplace
+ if canSpecializeRegExpReplace(boundExpr.asInstanceOf[RegExpReplace]) =>
+ " [specialized]"
+ case _ => ""
+ }
+ logInfo(
+ s"CometBatchKernelCodegen: compiled ${boundExpr.getClass.getSimpleName}$specialized " +
+ s"-> ${boundExpr.dataType} inputs=" +
+ inputSchema
+ .map(s => s"${s.vectorClass.getSimpleName}${if (s.nullable) "?" else ""}")
+ .mkString(","))
+ // Freshen references per kernel allocation. See the `CompiledKernel` scaladoc for why.
+ // `generateSource` is pure with respect to its inputs (no hidden state) and produces a
+ // layout-compatible references array each time because the expression and schema are
+ // fixed.
+ val freshReferences: () => Array[Any] = () =>
+ generateSource(boundExpr, inputSchema).references
+ CompiledKernel(clazz, freshReferences)
+ }
+
+ /**
+ * Can this `RegExpReplace` instance be handled by the specialized emitter? Requires a direct
+ * column reference as subject, non-null foldable pattern and replacement, and offset of 1.
+ * Other shapes fall back to the default `doGenCode` path.
+ */
+ private def canSpecializeRegExpReplace(rr: RegExpReplace): Boolean = {
+ val subjectIsBound =
+ rr.subject.isInstanceOf[BoundReference] && rr.subject.dataType == StringType
+ val patternOk =
+ rr.regexp.foldable && rr.regexp.dataType == StringType && rr.regexp.eval() != null
+ val replOk = rr.rep.foldable && rr.rep.dataType == StringType && rr.rep.eval() != null
+ val posIsOne = rr.pos match {
+ case Literal(v: Int, _) => v == 1
+ case _ => false
+ }
+ subjectIsBound && patternOk && replOk && posIsOne
+ }
+
+ /**
+ * Emit the per-row body for `RegExpReplace`. Per-row shape: read Arrow subject bytes, decode to
+ * Java `String`, run `Matcher.replaceAll` with a cached `Pattern` and the replacement String,
+ * re-encode to bytes, write to Arrow.
+ *
+ * ==Why this specialization exists==
+ *
+ * The default path runs `boundExpr.genCode(ctx)` and wraps it with kernel-side getter reads and
+ * a `UTF8String -> bytes -> Arrow` write. For `RegExpReplace` specifically, Spark's generated
+ * code does not stay in `UTF8String` space: `java.util.regex.Matcher` requires a
+ * `CharSequence`, so the generated code materializes a Java `String` from the input
+ * `UTF8String` (a UTF-8 decode, allocating a `char[]`), runs the matcher, then wraps the result
+ * String back into a `UTF8String` (a UTF-8 encode, allocating a `byte[]`). The per-row shape
+ * is:
+ *
+ * {{{
+ * default: Arrow bytes -> UTF8String -> String -> Matcher ->
+ * String -> UTF8String -> bytes -> Arrow
+ * }}}
+ *
+ * On a wide-match workload (every character of the row gets replaced, so the output is the full
+ * row length), the round trip added ~44% per-row cost versus a tight byte-oriented loop with
+ * shape:
+ *
+ * {{{
+ * specialized: Arrow bytes -> String -> Matcher -> String -> bytes -> Arrow
+ * }}}
+ *
+ * This specialization emits the byte-oriented shape directly. No `UTF8String` appears in the
+ * generated per-row loop. The expression remains a first-class citizen of the dispatcher
+ * (plan-time serde, schema-keyed caching, zero-config for the caller).
+ *
+ * ==When to add a specialization==
+ *
+ * The general rule: specialize when an expression's `doGenCode` output shape forces conversions
+ * that an Arrow-aware byte-oriented implementation does not pay. The common case is expressions
+ * whose implementation requires a Java `String` (anything using `java.util.regex` and some
+ * `DateTimeFormatter` expressions), because Spark's `UTF8String <-> String` round-trip is not
+ * free for wide outputs. Keep specializations minimal so comparisons stay honest. Avoid
+ * layering speculative optimizations; let the default-path optimization menu handle the common
+ * cases.
+ */
+ private def specializedRegExpReplaceBody(
+ ctx: CodegenContext,
+ rr: RegExpReplace,
+ inputSchema: Seq[ArrowColumnSpec]): String = {
+ val subjectOrd = rr.subject.asInstanceOf[BoundReference].ordinal
+ val subjectClass = inputSchema(subjectOrd).vectorClass
+ require(
+ subjectClass == classOf[VarCharVector],
+ "specializedRegExpReplaceBody expects VarCharVector at ordinal " +
+ s"$subjectOrd, got ${subjectClass.getSimpleName}")
+
+ val patternStr = rr.regexp.eval().toString
+ val replStr = rr.rep.eval().toString
+ val compiledPattern = java.util.regex.Pattern.compile(patternStr)
+
+ // addReferenceObj adds a class-level field initialized from references[] in the constructor,
+ // so the Pattern and replacement String are resolved once, not per row.
+ val patternRef =
+ ctx.addReferenceObj("pattern", compiledPattern, "java.util.regex.Pattern")
+ val replRef = ctx.addReferenceObj("replacement", replStr, "java.lang.String")
+
+ val sb = ctx.freshName("sb")
+ val s = ctx.freshName("s")
+ val r = ctx.freshName("r")
+ val rb = ctx.freshName("rb")
+
+ s"""
+ |if (this.col$subjectOrd.isNull(i)) {
+ | output.setNull(i);
+ |} else {
+ | byte[] $sb = this.col$subjectOrd.get(i);
+ | String $s = new String($sb, java.nio.charset.StandardCharsets.UTF_8);
+ | String $r = $patternRef.matcher($s).replaceAll($replRef);
+ | byte[] $rb = $r.getBytes(java.nio.charset.StandardCharsets.UTF_8);
+ | output.setSafe(i, $rb, 0, $rb.length);
+ |}
+ """.stripMargin
+ }
+
+ /**
+ * Per-row body for the default (non-specialized) path.
+ *
+ * For expressions that implement the `NullIntolerant` marker trait (null in any input -> null
+ * output), emits a short-circuit that skips expression evaluation entirely when any input
+ * column is null in the current row. This saves the full `ev.code` cost for null rows, not just
+ * the output setNull call. Does not change behavior, only performance.
+ *
+ * For other expressions, the standard shape applies: evaluate the expression, then check
+ * `ev.isNull` to decide between `setNull` and a write. Null semantics are handled internally by
+ * Spark's generated `ev.code`.
+ *
+ * `subExprsCode` is the CSE helper-invocation block (see the "Subexpression elimination"
+ * section of the object-level Scaladoc). It writes common subexpression results into class
+ * fields that `ev.code` reads, so it must run before `ev.code`. In the NullIntolerant short-
+ * circuit case it is placed inside the else branch, skipping CSE evaluation for null rows as
+ * well as main-body evaluation. In the default case it precedes `ev.code`. Empty string when
+ * CSE is disabled or the tree has no common subexpressions.
+ */
+ private def defaultBody(
+ boundExpr: Expression,
+ ev: ExprCode,
+ writeSnippet: String,
+ subExprsCode: String): String = {
+ boundExpr match {
+ case _ if isNullIntolerant(boundExpr) && allNullIntolerant(boundExpr) =>
+ // Every node from root to leaf is either NullIntolerant or a leaf. That transitively
+ // guarantees "any BoundReference null at this row -> whole expression null", so we can
+ // short-circuit on the union of input ordinals. Breaking the chain with a non-null-
+ // propagating node like `coalesce` or `if` produces the wrong result (coalesce(null,x)
+ // is x, not null), so the check above rejects those shapes and falls through to the
+ // default branch which runs Spark's own null-aware ev.code.
+ val inputOrdinals =
+ boundExpr.collect { case b: BoundReference => b.ordinal }.distinct
+ val nullCheck =
+ if (inputOrdinals.isEmpty) "false"
+ else inputOrdinals.map(ord => s"this.col$ord.isNull(i)").mkString(" || ")
+ s"""
+ |if ($nullCheck) {
+ | output.setNull(i);
+ |} else {
+ | $subExprsCode
+ | ${ev.code}
+ | $writeSnippet
+ |}
+ """.stripMargin
+ case _ =>
+ // Optimization: NonNullableOutputShortCircuit.
+ // When the bound expression declares `nullable = false`, the `if (ev.isNull)` branch is
+ // dead and HotSpot may or may not fold it (it depends on whether the expression's
+ // `doGenCode` made `ev.isNull` a `FalseLiteral` or a variable whose value is
+ // false-at-runtime but not a compile-time constant from Spark's side). Drop the guard
+ // at source level so we don't depend on JIT folding and keep the generated body
+ // minimal.
+ if (!boundExpr.nullable) {
+ s"""
+ |$subExprsCode
+ |${ev.code}
+ |$writeSnippet
+ """.stripMargin
+ } else {
+ s"""
+ |$subExprsCode
+ |${ev.code}
+ |if (${ev.isNull}) {
+ | output.setNull(i);
+ |} else {
+ | $writeSnippet
+ |}
+ """.stripMargin
+ }
+ }
+ }
+
+ /**
+ * True iff every node in the expression tree is either `NullIntolerant` or a leaf we can safely
+ * consider null-propagating (`BoundReference` and `Literal`). Used to gate the `NullIntolerant`
+ * short-circuit in [[defaultBody]]: the short-circuit collects `BoundReference` ordinals from
+ * the whole tree and skips `ev.code` when any of them is null, which is only correct when every
+ * path from a leaf to the root propagates nulls. A non- propagating node (`Coalesce`, `If`,
+ * `CaseWhen`, `Concat`, etc.) anywhere in the tree invalidates this assumption: `coalesce(null,
+ * x)` is `x`, not null, so pre-nulling on any input null would produce the wrong result.
+ */
+ private def allNullIntolerant(expr: Expression): Boolean =
+ !expr.exists {
+ case _: BoundReference | _: Literal => false
+ case other => !isNullIntolerant(other)
+ }
+}
diff --git a/common/src/main/scala/org/apache/comet/udf/CometBatchKernelCodegenInput.scala b/common/src/main/scala/org/apache/comet/udf/CometBatchKernelCodegenInput.scala
new file mode 100644
index 0000000000..b4cdfd4595
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/udf/CometBatchKernelCodegenInput.scala
@@ -0,0 +1,1095 @@
+/*
+ * 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.udf
+
+import scala.collection.mutable
+
+import org.apache.arrow.vector.{BigIntVector, BitVector, DateDayVector, DecimalVector, Float4Vector, Float8Vector, IntVector, SmallIntVector, TimeStampMicroTZVector, TimeStampMicroVector, TinyIntVector, VarBinaryVector, VarCharVector}
+import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampNTZType, TimestampType}
+
+import org.apache.comet.udf.CometBatchKernelCodegen.{ArrayColumnSpec, ArrowColumnSpec, MapColumnSpec, ScalarColumnSpec, StructColumnSpec}
+import org.apache.comet.vector.CometPlainVector
+
+/**
+ * Input-side emitters for the Arrow-direct codegen kernel. Everything that generates source for
+ * reading Arrow input into Spark's typed getter surface lives here: kernel field declarations,
+ * per-batch input casts, top-level typed-getter switches, nested `InputArray_${path}` /
+ * `InputStruct_${path}` / `InputMap_${path}` classes at every complex level, and the input-side
+ * type-support gate.
+ *
+ * ==Path encoding for nested complex types==
+ *
+ * Each position in a spec tree has a unique path string, used as the suffix on typed vector
+ * fields and as the identifier on nested classes. Starting from the column ordinal:
+ *
+ * - root: `col${ord}`
+ * - array element of `P`: `${P}_e`
+ * - struct field `fi` of `P`: `${P}_f${fi}`
+ * - map key of `P`: `${P}_k`
+ * - map value of `P`: `${P}_v`
+ *
+ * ==Nested-class composition==
+ *
+ * A nested class at path `P` represents a Spark `ArrayData`, `InternalRow`, or `MapData` view of
+ * its Arrow vector. For any complex child one level down, the class holds a pre-allocated
+ * instance of the corresponding inner nested class and routes `getArray` / `getStruct` / `getMap`
+ * / `keyArray` / `valueArray` calls to that instance after resetting it. N-deep nesting falls out
+ * of this: each level only knows about its immediate children.
+ *
+ * ==Unified reset protocol==
+ *
+ * `InputArray_${path}` and `InputMap_${path}` classes both take `reset(int startIdx, int length)`
+ * and simply capture the slice. Callers (kernel top-level switches, outer complex-getter routers,
+ * map `keyArray` / `valueArray` returns) compute `(startIdx, length)` from the appropriate parent
+ * offsets before calling `reset`. This unifies the view shape across list-backed arrays and map
+ * key/value slices. Structs stay flat-indexed: `InputStruct_${path}` has `reset(int rowIdx)` that
+ * just captures the outer row index.
+ *
+ * Paired with [[CometBatchKernelCodegenOutput]], which handles the symmetric output side.
+ */
+private[udf] object CometBatchKernelCodegenInput {
+
+ /**
+ * Input types the kernel has a typed getter for. Recursive: `ArrayType(inner)` supported when
+ * `inner` is supported; `StructType` when every field is; `MapType` when key and value types
+ * are both supported.
+ */
+ def isSupportedInputType(dt: DataType): Boolean = dt match {
+ case BooleanType | ByteType | ShortType | IntegerType | LongType => true
+ case FloatType | DoubleType => true
+ case _: DecimalType => true
+ case _: StringType | _: BinaryType => true
+ case DateType | TimestampType | TimestampNTZType => true
+ case ArrayType(inner, _) => isSupportedInputType(inner)
+ case st: StructType => st.fields.forall(f => isSupportedInputType(f.dataType))
+ case mt: MapType => isSupportedInputType(mt.keyType) && isSupportedInputType(mt.valueType)
+ case _ => false
+ }
+
+ /**
+ * Emit the kernel's typed vector-field declarations for every level of every input column's
+ * spec tree. Top-level complex columns additionally get an instance-field declaration for the
+ * pre-allocated nested class. Instance fields for nested-class children one level down live
+ * inside the parent nested class.
+ */
+ def emitInputFieldDecls(inputSchema: Seq[ArrowColumnSpec]): String = {
+ val lines = new mutable.ArrayBuffer[String]()
+ inputSchema.zipWithIndex.foreach { case (spec, ord) =>
+ val path = s"col$ord"
+ collectVectorFieldDecls(path, spec, lines)
+ collectTopLevelInstanceDecl(path, spec, lines)
+ }
+ lines.mkString("\n ")
+ }
+
+ /**
+ * Primitive Arrow vector classes that we wrap in [[CometPlainVector]] at the kernel's input-
+ * cast time. `CometPlainVector.get*` reads use `Platform.get*` against a `final long` buffer
+ * address, so JIT inlines them to branchless reads with no per-call `ArrowBuf` dereference.
+ * `CometPlainVector.getBoolean` also includes a bit-packed data-byte cache that collapses 8
+ * sequential bit reads to 1 byte read.
+ *
+ * Not wrapped: `DecimalVector` (kernel emits inline unsafe reads keyed on compile-time
+ * precision, so the fast/slow split stays branchless in the emitted Java rather than branching
+ * at runtime inside `CometPlainVector.getDecimal`), `VarCharVector` / `VarBinaryVector` (kernel
+ * emits inline unsafe reads to avoid the redundant `isNullAt` check inside
+ * `CometPlainVector.getUTF8String` / `getBinary`).
+ */
+ private val primitiveArrowClasses: Set[Class[_]] = Set(
+ classOf[BitVector],
+ classOf[TinyIntVector],
+ classOf[SmallIntVector],
+ classOf[IntVector],
+ classOf[BigIntVector],
+ classOf[Float4Vector],
+ classOf[Float8Vector],
+ classOf[DateDayVector],
+ classOf[TimeStampMicroVector],
+ classOf[TimeStampMicroTZVector])
+
+ private def wrapsInCometPlainVector(cls: Class[_]): Boolean =
+ primitiveArrowClasses.contains(cls)
+
+ /**
+ * Non-wrapped scalar columns that want a cached data-buffer address for inline unsafe reads.
+ * `DecimalVector` uses it for the short-precision fast path (`Platform.getLong`);
+ * `VarCharVector` / `VarBinaryVector` use it as the base address for `UTF8String.fromAddress` /
+ * `Platform.copyMemory`. See the unsafe-emitter block at the bottom of this file for why we
+ * inline rather than reuse `CometPlainVector`.
+ */
+ private def needsValueAddrField(cls: Class[_]): Boolean =
+ cls == classOf[DecimalVector] ||
+ cls == classOf[VarCharVector] ||
+ cls == classOf[VarBinaryVector]
+
+ /** Variable-width columns also want the offset-buffer address cached for `Platform.getInt`. */
+ private def needsOffsetAddrField(cls: Class[_]): Boolean =
+ cls == classOf[VarCharVector] || cls == classOf[VarBinaryVector]
+
+ /**
+ * Java method name for the null check on a column's typed field. Primitive scalars wrapped in
+ * [[CometPlainVector]] expose `isNullAt`; Arrow typed fields (complex containers,
+ * `DecimalVector`, `VarCharVector`, `VarBinaryVector`) expose `isNull`. Both read the validity
+ * bitmap.
+ */
+ private def nullCheckMethod(spec: ArrowColumnSpec): String = spec match {
+ case sc: ScalarColumnSpec if wrapsInCometPlainVector(sc.vectorClass) => "isNullAt"
+ case _ => "isNull"
+ }
+
+ private val cometPlainVectorName: String = classOf[CometPlainVector].getName
+
+ private def collectVectorFieldDecls(
+ path: String,
+ spec: ArrowColumnSpec,
+ out: mutable.ArrayBuffer[String]): Unit = spec match {
+ case sc: ScalarColumnSpec =>
+ // Primitive scalar columns (at any nesting depth) are wrapped in CometPlainVector so
+ // per-row reads go through JIT-inlined Platform.get* against a cached buffer address.
+ // DecimalVector / VarCharVector / VarBinaryVector stay on the Arrow typed field but
+ // cache data- and (variable-width) offset-buffer addresses for inline unsafe reads.
+ val fieldClass =
+ if (wrapsInCometPlainVector(sc.vectorClass)) cometPlainVectorName
+ else sc.vectorClass.getName
+ out += s"private $fieldClass $path;"
+ if (needsValueAddrField(sc.vectorClass)) {
+ out += s"private long ${path}_valueAddr;"
+ }
+ if (needsOffsetAddrField(sc.vectorClass)) {
+ out += s"private long ${path}_offsetAddr;"
+ }
+ case ar: ArrayColumnSpec =>
+ out += s"private ${classOf[ListVector].getName} $path;"
+ collectVectorFieldDecls(s"${path}_e", ar.element, out)
+ case st: StructColumnSpec =>
+ out += s"private ${classOf[StructVector].getName} $path;"
+ st.fields.zipWithIndex.foreach { case (f, fi) =>
+ collectVectorFieldDecls(s"${path}_f$fi", f.child, out)
+ }
+ case mp: MapColumnSpec =>
+ out += s"private ${classOf[MapVector].getName} $path;"
+ // Key and value vectors live at `${P}_k_e` / `${P}_v_e` so the `InputArray_${P}_k` /
+ // `InputArray_${P}_v` synthetic classes (which follow the array-element convention of
+ // reading from `${path}_e`) resolve their element reads correctly.
+ collectVectorFieldDecls(s"${path}_k_e", mp.key, out)
+ collectVectorFieldDecls(s"${path}_v_e", mp.value, out)
+ }
+
+ private def collectTopLevelInstanceDecl(
+ path: String,
+ spec: ArrowColumnSpec,
+ out: mutable.ArrayBuffer[String]): Unit = spec match {
+ case _: ScalarColumnSpec => ()
+ case _: ArrayColumnSpec =>
+ out += s"private final InputArray_$path ${path}_arrayData = new InputArray_$path();"
+ case _: StructColumnSpec =>
+ out += s"private final InputStruct_$path ${path}_structData = new InputStruct_$path();"
+ case _: MapColumnSpec =>
+ out += s"private final InputMap_$path ${path}_mapData = new InputMap_$path();"
+ }
+
+ /**
+ * Emit the per-batch cast statements. For a map column, casts the outer `MapVector`, then casts
+ * the inner `StructVector` (via a local variable) to extract key and value children via
+ * `getChildByOrdinal(0)` / `(1)`. For arrays, casts the outer `ListVector` and recurses via
+ * `getDataVector()`. For structs, casts the outer `StructVector` and recurses via
+ * `getChildByOrdinal(fi)`.
+ */
+ def emitInputCasts(inputSchema: Seq[ArrowColumnSpec]): String = {
+ val lines = new mutable.ArrayBuffer[String]()
+ inputSchema.zipWithIndex.foreach { case (spec, ord) =>
+ val path = s"col$ord"
+ collectCasts(path, spec, s"inputs[$ord]", lines)
+ }
+ lines.mkString("\n ")
+ }
+
+ private def collectCasts(
+ path: String,
+ spec: ArrowColumnSpec,
+ source: String,
+ out: mutable.ArrayBuffer[String]): Unit = spec match {
+ case sc: ScalarColumnSpec =>
+ if (wrapsInCometPlainVector(sc.vectorClass)) {
+ // Wrap in CometPlainVector so per-row reads go through Platform.get* against a final
+ // long buffer address. JIT inlines the one-liner getters, treating the address as a
+ // register-cached constant across the process loop. useDecimal128 = true matches
+ // Spark's 128-bit decimal storage.
+ out += s"this.$path = new $cometPlainVectorName($source, true);"
+ } else {
+ out += s"this.$path = (${sc.vectorClass.getName}) $source;"
+ }
+ if (needsValueAddrField(sc.vectorClass)) {
+ out += s"this.${path}_valueAddr = this.$path.getDataBuffer().memoryAddress();"
+ }
+ if (needsOffsetAddrField(sc.vectorClass)) {
+ out += s"this.${path}_offsetAddr = this.$path.getOffsetBuffer().memoryAddress();"
+ }
+ case ar: ArrayColumnSpec =>
+ out += s"this.$path = (${classOf[ListVector].getName}) $source;"
+ collectCasts(s"${path}_e", ar.element, s"this.$path.getDataVector()", out)
+ case st: StructColumnSpec =>
+ out += s"this.$path = (${classOf[StructVector].getName}) $source;"
+ st.fields.zipWithIndex.foreach { case (f, fi) =>
+ collectCasts(s"${path}_f$fi", f.child, s"this.$path.getChildByOrdinal($fi)", out)
+ }
+ case mp: MapColumnSpec =>
+ // MapVector's data vector is a StructVector with key at child 0 and value at child 1.
+ // Grab the struct through a local var and pull out the typed children. The key / value
+ // vectors live at the `_k_e` / `_v_e` paths so the synthetic `InputArray_${P}_k` /
+ // `InputArray_${P}_v` classes read them via the standard array-element convention.
+ val structLocal = s"${path}__mapStruct"
+ out += s"this.$path = (${classOf[MapVector].getName}) $source;"
+ out += s"${classOf[StructVector].getName} $structLocal = " +
+ s"(${classOf[StructVector].getName}) this.$path.getDataVector();"
+ collectCasts(s"${path}_k_e", mp.key, s"$structLocal.getChildByOrdinal(0)", out)
+ collectCasts(s"${path}_v_e", mp.value, s"$structLocal.getChildByOrdinal(1)", out)
+ }
+
+ /**
+ * Emit the kernel's typed-getter overrides. Spark's `InternalRow` provides the base virtual
+ * method; the `@Override` on a final class gives the JIT enough information to devirtualize.
+ * Each getter switches on the column ordinal so the call site (with an inlined constant ordinal
+ * from `BoundReference.genCode`) folds down to a single branch.
+ *
+ * `decimalTypeByOrdinal` lets the decimal getter specialize per ordinal: when a
+ * `BoundReference` of `DecimalType(precision <= 18)` is the only decimal read at that ordinal,
+ * the emitted case skips the `BigDecimal` allocation and reads the unscaled long directly.
+ *
+ * TODO(unsafe-readers): primitive `v.get(i)` performs a bounds check that is redundant given `i
+ * in [0, numRows)`. See `docs/source/contributor-guide/jvm_udf_dispatch.md#open-items`.
+ */
+ def emitTypedGetters(
+ inputSchema: Seq[ArrowColumnSpec],
+ decimalTypeByOrdinal: Map[Int, Option[DecimalType]]): String = {
+ val withOrd = inputSchema.zipWithIndex
+
+ val isNullCases = withOrd.map { case (spec, ord) =>
+ if (!spec.nullable) {
+ s" case $ord: return false;"
+ } else {
+ // CometPlainVector exposes `isNullAt`; Arrow-typed fields expose `isNull`. Both check
+ // the validity bitmap with the same semantics.
+ val method = spec.vectorClass match {
+ case cls if wrapsInCometPlainVector(cls) => "isNullAt"
+ case _ => "isNull"
+ }
+ s" case $ord: return this.col$ord.$method(this.rowIdx);"
+ }
+ }
+
+ val booleanCases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord) if cls == classOf[BitVector] =>
+ s" case $ord: return this.col$ord.getBoolean(this.rowIdx);"
+ }
+ val byteCases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord) if cls == classOf[TinyIntVector] =>
+ s" case $ord: return this.col$ord.getByte(this.rowIdx);"
+ }
+ val shortCases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord) if cls == classOf[SmallIntVector] =>
+ s" case $ord: return this.col$ord.getShort(this.rowIdx);"
+ }
+ val intCases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord)
+ if cls == classOf[IntVector] || cls == classOf[DateDayVector] =>
+ s" case $ord: return this.col$ord.getInt(this.rowIdx);"
+ }
+ val longCases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord)
+ if cls == classOf[BigIntVector] ||
+ cls == classOf[TimeStampMicroVector] ||
+ cls == classOf[TimeStampMicroTZVector] =>
+ s" case $ord: return this.col$ord.getLong(this.rowIdx);"
+ }
+ val floatCases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord) if cls == classOf[Float4Vector] =>
+ s" case $ord: return this.col$ord.getFloat(this.rowIdx);"
+ }
+ val doubleCases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord) if cls == classOf[Float8Vector] =>
+ s" case $ord: return this.col$ord.getDouble(this.rowIdx);"
+ }
+ val decimalCases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord) if cls == classOf[DecimalVector] =>
+ val known = decimalTypeByOrdinal.getOrElse(ord, None)
+ val valueAddr = s"this.col${ord}_valueAddr"
+ val slowField = s"this.col$ord"
+ val fastPath = emitDecimalFastBodyUnsafe(valueAddr, "this.rowIdx", " ")
+ val slowPath = emitDecimalSlowBody(slowField, "this.rowIdx", " ")
+ val body = known match {
+ case Some(dt) if dt.precision <= 18 => fastPath
+ case Some(_) => slowPath
+ case None =>
+ s""" if (precision <= 18) {
+ |$fastPath
+ | } else {
+ |$slowPath
+ | }""".stripMargin
+ }
+ s""" case $ord: {
+ |$body
+ | }""".stripMargin
+ }
+ val binaryCases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord) if cls == classOf[VarBinaryVector] =>
+ s""" case $ord: {
+ |${emitBinaryBodyUnsafe(
+ s"this.col${ord}_valueAddr",
+ s"this.col${ord}_offsetAddr",
+ "this.rowIdx",
+ " ")}
+ | }""".stripMargin
+ }
+ val utf8Cases = withOrd.collect {
+ case (ArrowColumnSpec(cls, _), ord) if cls == classOf[VarCharVector] =>
+ s""" case $ord: {
+ |${emitUtf8BodyUnsafe(
+ s"this.col${ord}_valueAddr",
+ s"this.col${ord}_offsetAddr",
+ "this.rowIdx",
+ " ")}
+ | }""".stripMargin
+ }
+
+ Seq(
+ emitOrdinalSwitch("public boolean isNullAt(int ordinal)", "isNullAt", isNullCases),
+ emitOrdinalSwitch("public boolean getBoolean(int ordinal)", "getBoolean", booleanCases),
+ emitOrdinalSwitch("public byte getByte(int ordinal)", "getByte", byteCases),
+ emitOrdinalSwitch("public short getShort(int ordinal)", "getShort", shortCases),
+ emitOrdinalSwitch("public int getInt(int ordinal)", "getInt", intCases),
+ emitOrdinalSwitch("public long getLong(int ordinal)", "getLong", longCases),
+ emitOrdinalSwitch("public float getFloat(int ordinal)", "getFloat", floatCases),
+ emitOrdinalSwitch("public double getDouble(int ordinal)", "getDouble", doubleCases),
+ emitOrdinalSwitch(
+ "public org.apache.spark.sql.types.Decimal getDecimal(" +
+ "int ordinal, int precision, int scale)",
+ "getDecimal",
+ decimalCases),
+ emitOrdinalSwitch("public byte[] getBinary(int ordinal)", "getBinary", binaryCases),
+ emitOrdinalSwitch(
+ "public org.apache.spark.unsafe.types.UTF8String getUTF8String(int ordinal)",
+ "getUTF8String",
+ utf8Cases)).mkString
+ }
+
+ /**
+ * Build a per-ordinal map of the `DecimalType` observed on `BoundReference`s in the bound
+ * expression. Used by [[emitTypedGetters]] to emit a compile-time-specialized `getDecimal` case
+ * per ordinal.
+ */
+ def decimalPrecisionByOrdinal(boundExpr: Expression): Map[Int, Option[DecimalType]] = {
+ boundExpr
+ .collect {
+ case b: BoundReference if b.dataType.isInstanceOf[DecimalType] =>
+ b.ordinal -> b.dataType.asInstanceOf[DecimalType]
+ }
+ .groupBy(_._1)
+ .map { case (ord, pairs) =>
+ val distinct = pairs.map(_._2).toSet
+ ord -> (if (distinct.size == 1) Some(distinct.head) else None)
+ }
+ }
+
+ /**
+ * Emit every nested class needed for every complex level of every input column. For an
+ * `ArrayColumnSpec` we emit `InputArray_${path}`; for a `StructColumnSpec`
+ * `InputStruct_${path}`; for a `MapColumnSpec` `InputMap_${path}` plus the `InputArray` classes
+ * for the key and value slices (because Spark's `MapData.keyArray()` / `valueArray()` return
+ * `ArrayData` - same view shape as any other array).
+ */
+ def emitNestedClasses(inputSchema: Seq[ArrowColumnSpec]): String = {
+ val out = new mutable.ArrayBuffer[String]()
+ inputSchema.zipWithIndex.foreach { case (spec, ord) =>
+ collectNestedClasses(s"col$ord", spec, out)
+ }
+ out.mkString("\n")
+ }
+
+ private def collectNestedClasses(
+ path: String,
+ spec: ArrowColumnSpec,
+ out: mutable.ArrayBuffer[String]): Unit = spec match {
+ case _: ScalarColumnSpec => ()
+ case ar: ArrayColumnSpec =>
+ out += emitArrayClass(path, ar)
+ collectNestedClasses(s"${path}_e", ar.element, out)
+ case st: StructColumnSpec =>
+ out += emitStructClass(path, st)
+ st.fields.zipWithIndex.foreach { case (f, fi) =>
+ collectNestedClasses(s"${path}_f$fi", f.child, out)
+ }
+ case mp: MapColumnSpec =>
+ out += emitMapClass(path, mp)
+ // Emit InputArray_${path}_k and InputArray_${path}_v - the ArrayData views returned by
+ // `MapData.keyArray()` / `valueArray()`. They follow the standard array-element
+ // convention: each reads from `${classPath}_e` which maps to the key / value vector
+ // emitted at `${path}_k_e` / `${path}_v_e` by [[collectVectorFieldDecls]]. Instance
+ // fields for complex key / value elements (one level deeper) live inside these array
+ // classes via [[instanceDeclaration]].
+ out += emitArrayClass(
+ s"${path}_k",
+ ArrayColumnSpec(nullable = true, elementSparkType = mp.keySparkType, element = mp.key))
+ out += emitArrayClass(
+ s"${path}_v",
+ ArrayColumnSpec(
+ nullable = true,
+ elementSparkType = mp.valueSparkType,
+ element = mp.value))
+ // Recurse into the key / value specs at their canonical paths (${path}_k_e /
+ // ${path}_v_e) so nested complex keys / values get their own nested classes.
+ collectNestedClasses(s"${path}_k_e", mp.key, out)
+ collectNestedClasses(s"${path}_v_e", mp.value, out)
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Shared helpers for complex-getter routing. A "list-backed child reset" computes
+ // `(startIdx, length)` for an inner instance from a ListVector / MapVector's offsets at a
+ // parent-provided index and calls `reset(startIdx, length)`.
+ // ---------------------------------------------------------------------------------------------
+
+ private def emitListBackedChildReset(
+ parentVectorPath: String,
+ indexExpr: String,
+ innerInstanceField: String): String =
+ s""" int __idx = $indexExpr;
+ | int __s = $parentVectorPath.getElementStartIndex(__idx);
+ | int __e = $parentVectorPath.getElementEndIndex(__idx);
+ | $innerInstanceField.reset(__s, __e - __s);""".stripMargin
+
+ /**
+ * Emit one `InputArray_${path}` nested class. Unified slice-based reset: callers pass
+ * `(startIdx, length)` directly.
+ *
+ * Key/value arrays of a map share this exact shape - the instance fields for their complex
+ * elements (if any) are emitted from [[emitArrayElementGetter]]; the vector fields they read
+ * from are at `${path}_e` (following the array-element path convention), which maps to
+ * `col${N}_k_e` or `col${N}_v_e` when the array represents a map key/value slice.
+ *
+ * NOTE: when this class is used for a map's key or value view and the underlying key/value is
+ * scalar, there is no `${path}_e` vector field - the map's key/value vector sits at `${path}`
+ * itself (e.g. `col0_k`). See [[emitArrayElementGetter]] for how that is handled: scalar
+ * element emission reads from `${path}_e`, but for map views the element vector IS the path
+ * itself. We rename the element path in [[emitMapClass]] below.
+ */
+ private def emitArrayClass(path: String, spec: ArrayColumnSpec): String = {
+ val baseClassName = classOf[CometArrayData].getName
+ val elemPath = s"${path}_e"
+ val innerInstance = instanceDeclaration(elemPath, spec.element)
+ val isNullAt =
+ s""" @Override
+ | public boolean isNullAt(int i) {
+ | return $elemPath.${nullCheckMethod(spec.element)}(startIndex + i);
+ | }""".stripMargin
+ val elementGetter = emitArrayElementGetter(path, spec)
+ s""" private final class InputArray_$path extends $baseClassName {
+ | private int startIndex;
+ | private int length;
+ |$innerInstance
+ |
+ | void reset(int startIdx, int len) {
+ | this.startIndex = startIdx;
+ | this.length = len;
+ | }
+ |
+ | @Override
+ | public int numElements() {
+ | return length;
+ | }
+ |
+ |$isNullAt
+ |
+ |$elementGetter
+ | }
+ |""".stripMargin
+ }
+
+ /**
+ * Emit the element getter body for a nested `InputArray_${path}`. Scalar element -> direct
+ * typed read. Complex element -> `getArray(i)` / `getStruct(i, n)` / `getMap(i)` that resets
+ * the inner instance.
+ */
+ private def emitArrayElementGetter(path: String, spec: ArrayColumnSpec): String = {
+ val elemPath = s"${path}_e"
+ spec.element match {
+ case _: ScalarColumnSpec =>
+ emitArrayElementScalarGetter(spec.elementSparkType, elemPath)
+ case _: ArrayColumnSpec =>
+ val reset = emitListBackedChildReset(elemPath, "startIndex + i", s"${elemPath}_arrayData")
+ s""" @Override
+ | public org.apache.spark.sql.catalyst.util.ArrayData getArray(int i) {
+ |$reset
+ | return ${elemPath}_arrayData;
+ | }""".stripMargin
+ case _: StructColumnSpec =>
+ s""" @Override
+ | public org.apache.spark.sql.catalyst.InternalRow getStruct(int i, int numFields) {
+ | ${elemPath}_structData.reset(startIndex + i);
+ | return ${elemPath}_structData;
+ | }""".stripMargin
+ case _: MapColumnSpec =>
+ val reset = emitListBackedChildReset(elemPath, "startIndex + i", s"${elemPath}_mapData")
+ s""" @Override
+ | public org.apache.spark.sql.catalyst.util.MapData getMap(int i) {
+ |$reset
+ | return ${elemPath}_mapData;
+ | }""".stripMargin
+ }
+ }
+
+ /**
+ * Emit the scalar-element getter override for a nested `InputArray_${path}`. Only the getter
+ * matching the element type is overridden; any other getter inherits the base class's
+ * `UnsupportedOperationException`.
+ */
+ private def emitArrayElementScalarGetter(elemType: DataType, childField: String): String =
+ elemType match {
+ case BooleanType =>
+ s""" @Override
+ | public boolean getBoolean(int i) {
+ | return $childField.getBoolean(startIndex + i);
+ | }""".stripMargin
+ case ByteType =>
+ s""" @Override
+ | public byte getByte(int i) {
+ | return $childField.getByte(startIndex + i);
+ | }""".stripMargin
+ case ShortType =>
+ s""" @Override
+ | public short getShort(int i) {
+ | return $childField.getShort(startIndex + i);
+ | }""".stripMargin
+ case IntegerType | DateType =>
+ s""" @Override
+ | public int getInt(int i) {
+ | return $childField.getInt(startIndex + i);
+ | }""".stripMargin
+ case LongType | TimestampType | TimestampNTZType =>
+ s""" @Override
+ | public long getLong(int i) {
+ | return $childField.getLong(startIndex + i);
+ | }""".stripMargin
+ case FloatType =>
+ s""" @Override
+ | public float getFloat(int i) {
+ | return $childField.getFloat(startIndex + i);
+ | }""".stripMargin
+ case DoubleType =>
+ s""" @Override
+ | public double getDouble(int i) {
+ | return $childField.getDouble(startIndex + i);
+ | }""".stripMargin
+ case dt: DecimalType =>
+ val body =
+ if (dt.precision <= 18) {
+ emitDecimalFastBodyUnsafe(s"${childField}_valueAddr", "startIndex + i", " ")
+ } else {
+ emitDecimalSlowBody(childField, "startIndex + i", " ")
+ }
+ s""" @Override
+ | public org.apache.spark.sql.types.Decimal getDecimal(
+ | int i, int precision, int scale) {
+ |$body
+ | }""".stripMargin
+ case _: StringType =>
+ s""" @Override
+ | public org.apache.spark.unsafe.types.UTF8String getUTF8String(int i) {
+ |${emitUtf8BodyUnsafe(
+ s"${childField}_valueAddr",
+ s"${childField}_offsetAddr",
+ "startIndex + i",
+ " ")}
+ | }""".stripMargin
+ case BinaryType =>
+ s""" @Override
+ | public byte[] getBinary(int i) {
+ |${emitBinaryBodyUnsafe(
+ s"${childField}_valueAddr",
+ s"${childField}_offsetAddr",
+ "startIndex + i",
+ " ")}
+ | }""".stripMargin
+ case other =>
+ throw new UnsupportedOperationException(
+ s"nested ArrayData: unsupported element type $other")
+ }
+
+ /**
+ * Emit the kernel's `@Override public ArrayData getArray(int ordinal)` method. Each case reads
+ * `(startIdx, length)` from the outer `ListVector`'s offsets at the current row and calls the
+ * pre-allocated instance's unified `reset(startIdx, length)`.
+ */
+ def emitGetArrayMethod(inputSchema: Seq[ArrowColumnSpec]): String = {
+ val cases = inputSchema.zipWithIndex.collect { case (_: ArrayColumnSpec, ord) =>
+ val reset =
+ emitListBackedChildReset(s"this.col$ord", "this.rowIdx", s"this.col${ord}_arrayData")
+ s""" case $ord: {
+ |$reset
+ | return this.col${ord}_arrayData;
+ | }""".stripMargin
+ }
+ if (cases.isEmpty) {
+ ""
+ } else {
+ s"""
+ | @Override
+ | public org.apache.spark.sql.catalyst.util.ArrayData getArray(int ordinal) {
+ | switch (ordinal) {
+ |${cases.mkString("\n")}
+ | default: throw new UnsupportedOperationException(
+ | "getArray out of range: " + ordinal);
+ | }
+ | }
+ |""".stripMargin
+ }
+ }
+
+ /**
+ * Emit one `InputStruct_${path}` nested class. Flat-indexed: `reset(int outerRowIdx)` just
+ * captures the index. Scalar getters switch on field ordinal; complex getters route to inner
+ * instances (offsets computed for array/map children; rowIdx passed through for struct
+ * children).
+ */
+ private def emitStructClass(path: String, spec: StructColumnSpec): String = {
+ val baseClassName = classOf[CometInternalRow].getName
+ val innerInstances = spec.fields.zipWithIndex
+ .flatMap { case (f, fi) =>
+ val fieldPath = s"${path}_f$fi"
+ Some(instanceDeclaration(fieldPath, f.child)).filter(_.nonEmpty)
+ }
+ .mkString("\n")
+ val isNullCases = spec.fields.zipWithIndex.map {
+ case (f, fi) if !f.nullable =>
+ s" case $fi: return false;"
+ case (f, fi) =>
+ s" case $fi: return ${path}_f$fi.${nullCheckMethod(f.child)}(this.rowIdx);"
+ }
+ val scalarGetters = emitStructScalarGetters(path, spec)
+ val complexGetters = emitStructComplexGetters(path, spec)
+ s""" private final class InputStruct_$path extends $baseClassName {
+ | private int rowIdx;
+ |$innerInstances
+ |
+ | void reset(int outerRowIdx) {
+ | this.rowIdx = outerRowIdx;
+ | }
+ |
+ | @Override
+ | public int numFields() {
+ | return ${spec.fields.length};
+ | }
+ |
+ | @Override
+ | public boolean isNullAt(int ordinal) {
+ | switch (ordinal) {
+ |${isNullCases.mkString("\n")}
+ | default: throw new UnsupportedOperationException(
+ | "InputStruct_$path.isNullAt out of range: " + ordinal);
+ | }
+ | }
+ |
+ |$scalarGetters
+ |$complexGetters
+ | }
+ |""".stripMargin
+ }
+
+ private def emitStructScalarGetters(path: String, spec: StructColumnSpec): String = {
+ val withOrd = spec.fields.zipWithIndex
+ val scalarOrd = withOrd.filter { case (f, _) => f.child.isInstanceOf[ScalarColumnSpec] }
+
+ def fieldReadScalar(fi: Int, dt: DataType): String = dt match {
+ case BooleanType =>
+ s" case $fi: return ${path}_f$fi.getBoolean(this.rowIdx);"
+ case ByteType =>
+ s" case $fi: return ${path}_f$fi.getByte(this.rowIdx);"
+ case ShortType =>
+ s" case $fi: return ${path}_f$fi.getShort(this.rowIdx);"
+ case IntegerType | DateType =>
+ s" case $fi: return ${path}_f$fi.getInt(this.rowIdx);"
+ case LongType | TimestampType | TimestampNTZType =>
+ s" case $fi: return ${path}_f$fi.getLong(this.rowIdx);"
+ case FloatType =>
+ s" case $fi: return ${path}_f$fi.getFloat(this.rowIdx);"
+ case DoubleType =>
+ s" case $fi: return ${path}_f$fi.getDouble(this.rowIdx);"
+ case BinaryType =>
+ s""" case $fi: {
+ |${emitBinaryBodyUnsafe(
+ s"${path}_f${fi}_valueAddr",
+ s"${path}_f${fi}_offsetAddr",
+ "this.rowIdx",
+ " ")}
+ | }""".stripMargin
+ case _: StringType =>
+ s""" case $fi: {
+ |${emitUtf8BodyUnsafe(
+ s"${path}_f${fi}_valueAddr",
+ s"${path}_f${fi}_offsetAddr",
+ "this.rowIdx",
+ " ")}
+ | }""".stripMargin
+ case _: DecimalType =>
+ throw new IllegalStateException("decimal handled separately")
+ case other =>
+ throw new UnsupportedOperationException(
+ s"nested InputStruct getter: unsupported field type $other")
+ }
+
+ val booleanCases =
+ scalarOrd.collect {
+ case (f, fi) if f.sparkType == BooleanType =>
+ fieldReadScalar(fi, BooleanType)
+ }
+ val byteCases =
+ scalarOrd.collect {
+ case (f, fi) if f.sparkType == ByteType =>
+ fieldReadScalar(fi, ByteType)
+ }
+ val shortCases =
+ scalarOrd.collect {
+ case (f, fi) if f.sparkType == ShortType =>
+ fieldReadScalar(fi, ShortType)
+ }
+ val intCases = scalarOrd.collect {
+ case (f, fi) if f.sparkType == IntegerType || f.sparkType == DateType =>
+ fieldReadScalar(fi, IntegerType)
+ }
+ val longCases = scalarOrd.collect {
+ case (f, fi)
+ if f.sparkType == LongType || f.sparkType == TimestampType ||
+ f.sparkType == TimestampNTZType =>
+ fieldReadScalar(fi, LongType)
+ }
+ val floatCases =
+ scalarOrd.collect {
+ case (f, fi) if f.sparkType == FloatType =>
+ fieldReadScalar(fi, FloatType)
+ }
+ val doubleCases =
+ scalarOrd.collect {
+ case (f, fi) if f.sparkType == DoubleType =>
+ fieldReadScalar(fi, DoubleType)
+ }
+ val binaryCases =
+ scalarOrd.collect {
+ case (f, fi) if f.sparkType == BinaryType =>
+ fieldReadScalar(fi, BinaryType)
+ }
+ val utf8Cases = scalarOrd.collect {
+ case (f, fi) if f.sparkType.isInstanceOf[StringType] => fieldReadScalar(fi, f.sparkType)
+ }
+
+ val decimalCases = scalarOrd.collect {
+ case (f, fi) if f.sparkType.isInstanceOf[DecimalType] =>
+ val dt = f.sparkType.asInstanceOf[DecimalType]
+ val field = s"${path}_f$fi"
+ val body =
+ if (dt.precision <= 18) {
+ emitDecimalFastBodyUnsafe(s"${field}_valueAddr", "this.rowIdx", " ")
+ } else {
+ emitDecimalSlowBody(field, "this.rowIdx", " ")
+ }
+ s""" case $fi: {
+ |$body
+ | }""".stripMargin
+ }
+
+ Seq(
+ structSwitch("public boolean getBoolean(int ordinal)", "getBoolean", booleanCases),
+ structSwitch("public byte getByte(int ordinal)", "getByte", byteCases),
+ structSwitch("public short getShort(int ordinal)", "getShort", shortCases),
+ structSwitch("public int getInt(int ordinal)", "getInt", intCases),
+ structSwitch("public long getLong(int ordinal)", "getLong", longCases),
+ structSwitch("public float getFloat(int ordinal)", "getFloat", floatCases),
+ structSwitch("public double getDouble(int ordinal)", "getDouble", doubleCases),
+ structSwitch(
+ "public org.apache.spark.sql.types.Decimal getDecimal(" +
+ "int ordinal, int precision, int scale)",
+ "getDecimal",
+ decimalCases),
+ structSwitch("public byte[] getBinary(int ordinal)", "getBinary", binaryCases),
+ structSwitch(
+ "public org.apache.spark.unsafe.types.UTF8String getUTF8String(int ordinal)",
+ "getUTF8String",
+ utf8Cases)).mkString
+ }
+
+ private def emitStructComplexGetters(path: String, spec: StructColumnSpec): String = {
+ val getArrayCases = spec.fields.zipWithIndex.collect {
+ case (f, fi) if f.child.isInstanceOf[ArrayColumnSpec] =>
+ val fieldPath = s"${path}_f$fi"
+ val reset = emitListBackedChildReset(fieldPath, "this.rowIdx", s"${fieldPath}_arrayData")
+ s""" case $fi: {
+ |$reset
+ | return ${fieldPath}_arrayData;
+ | }""".stripMargin
+ }
+ val getStructCases = spec.fields.zipWithIndex.collect {
+ case (f, fi) if f.child.isInstanceOf[StructColumnSpec] =>
+ val fieldPath = s"${path}_f$fi"
+ s""" case $fi: {
+ | ${fieldPath}_structData.reset(this.rowIdx);
+ | return ${fieldPath}_structData;
+ | }""".stripMargin
+ }
+ val getMapCases = spec.fields.zipWithIndex.collect {
+ case (f, fi) if f.child.isInstanceOf[MapColumnSpec] =>
+ val fieldPath = s"${path}_f$fi"
+ val reset = emitListBackedChildReset(fieldPath, "this.rowIdx", s"${fieldPath}_mapData")
+ s""" case $fi: {
+ |$reset
+ | return ${fieldPath}_mapData;
+ | }""".stripMargin
+ }
+ Seq(
+ structSwitch(
+ "public org.apache.spark.sql.catalyst.util.ArrayData getArray(int ordinal)",
+ "getArray",
+ getArrayCases),
+ structSwitch(
+ "public org.apache.spark.sql.catalyst.InternalRow getStruct(int ordinal, int numFields)",
+ "getStruct",
+ getStructCases),
+ structSwitch(
+ "public org.apache.spark.sql.catalyst.util.MapData getMap(int ordinal)",
+ "getMap",
+ getMapCases)).mkString
+ }
+
+ /**
+ * Emit one `InputMap_${path}` nested class. Holds the slice `(startIndex, length)` and routes
+ * `keyArray()` / `valueArray()` through pre-allocated `InputArray_${path}_k` /
+ * `InputArray_${path}_v` instances (emitted by [[collectNestedClasses]]).
+ */
+ private def emitMapClass(path: String, spec: MapColumnSpec): String = {
+ val _ = spec // key/value arrays declared via path convention below
+ val baseClassName = classOf[CometMapData].getName
+ val keyPath = s"${path}_k"
+ val valPath = s"${path}_v"
+ s""" private final class InputMap_$path extends $baseClassName {
+ | private int startIndex;
+ | private int length;
+ | private final InputArray_$keyPath ${keyPath}_arrayData = new InputArray_$keyPath();
+ | private final InputArray_$valPath ${valPath}_arrayData = new InputArray_$valPath();
+ |
+ | void reset(int startIdx, int len) {
+ | this.startIndex = startIdx;
+ | this.length = len;
+ | }
+ |
+ | @Override
+ | public int numElements() {
+ | return length;
+ | }
+ |
+ | @Override
+ | public org.apache.spark.sql.catalyst.util.ArrayData keyArray() {
+ | ${keyPath}_arrayData.reset(this.startIndex, this.length);
+ | return ${keyPath}_arrayData;
+ | }
+ |
+ | @Override
+ | public org.apache.spark.sql.catalyst.util.ArrayData valueArray() {
+ | ${valPath}_arrayData.reset(this.startIndex, this.length);
+ | return ${valPath}_arrayData;
+ | }
+ | }
+ |""".stripMargin
+ }
+
+ /**
+ * Emit the kernel's top-level `@Override public MapData getMap(int ordinal)` method when the
+ * input schema has at least one map-typed column at the top level; empty string otherwise.
+ */
+ def emitGetMapMethod(inputSchema: Seq[ArrowColumnSpec]): String = {
+ val cases = inputSchema.zipWithIndex.collect { case (_: MapColumnSpec, ord) =>
+ val reset =
+ emitListBackedChildReset(s"this.col$ord", "this.rowIdx", s"this.col${ord}_mapData")
+ s""" case $ord: {
+ |$reset
+ | return this.col${ord}_mapData;
+ | }""".stripMargin
+ }
+ if (cases.isEmpty) {
+ ""
+ } else {
+ s"""
+ | @Override
+ | public org.apache.spark.sql.catalyst.util.MapData getMap(int ordinal) {
+ | switch (ordinal) {
+ |${cases.mkString("\n")}
+ | default: throw new UnsupportedOperationException(
+ | "getMap out of range: " + ordinal);
+ | }
+ | }
+ |""".stripMargin
+ }
+ }
+
+ /**
+ * Return the inner-instance field declaration for one complex spec at the given path, or an
+ * empty string for a scalar spec. Used inside nested-class bodies to declare pre-allocated
+ * child-view instances.
+ */
+ private def instanceDeclaration(path: String, spec: ArrowColumnSpec): String = spec match {
+ case _: ScalarColumnSpec => ""
+ case _: ArrayColumnSpec =>
+ s" private final InputArray_$path ${path}_arrayData = new InputArray_$path();"
+ case _: StructColumnSpec =>
+ s" private final InputStruct_$path ${path}_structData = new InputStruct_$path();"
+ case _: MapColumnSpec =>
+ s" private final InputMap_$path ${path}_mapData = new InputMap_$path();"
+ }
+
+ private def structSwitch(methodSig: String, label: String, cases: Seq[String]): String = {
+ if (cases.isEmpty) {
+ ""
+ } else {
+ s"""
+ | @Override
+ | $methodSig {
+ | switch (ordinal) {
+ |${cases.mkString("\n")}
+ | default: throw new UnsupportedOperationException(
+ | "$label out of range: " + ordinal);
+ | }
+ | }
+ """.stripMargin
+ }
+ }
+
+ /**
+ * Emit the kernel's top-level `@Override public InternalRow getStruct(int ordinal, int
+ * numFields)` method when the input schema has at least one struct-typed column.
+ */
+ def emitGetStructMethod(inputSchema: Seq[ArrowColumnSpec]): String = {
+ val cases = inputSchema.zipWithIndex.collect { case (_: StructColumnSpec, ord) =>
+ s""" case $ord: {
+ | this.col${ord}_structData.reset(this.rowIdx);
+ | return this.col${ord}_structData;
+ | }""".stripMargin
+ }
+ if (cases.isEmpty) {
+ ""
+ } else {
+ s"""
+ | @Override
+ | public org.apache.spark.sql.catalyst.InternalRow getStruct(int ordinal, int numFields) {
+ | switch (ordinal) {
+ |${cases.mkString("\n")}
+ | default: throw new UnsupportedOperationException(
+ | "getStruct out of range: " + ordinal);
+ | }
+ | }
+ |""".stripMargin
+ }
+ }
+
+ private def emitOrdinalSwitch(methodSig: String, label: String, cases: Seq[String]): String = {
+ if (cases.isEmpty) {
+ ""
+ } else {
+ s"""
+ | @Override
+ | $methodSig {
+ | switch (ordinal) {
+ |${cases.mkString("\n")}
+ | default: throw new UnsupportedOperationException(
+ | "$label out of range: " + ordinal);
+ | }
+ | }
+ """.stripMargin
+ }
+ }
+
+ // -------------------------------------------------------------------------------------------
+ // Scalar-read body templates. Each helper emits the per-type read statements parameterized on
+ // a Java expression for the row/slot index (`idx`), the cached buffer address(es) for unsafe
+ // reads (`valueAddr`, `offsetAddr`), or the Arrow typed field (`field`) for the slow-path
+ // decimal case that still needs `getObject`. `ind` is the per-line indent prefix;
+ // continuation lines add four spaces. Callers wrap the output in switch cases or method
+ // overrides.
+ //
+ // The VarChar / VarBinary unsafe emitters below duplicate what CometPlainVector.getUTF8String
+ // and getBinary do today, with two differences: they skip CometPlainVector's internal
+ // isNullAt (redundant here because the kernel's caller already handled it) and they read the
+ // offset-buffer address from a kernel-cached field rather than re-dereferencing the ArrowBuf.
+ // Once apache/datafusion-comet#4280 (offsetBufferAddress caching) and #4279 (validity-bitmap
+ // byte cache) land, both differences stop mattering and `emitUtf8BodyUnsafe` /
+ // `emitBinaryBodyUnsafe` can be deleted in favor of `CometPlainVector` reuse for variable-
+ // width. The decimal-fast variant has its own motivation (compile-time precision
+ // specialization) unrelated to those issues.
+ // -------------------------------------------------------------------------------------------
+
+ /** Parenthesize `idx` when it contains whitespace, to keep `(long) idx * 16L` well-formed. */
+ private def castableIdx(idx: String): String = if (idx.contains(' ')) s"($idx)" else idx
+
+ private def emitDecimalSlowBody(field: String, idx: String, ind: String): String = {
+ val cont = ind + " "
+ s"""${ind}java.math.BigDecimal bd = $field.getObject($idx);
+ |${ind}return org.apache.spark.sql.types.Decimal$$.MODULE$$
+ |$cont.apply(bd, precision, scale);""".stripMargin
+ }
+
+ private def emitDecimalFastBodyUnsafe(valueAddr: String, idx: String, ind: String): String = {
+ val cont = ind + " "
+ val i = castableIdx(idx)
+ s"""${ind}long unscaled = org.apache.spark.unsafe.Platform.getLong(null,
+ |$cont$valueAddr + (long) $i * 16L);
+ |${ind}return org.apache.spark.sql.types.Decimal$$.MODULE$$
+ |$cont.createUnsafe(unscaled, precision, scale);""".stripMargin
+ }
+
+ private def emitUtf8BodyUnsafe(
+ valueAddr: String,
+ offsetAddr: String,
+ idx: String,
+ ind: String): String = {
+ val cont = ind + " "
+ val i = castableIdx(idx)
+ s"""${ind}int s = org.apache.spark.unsafe.Platform.getInt(null,
+ |$cont$offsetAddr + (long) $i * 4L);
+ |${ind}int e = org.apache.spark.unsafe.Platform.getInt(null,
+ |$cont$offsetAddr + ((long) $i + 1L) * 4L);
+ |${ind}return org.apache.spark.unsafe.types.UTF8String
+ |$cont.fromAddress(null, $valueAddr + s, e - s);""".stripMargin
+ }
+
+ private def emitBinaryBodyUnsafe(
+ valueAddr: String,
+ offsetAddr: String,
+ idx: String,
+ ind: String): String = {
+ val cont = ind + " "
+ val i = castableIdx(idx)
+ s"""${ind}int s = org.apache.spark.unsafe.Platform.getInt(null,
+ |$cont$offsetAddr + (long) $i * 4L);
+ |${ind}int e = org.apache.spark.unsafe.Platform.getInt(null,
+ |$cont$offsetAddr + ((long) $i + 1L) * 4L);
+ |${ind}int len = e - s;
+ |${ind}byte[] out = new byte[len];
+ |${ind}org.apache.spark.unsafe.Platform.copyMemory(null, $valueAddr + s, out,
+ |${cont}org.apache.spark.unsafe.Platform.BYTE_ARRAY_OFFSET, len);
+ |${ind}return out;""".stripMargin
+ }
+}
diff --git a/common/src/main/scala/org/apache/comet/udf/CometBatchKernelCodegenOutput.scala b/common/src/main/scala/org/apache/comet/udf/CometBatchKernelCodegenOutput.scala
new file mode 100644
index 0000000000..4dd4d02497
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/udf/CometBatchKernelCodegenOutput.scala
@@ -0,0 +1,370 @@
+/*
+ * 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.udf
+
+import org.apache.arrow.vector.{BigIntVector, BitVector, DateDayVector, DecimalVector, FieldVector, Float4Vector, Float8Vector, IntVector, SmallIntVector, TimeStampMicroTZVector, TimeStampMicroVector, TinyIntVector, VarBinaryVector, VarCharVector}
+import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampNTZType, TimestampType}
+
+import org.apache.comet.CometArrowAllocator
+
+/**
+ * Output-side emitters for the Arrow-direct codegen kernel. Everything that writes a computed
+ * value into an Arrow output vector lives here: [[allocateOutput]], [[emitOutputWriter]] (the
+ * entry point for the kernel's top-level write), [[emitWrite]] (recursive per-type write), the
+ * output vector-class lookup, and the output-side type-support gate.
+ *
+ * Paired with [[CometBatchKernelCodegenInput]], which handles the symmetric input side.
+ */
+private[udf] object CometBatchKernelCodegenOutput {
+
+ /**
+ * Output types [[allocateOutput]] and [[emitOutputWriter]] can materialize. Recursive: complex
+ * types are supported when their children are.
+ */
+ def isSupportedOutputType(dt: DataType): Boolean = dt match {
+ case BooleanType | ByteType | ShortType | IntegerType | LongType => true
+ case FloatType | DoubleType => true
+ case _: DecimalType => true
+ case _: StringType | _: BinaryType => true
+ case DateType | TimestampType | TimestampNTZType => true
+ case ArrayType(inner, _) => isSupportedOutputType(inner)
+ case st: StructType => st.fields.forall(f => isSupportedOutputType(f.dataType))
+ case mt: MapType =>
+ isSupportedOutputType(mt.keyType) && isSupportedOutputType(mt.valueType)
+ case _ => false
+ }
+
+ /**
+ * Allocate an Arrow output vector matching `dataType`. Delegates field and vector construction
+ * to [[Utils.toArrowField]] + `Field.createVector`, which is the pattern the rest of Comet uses
+ * to go Spark -> Arrow and handles complex-type wiring (including Arrow's non-null-key and
+ * non-null-entries invariants on `MapVector`).
+ *
+ * For variable-length scalar outputs (`StringType`, `BinaryType`), callers can pass
+ * `estimatedBytes` to pre-size the data buffer and avoid `setSafe` reallocation mid-loop. The
+ * hint is only applied when the root vector is `VarCharVector` or `VarBinaryVector`; inside a
+ * `ListVector` / `StructVector` / `MapVector`, the parent's `allocateNew` reallocates child
+ * buffers at default size, so a leaf hint would be lost.
+ *
+ * Closes the vector on any failure between construction and return so a partially-initialized
+ * tree does not leak buffers back to the allocator.
+ */
+ def allocateOutput(
+ dataType: DataType,
+ name: String,
+ numRows: Int,
+ estimatedBytes: Int = -1): FieldVector = {
+ val field = Utils.toArrowField(name, dataType, nullable = true, "UTC")
+ val vec = field.createVector(CometArrowAllocator).asInstanceOf[FieldVector]
+ try {
+ vec.setInitialCapacity(numRows)
+ vec match {
+ case v: VarCharVector if estimatedBytes > 0 =>
+ v.allocateNew(estimatedBytes.toLong, numRows)
+ case v: VarBinaryVector if estimatedBytes > 0 =>
+ v.allocateNew(estimatedBytes.toLong, numRows)
+ case _ =>
+ vec.allocateNew()
+ }
+ vec
+ } catch {
+ case t: Throwable =>
+ try vec.close()
+ catch { case _: Throwable => () }
+ throw t
+ }
+ }
+
+ /**
+ * Split output for a complex-type write: `setup` holds once-per-batch declarations (typed
+ * child-vector casts) and lives outside the per-row for-loop; `perRow` holds the statements
+ * executed for each row. Scalar writes have empty setup.
+ */
+ private case class OutputEmit(setup: String, perRow: String)
+
+ /**
+ * Returns `(concreteVectorClassName, batchSetup, perRowSnippet)` for the expression's output
+ * type at the root of the generated kernel. `output` is already cast to
+ * `concreteVectorClassName` in `process`'s prelude, so `emitWrite`'s complex-type branches can
+ * hoist child casts straight off `output` without re-casting it per row.
+ */
+ def emitOutputWriter(
+ dataType: DataType,
+ valueTerm: String,
+ ctx: CodegenContext): (String, String, String) = {
+ val cls = outputVectorClass(dataType)
+ val emit = emitWrite("output", "i", valueTerm, dataType, ctx)
+ (cls, emit.setup, emit.perRow)
+ }
+
+ /**
+ * Concrete Arrow vector class name for the given output type. The name is used to cast `outRaw`
+ * to the right type at the top of the generated `process` method, so that subsequent writes
+ * through `emitWrite` can call vector-specific methods without further casts.
+ */
+ private def outputVectorClass(dataType: DataType): String = dataType match {
+ case BooleanType => classOf[BitVector].getName
+ case ByteType => classOf[TinyIntVector].getName
+ case ShortType => classOf[SmallIntVector].getName
+ case IntegerType => classOf[IntVector].getName
+ case LongType => classOf[BigIntVector].getName
+ case FloatType => classOf[Float4Vector].getName
+ case DoubleType => classOf[Float8Vector].getName
+ case _: DecimalType => classOf[DecimalVector].getName
+ case _: StringType => classOf[VarCharVector].getName
+ case BinaryType => classOf[VarBinaryVector].getName
+ case DateType => classOf[DateDayVector].getName
+ case TimestampType => classOf[TimeStampMicroTZVector].getName
+ case TimestampNTZType => classOf[TimeStampMicroVector].getName
+ case _: ArrayType => classOf[ListVector].getName
+ case _: StructType => classOf[StructVector].getName
+ case _: MapType => classOf[MapVector].getName
+ case other =>
+ throw new UnsupportedOperationException(
+ s"CometBatchKernelCodegen.outputVectorClass: unsupported output type $other")
+ }
+
+ /**
+ * Composable write emitter. Returns an [[OutputEmit]] whose `setup` declares once-per-batch
+ * typed child-vector casts (hoisted above the `process` for-loop) and whose `perRow` writes the
+ * value produced by `source` into `targetVec` at index `idx`. `targetVec` is assumed to be
+ * already typed to the concrete Arrow vector class for `dataType` at the call site (via the
+ * prelude cast in `process` for the root, or via a setup cast declared by the caller for nested
+ * children).
+ *
+ * Scalars emit `perRow` only; complex types (`ArrayType` / `StructType` / `MapType`) emit both
+ * setup (child-vector casts) and perRow (loops, null guards, recursive writes). Inner
+ * `emitWrite` calls return their own setup, which the outer caller concatenates so child-of-
+ * child casts bubble up to the batch prelude.
+ */
+ private def emitWrite(
+ targetVec: String,
+ idx: String,
+ source: String,
+ dataType: DataType,
+ ctx: CodegenContext): OutputEmit = dataType match {
+ case BooleanType =>
+ OutputEmit("", s"$targetVec.set($idx, $source ? 1 : 0);")
+ case ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | DateType |
+ TimestampType | TimestampNTZType =>
+ // All scalar primitives and date/time types share the direct `set(idx, value)` shape.
+ // Spark's codegen already emits the correct primitive Java type for each; Arrow's
+ // typed vectors accept the matching primitive in their `set` overloads.
+ OutputEmit("", s"$targetVec.set($idx, $source);")
+ case dt: DecimalType =>
+ // Optimization: DecimalOutputShortFastPath.
+ // For precision <= 18 the unscaled value fits in a signed long; pass it straight to
+ // `DecimalVector.setSafe(int, long)` and skip the `java.math.BigDecimal` allocation
+ // `setSafe(int, BigDecimal)` requires. For p > 18 the BigDecimal path is unavoidable.
+ val write =
+ if (dt.precision <= 18) s"$targetVec.setSafe($idx, $source.toUnscaledLong());"
+ else s"$targetVec.setSafe($idx, $source.toJavaBigDecimal());"
+ OutputEmit("", write)
+ case _: StringType =>
+ // Optimization: Utf8OutputOnHeapShortcut.
+ // `UTF8String` is internally a `(base, offset, numBytes)` view. When the base is a
+ // `byte[]` (common case: Spark string functions allocate results on-heap), pass the
+ // existing byte[] directly to `VarCharVector.setSafe(int, byte[], int, int)` via the
+ // encoded offset and skip the redundant `getBytes()` allocation. Off-heap passthrough
+ // (rare on output side) falls back to `getBytes()`.
+ val bBase = ctx.freshName("utfBase")
+ val bLen = ctx.freshName("utfLen")
+ val bArr = ctx.freshName("utfArr")
+ OutputEmit(
+ "",
+ s"""Object $bBase = $source.getBaseObject();
+ |int $bLen = $source.numBytes();
+ |if ($bBase instanceof byte[]) {
+ | $targetVec.setSafe($idx, (byte[]) $bBase,
+ | (int) ($source.getBaseOffset()
+ | - org.apache.spark.unsafe.Platform.BYTE_ARRAY_OFFSET),
+ | $bLen);
+ |} else {
+ | byte[] $bArr = $source.getBytes();
+ | $targetVec.setSafe($idx, $bArr, 0, $bArr.length);
+ |}""".stripMargin)
+ case BinaryType =>
+ // Spark's BinaryType value is already a `byte[]`.
+ OutputEmit("", s"$targetVec.setSafe($idx, $source, 0, $source.length);")
+ case ArrayType(elementType, _) =>
+ // Complex-type output: recursive per-row write.
+ // Spark's `doGenCode` for ArrayType-returning expressions produces an `ArrayData` value
+ // (usually `GenericArrayData` / `UnsafeArrayData`). We iterate its elements, write each
+ // one into the Arrow `ListVector`'s child, and bracket with `startNewValue` /
+ // `endValue`. The element write recurses through `emitWrite` on the list's child vector,
+ // so any scalar we support becomes a valid array element. Nested complex types (Array of
+ // Array, Array of Struct) work by the same recursion. `targetVec` is a `ListVector` at
+ // the call site (either `output` at root or a hoisted child cast); we only need to cast
+ // its data vector, and that cast goes into setup.
+ val childVar = ctx.freshName("outListChild")
+ val childClass = outputVectorClass(elementType)
+ val arrVar = ctx.freshName("arr")
+ val nVar = ctx.freshName("n")
+ val childIdx = ctx.freshName("cidx")
+ val jVar = ctx.freshName("j")
+ val elemSource = emitSpecializedGetterExpr(arrVar, jVar, elementType)
+ val inner = emitWrite(childVar, s"$childIdx + $jVar", elemSource, elementType, ctx)
+ val setup =
+ (s"$childClass $childVar = ($childClass) $targetVec.getDataVector();" +:
+ Seq(inner.setup).filter(_.nonEmpty)).mkString("\n")
+ val perRow =
+ s"""org.apache.spark.sql.catalyst.util.ArrayData $arrVar = $source;
+ |int $nVar = $arrVar.numElements();
+ |int $childIdx = $targetVec.startNewValue($idx);
+ |for (int $jVar = 0; $jVar < $nVar; $jVar++) {
+ | if ($arrVar.isNullAt($jVar)) {
+ | $childVar.setNull($childIdx + $jVar);
+ | } else {
+ | ${inner.perRow}
+ | }
+ |}
+ |$targetVec.endValue($idx, $nVar);""".stripMargin
+ OutputEmit(setup, perRow)
+ case st: StructType =>
+ // Complex-type output: recursive per-row write to a StructVector.
+ // Spark's `doGenCode` for StructType-returning expressions produces an `InternalRow`
+ // value (`GenericInternalRow` / `UnsafeRow` / ScalaUDF encoder output). Typed child-vector
+ // casts are hoisted to setup (once per batch); the per-row body references the hoisted
+ // names. `StructVector` writes are flat-indexed (same `$idx` as the struct's outer slot).
+ //
+ // Branchless optimization: for each field whose `nullable == false` on the
+ // [[StructType]], we skip the `row.isNullAt($fi)` guard at source level. Non-nullable
+ // fields in Spark are a contract that the producer does not emit nulls for that field,
+ // and matching that contract here lets HotSpot emit a straight write path per field
+ // rather than a branch.
+ val rowVar = ctx.freshName("row")
+ val perField = st.fields.zipWithIndex.map { case (field, fi) =>
+ val childVar = ctx.freshName("outStructChild")
+ val childClass = outputVectorClass(field.dataType)
+ val childDecl =
+ s"$childClass $childVar = ($childClass) $targetVec.getChildByOrdinal($fi);"
+ val fieldSource = emitSpecializedGetterExpr(rowVar, fi.toString, field.dataType)
+ val inner = emitWrite(childVar, idx, fieldSource, field.dataType, ctx)
+ val write =
+ if (!field.nullable) {
+ inner.perRow
+ } else {
+ s"""if ($rowVar.isNullAt($fi)) {
+ | $childVar.setNull($idx);
+ |} else {
+ | ${inner.perRow}
+ |}""".stripMargin
+ }
+ val perFieldSetup = (Seq(childDecl) ++ Seq(inner.setup).filter(_.nonEmpty)).mkString("\n")
+ (perFieldSetup, write)
+ }
+ val setup = perField.map(_._1).mkString("\n")
+ val perFieldWrites = perField.map(_._2).mkString("\n")
+ val perRow =
+ s"""org.apache.spark.sql.catalyst.InternalRow $rowVar = $source;
+ |$targetVec.setIndexDefined($idx);
+ |$perFieldWrites""".stripMargin
+ OutputEmit(setup, perRow)
+ case mt: MapType =>
+ // Complex-type output: recursive per-row write to a MapVector.
+ // Spark's `doGenCode` for MapType-returning expressions produces a `MapData` value
+ // (`ArrayBasedMapData` / `UnsafeMapData` / ScalaUDF encoder output). Typed child-vector
+ // casts for the entries struct and the key/value children are hoisted to setup (once per
+ // batch); the per-row body references them.
+ //
+ // Per-row shape:
+ // 1. Read keyArray / valueArray from the MapData source.
+ // 2. Open a new map entry via `startNewValue(idx)`; returns the base index into the
+ // entries StructVector for this row's key/value pairs.
+ // 3. For each key/value pair: set the entries struct slot defined (map values can be
+ // null, but the struct slot itself is defined), write the key (always non-null by
+ // Spark/Arrow invariant), then write the value with a null-guard on
+ // `vals.isNullAt(j)`. Both writes recurse through `emitWrite`.
+ // 4. Close the map entry with `endValue(idx, n)`.
+ val entriesVar = ctx.freshName("outMapEntries")
+ val keyVar = ctx.freshName("outMapKey")
+ val valVar = ctx.freshName("outMapVal")
+ val mapSrc = ctx.freshName("mapSrc")
+ val keyArr = ctx.freshName("keyArr")
+ val valArr = ctx.freshName("valArr")
+ val nVar = ctx.freshName("n")
+ val childIdx = ctx.freshName("cidx")
+ val jVar = ctx.freshName("j")
+ val structClass = classOf[StructVector].getName
+ val keyClass = outputVectorClass(mt.keyType)
+ val valClass = outputVectorClass(mt.valueType)
+ val keySrcExpr = emitSpecializedGetterExpr(keyArr, jVar, mt.keyType)
+ val valSrcExpr = emitSpecializedGetterExpr(valArr, jVar, mt.valueType)
+ val keyEmit = emitWrite(keyVar, s"$childIdx + $jVar", keySrcExpr, mt.keyType, ctx)
+ val valEmit = emitWrite(valVar, s"$childIdx + $jVar", valSrcExpr, mt.valueType, ctx)
+ val setup =
+ (Seq(
+ s"$structClass $entriesVar = ($structClass) $targetVec.getDataVector();",
+ s"$keyClass $keyVar = ($keyClass) $entriesVar.getChildByOrdinal(0);",
+ s"$valClass $valVar = ($valClass) $entriesVar.getChildByOrdinal(1);") ++
+ Seq(keyEmit.setup, valEmit.setup).filter(_.nonEmpty)).mkString("\n")
+ val perRow =
+ s"""org.apache.spark.sql.catalyst.util.MapData $mapSrc = $source;
+ |org.apache.spark.sql.catalyst.util.ArrayData $keyArr = $mapSrc.keyArray();
+ |org.apache.spark.sql.catalyst.util.ArrayData $valArr = $mapSrc.valueArray();
+ |int $nVar = $mapSrc.numElements();
+ |int $childIdx = $targetVec.startNewValue($idx);
+ |for (int $jVar = 0; $jVar < $nVar; $jVar++) {
+ | $entriesVar.setIndexDefined($childIdx + $jVar);
+ | ${keyEmit.perRow}
+ | if ($valArr.isNullAt($jVar)) {
+ | $valVar.setNull($childIdx + $jVar);
+ | } else {
+ | ${valEmit.perRow}
+ | }
+ |}
+ |$targetVec.endValue($idx, $nVar);""".stripMargin
+ OutputEmit(setup, perRow)
+ case other =>
+ throw new UnsupportedOperationException(
+ s"CometBatchKernelCodegen.emitWrite: unsupported output type $other")
+ }
+
+ /**
+ * Java expression that reads a typed value out of a Spark `SpecializedGetters` reference (which
+ * both `ArrayData` and `InternalRow` implement) at a given ordinal/index. Used by the
+ * `ArrayType` and `StructType` branches of [[emitWrite]] to source each element / field for its
+ * recursive inner write.
+ */
+ private def emitSpecializedGetterExpr(target: String, idx: String, elemType: DataType): String =
+ elemType match {
+ case BooleanType => s"$target.getBoolean($idx)"
+ case ByteType => s"$target.getByte($idx)"
+ case ShortType => s"$target.getShort($idx)"
+ case IntegerType | DateType => s"$target.getInt($idx)"
+ case LongType | TimestampType | TimestampNTZType => s"$target.getLong($idx)"
+ case FloatType => s"$target.getFloat($idx)"
+ case DoubleType => s"$target.getDouble($idx)"
+ case dt: DecimalType => s"$target.getDecimal($idx, ${dt.precision}, ${dt.scale})"
+ case _: StringType => s"$target.getUTF8String($idx)"
+ case BinaryType => s"$target.getBinary($idx)"
+ case ArrayType(_, _) => s"$target.getArray($idx)"
+ case _: MapType => s"$target.getMap($idx)"
+ case _: StructType =>
+ val numFields = elemType.asInstanceOf[StructType].fields.length
+ s"$target.getStruct($idx, $numFields)"
+ case other =>
+ throw new UnsupportedOperationException(
+ s"CometBatchKernelCodegen.emitSpecializedGetterExpr: unsupported type $other")
+ }
+}
diff --git a/common/src/main/scala/org/apache/comet/udf/CometCodegenDispatchUDF.scala b/common/src/main/scala/org/apache/comet/udf/CometCodegenDispatchUDF.scala
new file mode 100644
index 0000000000..5be5dc25d5
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/udf/CometCodegenDispatchUDF.scala
@@ -0,0 +1,358 @@
+/*
+ * 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.udf
+
+import java.nio.ByteBuffer
+import java.util.{Collections, LinkedHashMap}
+import java.util.concurrent.atomic.AtomicLong
+
+import org.apache.arrow.vector.{BigIntVector, BitVector, DateDayVector, DecimalVector, Float4Vector, Float8Vector, IntVector, SmallIntVector, TimeStampMicroTZVector, TimeStampMicroVector, TinyIntVector, ValueVector, VarBinaryVector, VarCharVector}
+import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.types.{BinaryType, DataType, StringType}
+
+import org.apache.comet.udf.CometBatchKernelCodegen.{ArrayColumnSpec, ArrowColumnSpec, MapColumnSpec, ScalarColumnSpec, StructColumnSpec, StructFieldSpec}
+
+/**
+ * Arrow-direct codegen dispatcher. For each (bound Spark `Expression`, input Arrow schema) pair,
+ * compiles a specialized [[CometBatchKernel]] on first encounter and caches the compile.
+ * Subsequent batches with the same expression and schema reuse the cached compile.
+ *
+ * Arg 0 is a `VarBinaryVector` scalar carrying the serialized Expression bytes (produced on the
+ * driver by Spark's closure serializer). Args 1..N are the data columns the `BoundReference`s
+ * refer to, in ordinal order. The bytes self-describe the expression so the path works in cluster
+ * mode without executor-side state.
+ *
+ * Three caches compose at different scopes: the JVM-wide compile cache on the companion
+ * (`kernelCache`); a per-thread UDF instance map in `CometUdfBridge.INSTANCES`; and per-partition
+ * kernel instance state on this object (`activeKernel`, `activeKey`, `activePartition`) managed
+ * by [[ensureKernel]]. See `docs/source/contributor-guide/jvm_udf_dispatch.md` for the rationale
+ * and why none of the layers can be collapsed.
+ */
+class CometCodegenDispatchUDF extends CometUDF {
+
+ override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector = {
+ require(
+ inputs.length >= 1,
+ "CometCodegenDispatchUDF requires at least 1 input (serialized expression), " +
+ s"got ${inputs.length}")
+ val exprVec = inputs(0).asInstanceOf[VarBinaryVector]
+ require(
+ exprVec.getValueCount >= 1 && !exprVec.isNull(0),
+ "CometCodegenDispatchUDF requires non-null serialized expression bytes at arg 0")
+ val bytes = exprVec.get(0)
+
+ // TODO(dict-encoded): kernels assume materialized inputs; dict-encoded vectors would fail the
+ // cast in `specFor` below. See docs/source/contributor-guide/jvm_udf_dispatch.md#open-items.
+
+ val numDataCols = inputs.length - 1
+ val dataCols = new Array[ValueVector](numDataCols)
+ val specs = new Array[ArrowColumnSpec](numDataCols)
+ var di = 0
+ while (di < numDataCols) {
+ val v = inputs(di + 1)
+ dataCols(di) = v
+ specs(di) = specFor(v)
+ di += 1
+ }
+ val n = numRows
+ val specsSeq = specs.toIndexedSeq
+
+ val key = CometCodegenDispatchUDF.CacheKey(ByteBuffer.wrap(bytes), specsSeq)
+ val entry = CometCodegenDispatchUDF.lookupOrCompile(key, bytes, specsSeq)
+
+ val partitionId = CometCodegenDispatchUDF.currentPartitionIndex()
+ val kernel = ensureKernel(entry.compiled, key, partitionId)
+
+ val out = CometBatchKernelCodegen.allocateOutput(
+ entry.outputType,
+ "codegen_result",
+ n,
+ estimatedOutputBytes(entry.outputType, dataCols))
+ try {
+ kernel.process(dataCols, out, n)
+ out.setValueCount(n)
+ out
+ } catch {
+ case t: Throwable =>
+ try out.close()
+ catch { case _: Throwable => () }
+ throw t
+ }
+ }
+
+ /**
+ * Per-partition kernel instance cache. The dispatcher's compile cache (on the companion object)
+ * is JVM-wide and stores the compiled `GeneratedClass`. The kernel '''instance''', however,
+ * holds per-row mutable state for non-deterministic and stateful expressions (`Rand`'s
+ * `XORShiftRandom`, `MonotonicallyIncreasingID`'s counter, etc.). That state must advance
+ * across batches in one partition and reset across partitions. Allocating per batch (the prior
+ * model) reset state every batch and was wrong; allocating per partition is right.
+ *
+ * `CometCodegenDispatchUDF` is per-thread via `CometUdfBridge.INSTANCES`, and Spark tasks are
+ * single-threaded on a partition, so plain instance fields are safe without synchronisation. A
+ * different partition or a different cached expression flowing through the same thread triggers
+ * a fresh allocation; same partition + same expression reuses the kernel.
+ */
+ private var activeKernel: CometBatchKernel = _
+ private var activeKey: CometCodegenDispatchUDF.CacheKey = _
+ private var activePartition: Int = -1
+
+ private def ensureKernel(
+ compiled: CometBatchKernelCodegen.CompiledKernel,
+ key: CometCodegenDispatchUDF.CacheKey,
+ partitionId: Int): CometBatchKernel = {
+ if (activeKernel == null || activePartition != partitionId || activeKey != key) {
+ activeKernel = compiled.newInstance()
+ activeKernel.init(partitionId)
+ activeKey = key
+ activePartition = partitionId
+ }
+ activeKernel
+ }
+
+ /**
+ * Did any row in this Arrow vector set the null bit? The cache key carries this per column, so
+ * a batch with no nulls and a later batch with nulls map to different keys and different
+ * compiles, no correctness risk from flipping this. The tighter `nullable=false` compile lets
+ * the kernel emit `return false` from its `isNullAt` switch and, once paired with the
+ * BoundReference tree rewrite in `lookupOrCompile`, lets Spark's `BoundReference.genCode` skip
+ * the null branch at source level rather than relying on JIT constant-folding.
+ *
+ * Trade-off: if real workloads flip a column's nullability frequently across batches, each
+ * expression caches up to `2^numCols` variants and the bounded LRU churns. The common case is
+ * stable per-column nullability per query, which keeps variance at one kernel per expression.
+ */
+ private def nullable(v: ValueVector): Boolean = v.getNullCount != 0
+
+ /**
+ * Build the compile-time spec for one input Arrow vector. Recurses on complex types; scalars
+ * produce a [[ScalarColumnSpec]] carrying the concrete Arrow vector class and nullability.
+ * Spark `DataType`s on complex children come from [[Utils.fromArrowField]] so the Arrow ->
+ * Spark mapping stays in one place.
+ */
+ private def specFor(v: ValueVector): ArrowColumnSpec = v match {
+ case map: MapVector =>
+ // MapVector extends ListVector; match it first. Its data vector is a StructVector with
+ // child 0 = key and child 1 = value.
+ val struct = map.getDataVector.asInstanceOf[StructVector]
+ val keyVec = struct.getChildByOrdinal(0).asInstanceOf[ValueVector]
+ val valueVec = struct.getChildByOrdinal(1).asInstanceOf[ValueVector]
+ MapColumnSpec(
+ nullable = nullable(map),
+ keySparkType = Utils.fromArrowField(keyVec.getField),
+ valueSparkType = Utils.fromArrowField(valueVec.getField),
+ key = specFor(keyVec),
+ value = specFor(valueVec))
+ case list: ListVector =>
+ val child = list.getDataVector
+ ArrayColumnSpec(nullable(list), Utils.fromArrowField(child.getField), specFor(child))
+ case struct: StructVector =>
+ val fieldSpecs = (0 until struct.size()).map { fi =>
+ val childVec = struct.getChildByOrdinal(fi).asInstanceOf[ValueVector]
+ val field = struct.getField.getChildren.get(fi)
+ StructFieldSpec(
+ name = field.getName,
+ sparkType = Utils.fromArrowField(field),
+ nullable = field.isNullable,
+ child = specFor(childVec))
+ }
+ StructColumnSpec(nullable(struct), fieldSpecs)
+ case _: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector | _: BigIntVector |
+ _: Float4Vector | _: Float8Vector | _: DecimalVector | _: VarCharVector |
+ _: VarBinaryVector | _: DateDayVector | _: TimeStampMicroVector |
+ _: TimeStampMicroTZVector =>
+ ScalarColumnSpec(v.getClass.asInstanceOf[Class[_ <: ValueVector]], nullable(v))
+ case other =>
+ throw new UnsupportedOperationException(
+ s"CometCodegenDispatchUDF: unsupported Arrow vector ${other.getClass.getSimpleName}")
+ }
+
+ /**
+ * Estimate output byte capacity for variable-length output types. Sums the data-buffer sizes of
+ * variable-length input vectors as an upper bound for typical transform expressions (replace,
+ * upper, lower, substring, concat on the same inputs). Underestimates are still corrected by
+ * `setSafe`; this just reduces the odds of mid-loop reallocation.
+ */
+ private def estimatedOutputBytes(outputType: DataType, dataCols: Array[ValueVector]): Int = {
+ outputType match {
+ case _: StringType | _: BinaryType =>
+ var sum = 0
+ var i = 0
+ while (i < dataCols.length) {
+ dataCols(i) match {
+ case v: VarCharVector => sum += v.getDataBuffer.writerIndex().toInt
+ case v: VarBinaryVector => sum += v.getDataBuffer.writerIndex().toInt
+ case _ => // no size hint for fixed-width vector types
+ }
+ i += 1
+ }
+ sum
+ case _ => -1
+ }
+ }
+}
+
+object CometCodegenDispatchUDF {
+
+ private val CacheCapacity: Int = 128
+
+ /**
+ * 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. TODO(perf-cache-key): see
+ * `docs/source/contributor-guide/jvm_udf_dispatch.md#open-items` for possible optimizations if
+ * a workload makes this hot.
+ */
+ final case class CacheKey(bytesKey: ByteBuffer, specs: IndexedSeq[ArrowColumnSpec])
+
+ private case class CacheEntry(
+ compiled: CometBatchKernelCodegen.CompiledKernel,
+ outputType: DataType)
+
+ private val kernelCache: java.util.Map[CacheKey, CacheEntry] =
+ Collections.synchronizedMap(
+ new LinkedHashMap[CacheKey, CacheEntry](CacheCapacity, 0.75f, true) {
+ override def removeEldestEntry(
+ eldest: java.util.Map.Entry[CacheKey, CacheEntry]): Boolean =
+ size() > CacheCapacity
+ })
+
+ // Observability counters. Incremented under the `kernelCache.synchronized` block in
+ // `lookupOrCompile` so counter increments and cache mutations cannot interleave. Read via
+ // [[stats]]; reset via [[resetStats]] for tests.
+ private val compileCount = new AtomicLong(0)
+ private val cacheHitCount = new AtomicLong(0)
+
+ /**
+ * Snapshot of dispatcher cache counters and current size. Intended for tests, logging, and
+ * future integration with Spark SQL metrics. Not thread-synchronized across the three fields
+ * (each read is atomic, but they are not read atomically together); snapshots taken during
+ * concurrent activity may show a consistent individual-field view but a slightly inconsistent
+ * combined view. Fine for reporting, not for assertions that require cross-field invariants.
+ */
+ final case class DispatcherStats(compileCount: Long, cacheHitCount: Long, cacheSize: Int) {
+ def totalLookups: Long = compileCount + cacheHitCount
+ def hitRate: Double =
+ if (totalLookups == 0) 0.0 else cacheHitCount.toDouble / totalLookups.toDouble
+ }
+
+ /** Returns a snapshot of cache counters and current size. Cheap; safe to call anytime. */
+ def stats(): DispatcherStats =
+ DispatcherStats(compileCount.get(), cacheHitCount.get(), kernelCache.size())
+
+ /** Reset counters to zero. Leaves the compile cache intact. Intended for tests. */
+ def resetStats(): Unit = {
+ compileCount.set(0)
+ cacheHitCount.set(0)
+ }
+
+ /**
+ * Test-facing snapshot of compiled kernel signatures currently in the cache. Each entry is the
+ * pair `(input Arrow vector classes in ordinal order, output Spark DataType)` the kernel
+ * compiled against. Lets tests assert that the dispatcher actually specialized on the types it
+ * was expected to, not just that the query returned a correct result (which Spark would do
+ * regardless of how the kernel was shaped).
+ *
+ * Drops the `ArrowColumnSpec.nullable` bit to keep assertions robust to per-batch nullability
+ * variance: test data with no nulls compiles with `nullable=false` and the same expression run
+ * against data with nulls would cache a second variant. Tests assert on vector class and output
+ * type; both variants satisfy the same assertion.
+ */
+ def snapshotCompiledSignatures(): Set[(IndexedSeq[Class[_ <: ValueVector]], DataType)] = {
+ kernelCache.synchronized {
+ import scala.jdk.CollectionConverters._
+ kernelCache
+ .entrySet()
+ .asScala
+ .iterator
+ .map { e =>
+ (e.getKey.specs.map(_.vectorClass), e.getValue.outputType)
+ }
+ .toSet
+ }
+ }
+
+ private def lookupOrCompile(
+ key: CacheKey,
+ bytes: Array[Byte],
+ specs: IndexedSeq[ArrowColumnSpec]): CacheEntry = {
+ kernelCache.synchronized {
+ val existing = kernelCache.get(key)
+ if (existing != null) {
+ cacheHitCount.incrementAndGet()
+ existing
+ } else {
+ // Use a classloader that can see Spark classes. The Comet native runtime calls us on a
+ // Tokio worker thread where the context classloader may not be set to Spark's task
+ // loader, so fall back to the loader that loaded `Expression` itself if needed.
+ val loader = Option(Thread.currentThread().getContextClassLoader)
+ .getOrElse(classOf[Expression].getClassLoader)
+ val rawExpr = SparkEnv.get.closureSerializer
+ .newInstance()
+ .deserialize[Expression](ByteBuffer.wrap(bytes), loader)
+ // Tighten BoundReference.nullable based on the observed batch. The plan-time value is
+ // conservative (the column may be null somewhere in the query's execution), but for
+ // this specific batch we know. Rewriting lets Spark's `BoundReference.genCode` skip the
+ // `isNull` branch at source level rather than leaving it to JIT constant-folding.
+ // Correctness is preserved by the cache key: a later batch with nulls on this column has
+ // a different `specs`, so it hits a different kernel compiled with nullable=true.
+ val boundExpr = rewriteBoundReferences(rawExpr, specs)
+ val compiled = CometBatchKernelCodegen.compile(boundExpr, specs)
+ val entry = CacheEntry(compiled, boundExpr.dataType)
+ kernelCache.put(key, entry)
+ compileCount.incrementAndGet()
+ entry
+ }
+ }
+ }
+
+ /**
+ * Walk the bound expression tree and rewrite any `BoundReference(ord, dt, nullable=true)` to
+ * `nullable=false` when the corresponding input column in `specs` is non-nullable for this
+ * batch. Only tightens; never relaxes. Expressions outside the `BoundReference` leaves are
+ * unchanged.
+ */
+ private def rewriteBoundReferences(
+ expr: Expression,
+ specs: IndexedSeq[ArrowColumnSpec]): Expression = {
+ expr.transform {
+ case BoundReference(ord, dt, true)
+ if ord >= 0 && ord < specs.length && !specs(ord).nullable =>
+ BoundReference(ord, dt, nullable = false)
+ // Fall through unchanged: non-BoundReference nodes and BoundReferences that are already
+ // non-nullable or point at a nullable column in this batch.
+ case other => other
+ }
+ }
+
+ /**
+ * Partition index for the generated kernel's `init`. Expressions whose `doGenCode` calls
+ * `addPartitionInitializationStatement` (e.g. `Rand`, `Randn`, `Uuid`) reseed mutable state
+ * from this. Falls back to 0 when the dispatcher is exercised outside a Spark task (unit tests)
+ * so an absent `TaskContext` does not fail the call; the result is still deterministic for that
+ * fallback.
+ */
+ private def currentPartitionIndex(): Int =
+ Option(TaskContext.get()).map(_.partitionId()).getOrElse(0)
+}
diff --git a/common/src/main/scala/org/apache/comet/udf/CometInternalRow.scala b/common/src/main/scala/org/apache/comet/udf/CometInternalRow.scala
new file mode 100644
index 0000000000..0007499ea1
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/udf/CometInternalRow.scala
@@ -0,0 +1,96 @@
+/*
+ * 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.udf
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampNTZType, TimestampType}
+import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
+
+import org.apache.comet.shims.CometInternalRowShim
+
+/**
+ * Shim base for Comet-owned [[InternalRow]] getters used by the Arrow-direct codegen kernel.
+ *
+ * Provides `throw new UnsupportedOperationException` defaults for every abstract method declared
+ * by `InternalRow` and `SpecializedGetters`. Concrete subclasses (`CometBatchKernel` and its
+ * generated subclasses) override only the getters they actually support for their input shape.
+ *
+ * Purpose: keep subclasses free of boilerplate throws, and absorb forward-compat breakage if
+ * Spark adds abstract methods to `InternalRow` in a future version. Add the defaulted override
+ * here once, all subclasses recompile.
+ */
+abstract class CometInternalRow extends InternalRow with CometInternalRowShim {
+
+ override def numFields: Int = unsupported("numFields")
+ override def isNullAt(ordinal: Int): Boolean = unsupported("isNullAt")
+
+ override def getBoolean(ordinal: Int): Boolean = unsupported("getBoolean")
+ override def getByte(ordinal: Int): Byte = unsupported("getByte")
+ override def getShort(ordinal: Int): Short = unsupported("getShort")
+ override def getInt(ordinal: Int): Int = unsupported("getInt")
+ override def getLong(ordinal: Int): Long = unsupported("getLong")
+ override def getFloat(ordinal: Int): Float = unsupported("getFloat")
+ override def getDouble(ordinal: Int): Double = unsupported("getDouble")
+ override def getDecimal(ordinal: Int, precision: Int, scale: Int): Decimal =
+ unsupported("getDecimal")
+ override def getUTF8String(ordinal: Int): UTF8String = unsupported("getUTF8String")
+ override def getBinary(ordinal: Int): Array[Byte] = unsupported("getBinary")
+ override def getInterval(ordinal: Int): CalendarInterval = unsupported("getInterval")
+ override def getStruct(ordinal: Int, numFields: Int): InternalRow = unsupported("getStruct")
+ override def getArray(ordinal: Int): ArrayData = unsupported("getArray")
+ override def getMap(ordinal: Int): MapData = unsupported("getMap")
+
+ /**
+ * Generic `get(ordinal, dataType)` dispatcher. Required because `SpecializedGetters` declares
+ * it abstract and some Spark codegen paths (notably `SafeProjection` for deserializing
+ * `ScalaUDF` struct arguments) call it instead of the typed getter. Dispatches to the typed
+ * getter matching `dataType`; a null entry returns `null` outright. Unsupported types fall
+ * through to the shared throw.
+ */
+ override def get(ordinal: Int, dataType: DataType): AnyRef = {
+ if (isNullAt(ordinal)) return null
+ dataType match {
+ case BooleanType => java.lang.Boolean.valueOf(getBoolean(ordinal))
+ case ByteType => java.lang.Byte.valueOf(getByte(ordinal))
+ case ShortType => java.lang.Short.valueOf(getShort(ordinal))
+ case IntegerType | DateType => java.lang.Integer.valueOf(getInt(ordinal))
+ case LongType | TimestampType | TimestampNTZType =>
+ java.lang.Long.valueOf(getLong(ordinal))
+ case FloatType => java.lang.Float.valueOf(getFloat(ordinal))
+ case DoubleType => java.lang.Double.valueOf(getDouble(ordinal))
+ case _: StringType => getUTF8String(ordinal)
+ case BinaryType => getBinary(ordinal)
+ case dt: DecimalType => getDecimal(ordinal, dt.precision, dt.scale)
+ case st: StructType => getStruct(ordinal, st.size)
+ case _: ArrayType => getArray(ordinal)
+ case _: MapType => getMap(ordinal)
+ case other => unsupported(s"get for dataType $other")
+ }
+ }
+
+ override def setNullAt(i: Int): Unit = unsupported("setNullAt")
+ override def update(i: Int, value: Any): Unit = unsupported("update")
+ override def copy(): InternalRow = unsupported("copy")
+
+ protected def unsupported(method: String): Nothing =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: $method not implemented for this row shape")
+}
diff --git a/common/src/main/scala/org/apache/comet/udf/CometLambdaRegistry.scala b/common/src/main/scala/org/apache/comet/udf/CometLambdaRegistry.scala
deleted file mode 100644
index 5e020ae74a..0000000000
--- a/common/src/main/scala/org/apache/comet/udf/CometLambdaRegistry.scala
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * 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.udf
-
-import java.util.UUID
-import java.util.concurrent.ConcurrentHashMap
-
-import org.apache.spark.sql.catalyst.expressions.Expression
-
-/**
- * Thread-safe registry bridging plan-time Spark expressions to execution-time UDF lookup. At plan
- * time the serde layer registers a lambda expression under a unique key; at execution time the
- * UDF retrieves it by that key (passed as a scalar argument).
- */
-object CometLambdaRegistry {
-
- private val registry = new ConcurrentHashMap[String, Expression]()
-
- def register(expression: Expression): String = {
- val key = UUID.randomUUID().toString
- registry.put(key, expression)
- key
- }
-
- def get(key: String): Expression = {
- val expr = registry.get(key)
- if (expr == null) {
- throw new IllegalStateException(
- s"Lambda expression not found in registry for key: $key. " +
- "This indicates a lifecycle issue between plan creation and execution.")
- }
- expr
- }
-
- def remove(key: String): Unit = {
- registry.remove(key)
- }
-
- // Visible for testing
- def size(): Int = registry.size()
-}
diff --git a/common/src/main/scala/org/apache/comet/udf/CometMapData.scala b/common/src/main/scala/org/apache/comet/udf/CometMapData.scala
new file mode 100644
index 0000000000..fc99844110
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/udf/CometMapData.scala
@@ -0,0 +1,50 @@
+/*
+ * 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.udf
+
+import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
+
+/**
+ * Shim base for Comet-owned [[MapData]] views used by the Arrow-direct codegen kernel. Provides
+ * `UnsupportedOperationException` defaults for every abstract method on `MapData`; the codegen-
+ * emitted `InputMap_${path}` subclass overrides `numElements`, `keyArray`, and `valueArray`.
+ *
+ * Pairs with [[CometArrayData]] and [[CometInternalRow]]. `MapData` does not extend
+ * `SpecializedGetters` (unlike `ArrayData` / `InternalRow`), so no version-specific shim is
+ * needed here.
+ */
+abstract class CometMapData extends MapData {
+
+ override def numElements(): Int = unsupported("numElements")
+ override def keyArray(): ArrayData = unsupported("keyArray")
+ override def valueArray(): ArrayData = unsupported("valueArray")
+ override def copy(): MapData = unsupported("copy")
+
+ override def toString(): String = {
+ val n =
+ try numElements().toString
+ catch { case _: Throwable => "?" }
+ s"${getClass.getSimpleName}(numElements=$n)"
+ }
+
+ protected def unsupported(method: String): Nothing =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: $method not implemented for this map shape")
+}
diff --git a/common/src/main/scala/org/apache/comet/udf/CometUDF.scala b/common/src/main/scala/org/apache/comet/udf/CometUDF.scala
index 29186f0a2c..98cb519c1b 100644
--- a/common/src/main/scala/org/apache/comet/udf/CometUDF.scala
+++ b/common/src/main/scala/org/apache/comet/udf/CometUDF.scala
@@ -27,11 +27,16 @@ 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 the longest input.
+ * - The returned vector's length must match `numRows`.
*
- * Implementations must have a public no-arg constructor and must be stateless: a single instance
- * per class is cached and shared across native worker threads for the lifetime of the JVM.
+ * `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
+ * `numRows`; UDFs that may be called with zero data columns (e.g. a zero-arg ScalaUDF through the
+ * codegen dispatcher) need `numRows` to know how many rows to produce.
+ *
+ * Implementations must have a public no-arg constructor and should be stateless: instances are
+ * cached per executor thread for the lifetime of the JVM.
*/
trait CometUDF {
- def evaluate(inputs: Array[ValueVector]): ValueVector
+ def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector
}
diff --git a/common/src/main/scala/org/apache/spark/comet/CometTaskContextShim.scala b/common/src/main/scala/org/apache/spark/comet/CometTaskContextShim.scala
new file mode 100644
index 0000000000..9218fc5e78
--- /dev/null
+++ b/common/src/main/scala/org/apache/spark/comet/CometTaskContextShim.scala
@@ -0,0 +1,41 @@
+/*
+ * 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.comet
+
+import org.apache.spark.TaskContext
+
+/**
+ * Package-private access shim for `TaskContext.setTaskContext` / `TaskContext.unset`.
+ *
+ * Both methods are declared `protected[spark]` on Spark's `TaskContext` companion, so they are
+ * reachable from code inside the `org.apache.spark` package tree but not from `org.apache.comet`.
+ * The Comet JVM UDF bridge needs to set the thread-local `TaskContext` on its caller thread (a
+ * Tokio worker thread with no `TaskContext`) so the user's UDF body and any partition-sensitive
+ * built-ins (`Rand`, `Uuid`, `MonotonicallyIncreasingID`, etc.) see the driving Spark task's
+ * `TaskContext`. This shim lives in `org.apache.spark.comet` so it can call through to the
+ * protected methods, and exposes plain public forwarders the bridge (which lives in
+ * `org.apache.comet.udf`) can use.
+ */
+object CometTaskContextShim {
+
+ def set(taskContext: TaskContext): Unit = TaskContext.setTaskContext(taskContext)
+
+ def unset(): Unit = TaskContext.unset()
+}
diff --git a/common/src/main/spark-3.x/org/apache/comet/shims/CometExprTraitShim.scala b/common/src/main/spark-3.x/org/apache/comet/shims/CometExprTraitShim.scala
new file mode 100644
index 0000000000..3d039879d5
--- /dev/null
+++ b/common/src/main/spark-3.x/org/apache/comet/shims/CometExprTraitShim.scala
@@ -0,0 +1,42 @@
+/*
+ * 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.shims
+
+import org.apache.spark.sql.catalyst.expressions.{Expression, NullIntolerant}
+
+/**
+ * Per-profile view of expression traits that shifted shape across Spark versions. Spark 3.x has a
+ * `NullIntolerant` marker trait and no scalar-expression `Stateful` concept at all (the notion
+ * was added in 4.x as a boolean method on `Expression`). Routing checks through one shim lets the
+ * dispatcher ask "is this expression null-intolerant / stateful" without sprinkling version
+ * pattern matches through the codebase.
+ */
+trait CometExprTraitShim {
+ def isNullIntolerant(expr: Expression): Boolean = expr.isInstanceOf[NullIntolerant]
+
+ // No scalar `Stateful` trait in 3.x. Aggregate/window/generator stateful cases are rejected
+ // elsewhere in `canHandle`, so treating all scalar expressions as non-stateful here is
+ // conservative-correct on this profile.
+ def isStateful(expr: Expression): Boolean = false
+
+ // No collation / `ResolvedCollation` concept in 3.x, so no `Unevaluable` leaf slips past the
+ // dispatcher's guard here.
+ def isCodegenInertUnevaluable(expr: Expression): Boolean = false
+}
diff --git a/common/src/main/spark-3.x/org/apache/comet/shims/CometInternalRowShim.scala b/common/src/main/spark-3.x/org/apache/comet/shims/CometInternalRowShim.scala
new file mode 100644
index 0000000000..e71d301d48
--- /dev/null
+++ b/common/src/main/spark-3.x/org/apache/comet/shims/CometInternalRowShim.scala
@@ -0,0 +1,29 @@
+/*
+ * 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.shims
+
+/**
+ * Per-profile extension point mixed into `CometInternalRow` and `CometArrayData`. Spark 4.x added
+ * new abstract getters on `SpecializedGetters` (`getVariant` in 4.0, `getGeography` and
+ * `getGeometry` in 4.1) that both `InternalRow` and `ArrayData` concrete subclasses must
+ * implement. Spark 3.x has none of these; this trait is empty so the shared classes compile
+ * unchanged on that profile.
+ */
+trait CometInternalRowShim
diff --git a/common/src/main/spark-4.0/org/apache/comet/shims/CometInternalRowShim.scala b/common/src/main/spark-4.0/org/apache/comet/shims/CometInternalRowShim.scala
new file mode 100644
index 0000000000..20c6d47816
--- /dev/null
+++ b/common/src/main/spark-4.0/org/apache/comet/shims/CometInternalRowShim.scala
@@ -0,0 +1,36 @@
+/*
+ * 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.shims
+
+import org.apache.spark.unsafe.types.VariantVal
+
+/**
+ * Throwing-default implementations for `SpecializedGetters` methods that were added in Spark 4.0:
+ * `getVariant`. The Janino-generated kernel subclasses `CometInternalRow` (rows) and
+ * `CometArrayData` (array inputs), and each must satisfy every abstract method on the interface;
+ * without these defaults the compiled class fails its abstract-method check at class-load time.
+ * `GeographyVal` and `GeometryVal` were added in 4.1, so this profile's shim does not override
+ * those getters.
+ */
+trait CometInternalRowShim {
+ def getVariant(ordinal: Int): VariantVal =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: getVariant not supported")
+}
diff --git a/common/src/main/spark-4.1/org/apache/comet/shims/CometInternalRowShim.scala b/common/src/main/spark-4.1/org/apache/comet/shims/CometInternalRowShim.scala
new file mode 100644
index 0000000000..3d277e7505
--- /dev/null
+++ b/common/src/main/spark-4.1/org/apache/comet/shims/CometInternalRowShim.scala
@@ -0,0 +1,43 @@
+/*
+ * 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.shims
+
+import org.apache.spark.unsafe.types.{GeographyVal, GeometryVal, VariantVal}
+
+/**
+ * Throwing-default implementations for `SpecializedGetters` methods added in Spark 4.x:
+ * `getVariant` (4.0), `getGeography` and `getGeometry` (4.1). The Janino-generated kernel
+ * subclasses `CometInternalRow` (rows) and `CometArrayData` (array inputs), and each must satisfy
+ * every abstract method on the interface; without these defaults the compiled class fails its
+ * abstract-method check at class-load time.
+ */
+trait CometInternalRowShim {
+ def getVariant(ordinal: Int): VariantVal =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: getVariant not supported")
+
+ def getGeography(ordinal: Int): GeographyVal =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: getGeography not supported")
+
+ def getGeometry(ordinal: Int): GeometryVal =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: getGeometry not supported")
+}
diff --git a/common/src/main/spark-4.2/org/apache/comet/shims/CometInternalRowShim.scala b/common/src/main/spark-4.2/org/apache/comet/shims/CometInternalRowShim.scala
new file mode 100644
index 0000000000..3d277e7505
--- /dev/null
+++ b/common/src/main/spark-4.2/org/apache/comet/shims/CometInternalRowShim.scala
@@ -0,0 +1,43 @@
+/*
+ * 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.shims
+
+import org.apache.spark.unsafe.types.{GeographyVal, GeometryVal, VariantVal}
+
+/**
+ * Throwing-default implementations for `SpecializedGetters` methods added in Spark 4.x:
+ * `getVariant` (4.0), `getGeography` and `getGeometry` (4.1). The Janino-generated kernel
+ * subclasses `CometInternalRow` (rows) and `CometArrayData` (array inputs), and each must satisfy
+ * every abstract method on the interface; without these defaults the compiled class fails its
+ * abstract-method check at class-load time.
+ */
+trait CometInternalRowShim {
+ def getVariant(ordinal: Int): VariantVal =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: getVariant not supported")
+
+ def getGeography(ordinal: Int): GeographyVal =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: getGeography not supported")
+
+ def getGeometry(ordinal: Int): GeometryVal =
+ throw new UnsupportedOperationException(
+ s"${getClass.getSimpleName}: getGeometry not supported")
+}
diff --git a/common/src/main/spark-4.x/org/apache/comet/shims/CometExprTraitShim.scala b/common/src/main/spark-4.x/org/apache/comet/shims/CometExprTraitShim.scala
new file mode 100644
index 0000000000..2d86258014
--- /dev/null
+++ b/common/src/main/spark-4.x/org/apache/comet/shims/CometExprTraitShim.scala
@@ -0,0 +1,43 @@
+/*
+ * 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.shims
+
+import org.apache.spark.sql.catalyst.expressions.{Expression, ResolvedCollation}
+
+/**
+ * Spark 4.x replaced the `NullIntolerant` marker trait with a boolean method on `Expression`, and
+ * introduced a `stateful` boolean method covering scalar expressions that carry per-row state
+ * (e.g. `Rand`, `Uuid`). Neither concept exists as a trait in 4.x, so pattern matches against
+ * them would fail to compile. This shim routes the checks through the method form.
+ */
+trait CometExprTraitShim {
+ def isNullIntolerant(expr: Expression): Boolean = expr.nullIntolerant
+ def isStateful(expr: Expression): Boolean = expr.stateful
+
+ // `ResolvedCollation` is an `Unevaluable` leaf that only lives in `Collate.collation` as a
+ // type-level marker. `Collate.genCode` passes through to its child and never touches the
+ // collation slot, so the leaf is never invoked in generated code. Spark 4.1 analyzes it away,
+ // but 4.0 leaves it in the tree, so the dispatcher's `Unevaluable` guard trips on 4.0 without
+ // this exemption.
+ def isCodegenInertUnevaluable(expr: Expression): Boolean = expr match {
+ case _: ResolvedCollation => true
+ case _ => false
+ }
+}
diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md
index 77c73d68da..c2fbec4f54 100644
--- a/docs/source/contributor-guide/index.md
+++ b/docs/source/contributor-guide/index.md
@@ -47,6 +47,7 @@ Benchmarking Guide
Adding a New Operator
Adding a New Expression
Supported Spark Expressions
+JVM UDF Dispatch
Tracing
Profiling
Comet SQL Tests
diff --git a/docs/source/contributor-guide/jvm_udf_dispatch.md b/docs/source/contributor-guide/jvm_udf_dispatch.md
new file mode 100644
index 0000000000..f68753b9f4
--- /dev/null
+++ b/docs/source/contributor-guide/jvm_udf_dispatch.md
@@ -0,0 +1,336 @@
+
+
+# JVM UDF dispatch
+
+Comet offloads expressions that lack a native DataFusion implementation, or whose native implementation diverges from Spark's semantics, to JVM-side code that operates on Arrow batches passed through the C Data Interface. This preserves Spark compatibility on expressions that would otherwise force a whole-plan fallback to Spark. The tradeoff is a JNI roundtrip and per-batch JVM execution.
+
+The dispatch path is **Arrow-direct codegen via `CometCodegenDispatchUDF`** - one generic dispatcher that compiles a specialized kernel per bound Spark `Expression` plus input schema. Per-expression specialized emitters inside the dispatcher cover the cases where the default `doGenCode` output pays avoidable conversions; see [Specialized emitters](#specialized-emitters) below.
+
+The JNI bridge (`CometUdfBridge`) and proto schema (`JvmScalarUdf`) are generic enough to carry any `CometUDF` implementation, but the codebase today contains one: `CometCodegenDispatchUDF`.
+
+## Arrow-direct codegen via `CometCodegenDispatchUDF`
+
+One UDF class handles any scalar Spark `Expression` in the supported type surface. For each `(boundExpr, inputSchema)` pair, it compiles a specialized `CometBatchKernel` subclass via Janino that fuses Arrow input reads, expression evaluation, and Arrow output writes into one method. The kernel is cached in a JVM-wide LRU.
+
+### Transport
+
+At plan time the serde binds the expression tree to its leaf `AttributeReference`s, serializes the bound `Expression` via Spark's closure serializer, and emits a `JvmScalarUdf` proto whose argument 0 is a `Literal(bytes, BinaryType)` holding the serialized Expression. Arguments 1..N are the raw data columns the `BoundReference`s refer to, in ordinal order.
+
+At execute time, `CometCodegenDispatchUDF.evaluate` reads the bytes from the `VarBinaryVector` at arg 0, computes a cache key from (bytes, per-column Arrow vector class, per-column nullability), and either reuses a cached `CompiledKernel` or compiles one on the miss path.
+
+The self-describing proto removes the driver-side state the original prototype relied on. Cluster-mode executors deserialize and compile locally.
+
+**Classloader caveat.** The Comet native runtime calls the UDF on a Tokio worker thread whose context classloader may not be Spark's task loader. `SparkEnv.get.closureSerializer.newInstance().deserialize[Expression](bytes)` without an explicit loader fails with `ClassNotFoundException` on Spark's expression classes. The dispatcher passes an explicit loader, falling back to the loader that loaded `Expression` if the thread context is null.
+
+### Compilation
+
+`CometBatchKernelCodegen.compile(boundExpr, inputSchema)` generates a Java source for a `SpecificCometBatchKernel` that:
+
+- Extends `CometBatchKernel`, which extends `CometInternalRow`, which extends Spark's `InternalRow`. The kernel **is** the `InternalRow` that Spark's `BoundReference.genCode` reads from.
+- Sets `ctx.INPUT_ROW = "row"` at compile time and aliases `InternalRow row = this;` inside `process`, so Spark's generated body calls `row.getUTF8String(ord)` which resolves to the kernel's own typed getter. The getter is final, the ordinal is constant at the call site, and JIT devirtualizes and folds the switch. `row` rather than `this` because Spark's `splitExpressions` uses INPUT_ROW as a helper-method parameter name and `this` is a reserved Java keyword.
+- Carries typed input fields `col0 .. colN`, one per bound column, cast at the top of `process` from the generic `ValueVector[]` to the concrete Arrow class baked in at compile time.
+- Emits `isNullAt(ordinal)` and `getUTF8String(ordinal)` overrides whose switch cases are specialized per column. A column marked non-nullable compiles to `return false;`; a `VarCharVector` compiles to a zero-copy `UTF8String.fromAddress` read against the Arrow data buffer.
+- Overrides `init(int partitionIndex)` with the statements collected by `ctx.addPartitionInitializationStatement`. Non-deterministic expressions (`Rand`, `Randn`, `Uuid`) register statements that reseed mutable state from `partitionIndex`; deterministic expressions leave `init` empty.
+- Processes the batch in a tight loop that sets `this.rowIdx = i`, runs the expression body (either `boundExpr.genCode` for the default path or a specialized emitter), and writes to the typed output vector.
+
+### Specialized emitters
+
+For expressions whose `doGenCode` forces conversions that a tighter byte-oriented loop could skip, the dispatcher has per-expression overrides that emit custom Java while staying inside the framework (same cache, same bridge, same serde entry). Today that is `RegExpReplace`: the default path goes `Arrow bytes -> UTF8String -> String -> Matcher -> String -> UTF8String -> bytes -> Arrow` because `java.util.regex.Matcher` requires a `CharSequence`. The specialized emitter writes the byte-oriented shape directly (`Arrow bytes -> String -> Matcher -> String -> bytes -> Arrow`). The `UTF8String` round-trip costs measurable time on wide-match workloads; see `specializedRegExpReplaceBody` for the benchmark rationale.
+
+Precedent for adding new specializations: match when an expression's `doGenCode` pays conversions an Arrow-aware byte-oriented loop would avoid. Keep the specialization minimal (no speculative layering beyond the conversions it exists to skip) so its value over the default path stays legible.
+
+### Caching
+
+Three cache layers compose at three different scopes. None is redundant: collapsing any pair would either lose correctness or pay an avoidable cost.
+
+1. **JVM-wide compile cache.** Value is `CompiledKernel(factory: GeneratedClass, freshReferences: () => Array[Any])`, keyed by `(ByteBuffer.wrap(bytes), IndexedSeq[ArrowColumnSpec])`. Bounded LRU via `Collections.synchronizedMap(LinkedHashMap(accessOrder=true))` with `removeEldestEntry`, capacity 128. Same shape as `IcebergPlanDataInjector.commonCache` in `spark/src/main/scala/org/apache/spark/sql/comet/operators.scala`. Amortizes the Janino compile cost across every thread and every query in the JVM.
+
+2. **Per-thread UDF instance cache.** `CometUdfBridge.INSTANCES` is a `ThreadLocal>` that hands each task thread its own `CometCodegenDispatchUDF`. Keeps cache layer 3's instance fields safe without synchronization.
+
+3. **Per-partition kernel instance cache.** Plain mutable fields (`activeKernel`, `activeKey`, `activePartition`) on each UDF instance, managed by `ensureKernel`. The compiled `GeneratedClass` produces a kernel instance, and the kernel carries per-row mutable state (`Rand`'s `XORShiftRandom`, `MonotonicallyIncreasingID`'s counter, `addMutableState` fields) that must advance across batches within a partition and reset across partitions. `ensureKernel` allocates a fresh kernel and calls `init(partitionIndex)` only when the partition or cache key changes; otherwise the same kernel handles every batch in the partition.
+
+Matches Spark `WholeStageCodegenExec`: compile once per plan, instantiate per partition, init, iterate.
+
+#### Why `freshReferences` is a closure, not a cached array
+
+`CompiledKernel` holds a closure that regenerates `references: Array[Any]` each time a new kernel is allocated, rather than caching a single shared array. Reason: some expressions (notably `ScalaUDF`) embed stateful Spark `ExpressionEncoder` serializers into `references` via `ctx.addReferenceObj`. Those serializers reuse an internal `UnsafeRow` / `byte[]` buffer per `.apply(...)` call and are not thread-safe. If two kernels on different partitions shared one serializer instance, they would race on that buffer and return garbage.
+
+Re-running `genCode(ctx)` per kernel allocation costs microseconds; Janino compile costs milliseconds. Caching only the expensive piece preserves correctness cheaply. A future optimization would be to distinguish expressions whose references are all immutable (most non-UDF expressions) from those that embed stateful converters, and cache the array in the immutable case; not worth the complexity today.
+
+### Plan-time dispatchability
+
+`CometBatchKernelCodegen.canHandle(boundExpr)` runs at serde time. It returns `None` when the dispatcher can compile the expression, `Some(reason)` when it cannot. Checks:
+
+- Output `dataType` is in the scalar set `allocateOutput` and `emitOutputWriter` cover.
+- No `AggregateFunction` or `Generator` anywhere in the tree (scalar-only bridge).
+- Every `BoundReference`'s data type is in the input set `emitTypedGetters` has a getter for.
+
+The serde calls `withInfo(original, reason) + None` on a `Some` result, so Spark falls back rather than the kernel compiler crashing at execute time. Intermediate node types are not checked - `doGenCode` materializes them in local variables; only leaves (row reads) and the root (output write) touch Arrow.
+
+### Observability
+
+`CometCodegenDispatchUDF.stats()` returns `DispatcherStats(compileCount, cacheHitCount, cacheSize)`. `hitRate` is derived. `resetStats()` clears the counters (not the cache) for test isolation.
+
+Counters are not yet surfaced anywhere user-visible. Candidates for future wiring: Spark SQL metrics on the hosting operator, a JMX MBean, a Spark accumulator, or a periodic log line.
+
+## User-defined scalar functions (ScalaUDF)
+
+The codegen dispatcher routes scalar `org.apache.spark.sql.catalyst.expressions.ScalaUDF` expressions through the same compile + per-partition-kernel pipeline as the regex serdes. The serde is `CometScalaUDF` in `spark/src/main/scala/org/apache/comet/serde/scalaUdf.scala`, registered in `QueryPlanSerde.miscExpressions`.
+
+Why it works without per-UDF handling: Spark's `ScalaUDF.doGenCode` already emits compilable Java that calls the user function via `ctx.addReferenceObj`. The compile path runs `boundExpr.genCode(ctx)` and picks this up unchanged. The serialized-bytes transport carries the function reference through Spark's closure serializer, the same machinery Spark uses to ship UDFs to executors. Per-partition kernel caching handles `ScalaUDF`'s `stateful=true`.
+
+Without this serde, any `ScalaUDF` in a plan forces Comet to fall back to Spark for the whole plan, losing acceleration on the surrounding operators. With it, scalar UDFs whose types fit the supported surface stay on the Comet path behind one JNI hop.
+
+### What's covered
+
+| What users write | Spark expression class | Route through codegen |
+| --------------------------------------------------------------- | ------------------------------------------------------ | ------------------------------------------------------------- |
+| `udf((x: T) => ...)` or `spark.udf.register` (Scala) | `ScalaUDF` | yes |
+| `spark.udf.register("f", new UDF1[...]{...})` (Java) | `ScalaUDF` (Spark wraps the Java functional interface) | yes, transparently |
+| `CREATE FUNCTION foo AS 'com.example.MyUDF'` (SQL registration) | `ScalaUDF` | yes, if the user class is reachable on the executor classpath |
+
+### What's not covered
+
+| What users write | Spark expression class | Why not |
+| ------------------------------- | --------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ |
+| Aggregate UDF | `ScalaAggregator`, `TypedImperativeAggregate`, old `UserDefinedAggregateFunction` | accumulator-based; needs a different bridge contract (accumulate + merge + finalize) |
+| Table UDF / generator | `UserDefinedTableFunction` | 1 row -> N rows; `canHandle` rejects `Generator` |
+| Python `@udf` | `PythonUDF` | subprocess runtime, not JVM |
+| Pandas `@pandas_udf` | `PandasUDF` | Arrow-via-subprocess runtime |
+| Hive `GenericUDF` / `SimpleUDF` | `HiveGenericUDF` / `HiveSimpleUDF` | separate expression classes; would need their own serde |
+
+### Constraints within the ScalaUDF path
+
+- Input and output types must be in the supported surface (see [Type surface](#type-surface)). Nested types (`Struct`, `Array`, `Map`) are supported when their element types are supported.
+- The user function must be closure-serializable. This is Spark's own requirement; a function that works with Spark's executor execution works here.
+- User functions that touch `TaskContext` internals, accumulators, or broadcast variables in unusual ways may misbehave; the common case works.
+- Stateful behavior: per-partition kernel caching resets kernel instance state on partition boundary, matching the contract most user UDFs assume (and matching Spark's own re-instantiation on some paths). UDFs that rely on long-lived JVM-wide state across partitions in the same executor see that state reset more often than before, which is rare and usually a latent bug in the UDF.
+
+### Mode knob interaction
+
+`spark.comet.exec.codegenDispatch.mode` controls routing:
+
+- `auto` (default) and `force`: ScalaUDFs go through the codegen dispatcher.
+- `disabled`: `CometScalaUDF.convert` returns `None` and the plan falls back to Spark.
+
+There is no non-codegen Comet path for arbitrary user functions.
+
+## Type surface
+
+### Input (kernel getters)
+
+All scalar Spark types that map to a single Arrow vector:
+
+| Spark type | Arrow vector class | `InternalRow` getter |
+| ----------------------------------------- | ---------------------------------------------------------- | -------------------------------------------------------- |
+| BooleanType | BitVector | `getBoolean` |
+| ByteType | TinyIntVector | `getByte` |
+| ShortType | SmallIntVector | `getShort` |
+| IntegerType, DateType | IntVector, DateDayVector | `getInt` |
+| LongType, TimestampType, TimestampNTZType | BigIntVector, TimeStampMicroVector, TimeStampMicroTZVector | `getLong` |
+| FloatType | Float4Vector | `getFloat` |
+| DoubleType | Float8Vector | `getDouble` |
+| DecimalType | DecimalVector | `getDecimal(ord, precision, scale)` |
+| StringType | VarCharVector | `getUTF8String` (zero-copy via `UTF8String.fromAddress`) |
+| BinaryType | VarBinaryVector | `getBinary` (allocates `byte[]`) |
+
+Widening: add cases to `CometBatchKernelCodegen.emitTypedGetters` and accept the new vector classes in `CometCodegenDispatchUDF.evaluate`'s input pattern match.
+
+### Output (writers + allocators)
+
+All scalar Spark types that map to a single Arrow vector: `Boolean`, `Byte`, `Short`, `Int`, `Long`, `Float`, `Double`, `Decimal`, `String`, `Binary`, `Date`, `Timestamp`, `TimestampNTZ`. Mirrors `ArrowWriters.createFieldWriter` so producer and consumer sides stay aligned. Widen by adding cases to `CometBatchKernelCodegen.allocateOutput` and `emitOutputWriter`.
+
+### Complex types
+
+`ArrayType`, `StructType`, and `MapType` are supported as both input and output, including arbitrary nesting (`Array>`, `Array`, `Struct