cleanup round 2

Author: mbutrovich

spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala

@@ -0,0 +1,82 @@
+/*
+ * 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
+
+import org.apache.arrow.vector.ValueVector
+import org.apache.spark.sql.types.DataType
+
+import org.apache.comet.udf.codegen.CometScalaUDFCodegen
+
+/**
+ * Shared assertions for the codegen-dispatcher test suites. Mix in alongside `CometTestBase`.
+ */
+trait CometCodegenAssertions {
+
+  /** Asserts the dispatcher actually ran during `f`, guarding against silent serde fallback. */
+  protected def assertCodegenRan(f: => Unit): Unit = {
+    CometScalaUDFCodegen.resetStats()
+    f
+    val after = CometScalaUDFCodegen.stats()
+    assert(
+      after.compileCount + after.cacheHitCount >= 1,
+      s"expected codegen dispatcher activity, got $after")
+  }
+
+  /**
+   * Asserts the composed subtree fused into one kernel signature, not N (one per sub-expression).
+   * Uses the JVM-wide signature set rather than `compileCount` because per-task `boundExpr`
+   * isolation makes multi-partition queries trip `compileCount > 1` even when the bytecode is
+   * shared.
+   */
+  protected def assertOneKernelForSubtree(f: => Unit): Unit = {
+    CometScalaUDFCodegen.resetStats()
+    val sigsBefore = CometScalaUDFCodegen.snapshotCompiledSignatures()
+    f
+    val sigsAfter = CometScalaUDFCodegen.snapshotCompiledSignatures()
+    val grew = sigsAfter.size - sigsBefore.size
+    assert(
+      grew <= 1,
+      s"expected <= 1 new compiled-kernel signature for the composed subtree, grew by $grew; " +
+        s"new=${sigsAfter -- sigsBefore}")
+    val after = CometScalaUDFCodegen.stats()
+    assert(
+      after.compileCount + after.cacheHitCount >= 1,
+      s"expected codegen dispatcher activity, got $after")
+  }
+
+  /**
+   * Asserts a kernel matching the given input Arrow vector classes and output type sits in the
+   * JVM-wide signature set. Pair with `assertCodegenRan` since the set is append-only. Compares
+   * by simple name because `common` shades `org.apache.arrow`.
+   */
+  protected def assertKernelSignaturePresent(
+      inputs: Seq[Class[_ <: ValueVector]],
+      output: DataType): Unit = {
+    val sigs = CometScalaUDFCodegen.snapshotCompiledSignatures()
+    val expectedNames = inputs.map(_.getSimpleName).toIndexedSeq
+    val present = sigs.exists { case (cached, dt) =>
+      dt == output && cached.map(_.getSimpleName) == expectedNames
+    }
+    assert(
+      present,
+      s"expected kernel signature $expectedNames -> $output; " +
+        s"cache had ${sigs.map { case (c, d) => (c.map(_.getSimpleName), d) }}")
+  }
+}

spark/src/test/scala/org/apache/comet/CometCodegenFuzzSuite.scala

@@ -33,7 +33,6 @@ import org.apache.spark.sql.types._
 
 import org.apache.comet.DataTypeSupport.isComplexType
 import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions}
-import org.apache.comet.udf.codegen.CometScalaUDFCodegen
 
 /**
  * Randomized end-to-end tests for the Arrow-direct codegen dispatcher: schema-driven coverage of
@@ -42,7 +41,10 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen
  * (not [[CometFuzzTestBase]]) because the base's `shuffle` x `nativeC2R` cross-product is
  * irrelevant for projection-only queries.
  */
