feat(experimental): ScalaUDF and Java UDF support via Janino codegen

.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

.github/workflows/pr_build_linux.yml

@@ -307,6 +307,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: |
@@ -385,6 +386,9 @@ 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
+              org.apache.comet.CometCodegenHOFSuite
           - name: "sql"
             value: |
               org.apache.spark.sql.CometToPrettyStringSuite

.github/workflows/pr_build_macos.yml

@@ -155,6 +155,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: |
@@ -232,6 +233,9 @@ 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
+              org.apache.comet.CometCodegenHOFSuite
           - name: "sql"
             value: |
               org.apache.spark.sql.CometToPrettyStringSuite

common/src/main/java/org/apache/comet/codegen/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.codegen;
+
+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.
+ *
+ * <p>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.
+   *
+   * <p>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) {}
+}

common/src/main/scala/org/apache/comet/CometConf.scala

@@ -380,6 +380,17 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_SCALA_UDF_CODEGEN_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.exec.scalaUDF.codegen.enabled")
+      .category(CATEGORY_EXEC)
+      .doc(
+        "Whether to route Spark `ScalaUDF` expressions through Comet's Arrow-direct codegen " +
+          "dispatcher. When enabled, a supported ScalaUDF is compiled into a per-batch kernel " +
+          "that reads and writes Arrow vectors directly from native execution. When disabled, " +
+          "plans containing a ScalaUDF fall back to Spark for the enclosing operator.")
+      .booleanConf
+      .createWithDefault(true)
+
   val COMET_EXEC_SHUFFLE_WITH_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.native.shuffle.partitioning.hash.enabled")
       .category(CATEGORY_SHUFFLE)

common/src/main/scala/org/apache/comet/codegen/CometArrayData.scala

@@ -0,0 +1,106 @@
+/*
+ * 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.codegen
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
+
+import org.apache.comet.shims.CometInternalRowShim
+
+/**
+ * Throwing-default base for [[ArrayData]] in the Arrow-direct codegen kernel. Subclasses override
+ * only the getters their element type needs (e.g. `numElements`, `isNullAt`, `getUTF8String` for
+ * an `ArrayType(StringType)` input).
+ *
+ * Consumer: `InputArray_${path}` nested classes the input emitter generates per `ArrayType` input
+ * column. They back `getArray(ord)` plus the recursion for `Array<Array<...>>` and array-typed
+ * map keys / struct fields.
+ *
+ * `ArrayData` and [[CometInternalRow]]'s [[InternalRow]] are sibling abstract classes in Spark
+ * (both extend `SpecializedGetters`, neither inherits the other), so a base aimed at one cannot
+ * serve the other. The dispatch body that '''is''' shared between them lives in
+ * [[CometSpecializedGettersDispatch]]. The third sibling, [[CometMapData]], backs `InputMap_*`
+ * and routes `keyArray()` / `valueArray()` through `CometArrayData` instances.
+ *
+ * Mixes in [[CometInternalRowShim]] for the same reason `CometInternalRow` does: Spark 4.x adds
+ * abstract `SpecializedGetters` methods (`getVariant`, `getGeography`, `getGeometry`) that both
+ * `InternalRow` and `ArrayData` inherit; the per-profile shim provides throwing defaults.
+ */
+abstract class CometArrayData extends ArrayData with CometInternalRowShim {
+
+  override def getInterval(ordinal: Int): CalendarInterval = unsupported("getInterval")
+
+  override def get(ordinal: Int, dataType: DataType): AnyRef =
+    CometSpecializedGettersDispatch.get(this, ordinal, dataType)
+
+  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 getStruct(ordinal: Int, numFields: Int): InternalRow = unsupported("getStruct")
+
+  override def getArray(ordinal: Int): ArrayData = unsupported("getArray")
+
+  override def getMap(ordinal: Int): MapData = unsupported("getMap")
+
+  override def setNullAt(i: Int): Unit = unsupported("setNullAt")
+
+  protected def unsupported(method: String): Nothing =
+    throw new UnsupportedOperationException(
+      s"${getClass.getSimpleName}: $method not implemented for this array shape")
+
+  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)"
+  }
+
+  override def numElements(): Int = unsupported("numElements")
+}