-class CometCodegenFuzzSuite extends CometTestBase with AdaptiveSparkPlanHelper {
+class CometCodegenFuzzSuite
+    extends CometTestBase
+    with AdaptiveSparkPlanHelper
+    with CometCodegenAssertions {
 
   /** Random schema with primitives plus shallow arrays and structs. No maps, no deep nesting. */
   private var mixedTypesFilename: String = _
@@ -120,20 +122,6 @@ class CometCodegenFuzzSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     super.sparkConf
       .set(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key, "true")
 
-  /**
-   * Resets dispatcher stats, runs `f`, then asserts the codegen path actually ran for at least
-   * one batch. Without this, a silent serde fallback would let the fuzz pass trivially because
-   * both Spark and whatever-Comet-ran-instead agree with Spark.
-   */
-  private def assertCodegenRan(f: => Unit): Unit = {
-    CometScalaUDFCodegen.resetStats()
-    f
-    val after = CometScalaUDFCodegen.stats()
-    assert(
-      after.compileCount + after.cacheHitCount >= 1,
-      s"expected at least one codegen dispatcher invocation during this query, got $after")
-  }
-
   /**
    * Identity ScalaUDF for one of the 14 primitive types in
    * [[org.apache.comet.testing.SchemaGenOptions.defaultPrimitiveTypes]]. Returns the registered

spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala

@@ -23,28 +23,29 @@ import org.apache.spark.SparkConf
 import org.apache.spark.sql.CometTestBase
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 
-import org.apache.comet.udf.codegen.CometScalaUDFCodegen
-
 /**
  * Higher-order function regression coverage for the codegen dispatcher.
  *
  * Spark's HOFs (`ArrayTransform`, `ArrayFilter`, `ArrayAggregate`, `ArrayExists`, `ZipWith`,
  * `MapFilter`, etc.) all extend `CodegenFallback`. The dispatcher's `canHandle` admits them.
  * `CodegenFallback.doGenCode` emits a single `((Expression) references[N]).eval(row)` call site
- * per HOF; at runtime the kernel dispatches to `Expression.eval(InternalRow)`, which iterates the
- * array, mutates `NamedLambdaVariable.value`'s `AtomicReference` per element, and recursively
- * evaluates the lambda body. Lambda-body leaf reads resolve through the kernel's own typed Arrow
- * getters since the kernel '''is''' an `InternalRow`.
+ * per HOF; the kernel dispatches to `Expression.eval(InternalRow)`, which iterates the array,
+ * mutates `NamedLambdaVariable.value`'s `AtomicReference` per element, and recursively evaluates
+ * the lambda body. Lambda-body leaf reads resolve through the kernel's typed Arrow getters since
+ * the kernel '''is''' an `InternalRow`.
  *
  * Cost model: per-row interpreted-eval inside the HOF subtree. Surrounding native operators stay
  * native; surrounding non-HOF expressions stay codegen.
  *
  * Critical invariant: each Spark task gets its own `boundExpr` Java object. The dispatcher's
  * compile cache lives on the per-task instance, not the companion, so concurrent partitions
- * cannot race on a shared `NamedLambdaVariable.value`. Mirrors Spark's per-task closure-
- * deserialize model. The two-collects test below regresses this.
+ * cannot race on a shared `NamedLambdaVariable.value`. The two-collects test below regresses
+ * this.
  */
-class CometCodegenHOFSuite extends CometTestBase with AdaptiveSparkPlanHelper {
+class CometCodegenHOFSuite
+    extends CometTestBase
+    with AdaptiveSparkPlanHelper
+    with CometCodegenAssertions {
 
   override protected def sparkConf: SparkConf =
     super.sparkConf
@@ -58,15 +59,6 @@ class CometCodegenHOFSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
-  private def assertCodegenRan(f: => Unit): Unit = {
-    CometScalaUDFCodegen.resetStats()
-    f
-    val after = CometScalaUDFCodegen.stats()
-    assert(
-      after.compileCount + after.cacheHitCount >= 1,
-      s"expected dispatcher activity, got $after")
-  }
-
   test("ArrayTransform inside identity ScalaUDF over Array<Int>") {
     // Regresses the simplest HOF shape: `idArr(transform(a, x -> x + 1))`. Tree contains one
     // CodegenFallback HOF; the kernel splices its interpreted-eval call site into the per-row

spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala

@@ -569,13 +569,6 @@ class CometCodegenSourceSuite extends AnyFunSuite {
       s"expected BigDecimal slow path for p>18 element; got:\n$src")
   }
 
-  // ============================================================================================
-  // Nested-type tests. Each case verifies that a complex-within-complex shape emits a full
-  // nested-class tree (outer + inner), wired together through the path-suffix naming
-  // convention: `_e` for array element, `_f${fi}` for struct field fi. Scalar-element / scalar-
-  // field leaves reuse the typed-getter templates already covered by the single-depth tests.
-  // ============================================================================================
-
   private def generate(expr: Expression, specs: IndexedSeq[ArrowColumnSpec]): String =
     CometBatchKernelCodegen.generateSource(expr, specs).body
 
@@ -751,18 +744,16 @@ class CometCodegenSourceSuite extends AnyFunSuite {
     }
   }
 
-  // ============================================================================================
-  // Null-guard emission for nested reference-typed getters. Spark's
-  // `CodeGenerator.setArrayElement` only emits an `isNullAt` check before `update(i, getX(j))`
-  // for primitive elements. For reference types (Decimal / String / Binary / Struct / Array /
-  // Map) it relies on the source's `getX` to return null on null positions itself, matching
-  // `ColumnarArray.getBinary` and friends. The emitter prepends `if (isNullAt(...)) return null;`
-  // to those getters when the element/field is nullable.
-  //
-  // Runtime regressions for the leaf reference types live in `CometCodegenSuite`; complex-type
-  // (Struct/Array/Map) coverage runs through HOFs in `CometCodegenHOFSuite`.
-  // ============================================================================================
-
+  /**
+   * Null-guard emission for nested reference-typed getters. Spark's
+   * `CodeGenerator.setArrayElement` only emits an `isNullAt` check before `update(i, getX(j))`
+   * for primitive elements. For reference types it relies on the source's `getX` to return null
+   * on null positions itself, matching `ColumnarArray.getBinary`. The emitter prepends `if
+   * (isNullAt(...)) return null;` when the element / field is nullable.
+   *
+   * Runtime regressions for the leaf reference types live in `CometCodegenSuite`; complex-type
+   * (Struct/Array/Map) coverage runs through HOFs in `CometCodegenHOFSuite`.
+   */
   private val nullableIntStruct = StructColumnSpec(
     nullable = true,
     fields = Seq(