test: add helpers to assert whether an expression ran natively or via codegen dispatch

.github/workflows/pr_build_linux.yml

@@ -393,6 +393,7 @@ jobs:
             value: |
               org.apache.comet.CometExpressionSuite
               org.apache.comet.CometSqlFileTestSuite
+              org.apache.comet.SqlFileTestParserSuite
               org.apache.comet.CometExpressionCoverageSuite
               org.apache.comet.CometHashExpressionSuite
               org.apache.comet.CometTemporalExpressionSuite

.github/workflows/pr_build_macos.yml

@@ -209,6 +209,7 @@ jobs:
             value: |
               org.apache.comet.CometExpressionSuite
               org.apache.comet.CometSqlFileTestSuite
+              org.apache.comet.SqlFileTestParserSuite
               org.apache.comet.CometExpressionCoverageSuite
               org.apache.comet.CometHashExpressionSuite
               org.apache.comet.CometTemporalExpressionSuite

docs/source/contributor-guide/adding_a_new_expression.md

@@ -409,6 +409,7 @@ For full documentation on the test file format, including directives like `Confi
 - **Cover both column references and literals.** Comet often uses different code paths for each. The Comet SQL Tests suite automatically disables constant folding, so all-literal queries are evaluated natively.
 - **Include edge cases** such as `NULL`, empty strings, boundary values, `NaN`, and multibyte UTF-8 characters.
 - **Keep one file per expression** to make failures easy to locate.
+- **Pin the mechanism when it depends on the input.** If your serde routes some input types natively and others through the JVM codegen dispatcher, say so with `expect_native(...)` and `expect_dispatch(...)`. A plain `query` cannot tell the two apart, so without these a later change that swaps one for the other passes silently. See [Comet SQL Tests](sql-file-tests.md).
 
 ##### Comet Scala Tests (alternative)
 
@@ -436,6 +437,34 @@ test("unhex") {
 }
 ```
 
+`checkSparkAnswerAndOperator` verifies that results match Spark and that Comet did not fall back
+for the operator. It does not distinguish _how_ Comet evaluated the expression. Comet has two
+accelerated paths, a native DataFusion expression and the JVM codegen dispatcher (Spark's own
+`doGenCode` compiled into an Arrow batch kernel), and both produce Spark-matching results by
+construction.
+
+If your serde picks between them, for example returning `Unsupported` for one input type and
+relying on `CodegenDispatchFallback` to keep it in the pipeline, use `checkSparkAnswerAndImpl` so
+the choice is asserted rather than assumed:
+
+```scala
+test("bit_length dispatches on binary and stays native on string") {
+  withTable("t") {
+    sql("CREATE TABLE t (b BINARY, s STRING) USING parquet")
+    sql("INSERT INTO t VALUES (X'48656c6c6f', 'hello'), (NULL, NULL)")
+
+    checkSparkAnswerAndImpl(sql("SELECT bit_length(b) FROM t"), dispatched = Seq("bit_length"))
+    checkSparkAnswerAndImpl(sql("SELECT bit_length(s) FROM t"), native = Seq("bit_length"))
+  }
+}
+```
+
+Names are the expression's `prettyName` lowercased. Naming an expression asserts both that it ran
+through the expected mechanism and that it did not run through the other one, so a serde that
+later gains a native path (or loses its dispatcher route) fails here instead of passing unnoticed.
+Prefer the SQL file equivalents, `expect_native(...)` and `expect_dispatch(...)`, when the test
+fits in a fixture.
+
 When writing Comet Scala Tests with literal values (e.g., `SELECT my_func('literal')`), Spark's constant folding optimizer may evaluate the expression at planning time, bypassing Comet. To prevent this, disable constant folding:
 
 ```scala

docs/source/contributor-guide/sql-file-tests.md

@@ -192,6 +192,35 @@ query expect_fallback(unsupported expression)
 SELECT unsupported_func(v) FROM test_table
 ```
 
+#### `query expect_dispatch(<names>)` / `query expect_native(<names>)`
+
+Checks results and coverage like a plain `query`, and additionally asserts how Comet evaluated
+the named expressions.
+
+Comet runs an expression either natively (a DataFusion expression) or through the JVM codegen
+dispatcher (Spark's own `doGenCode` compiled into an Arrow batch kernel). Both produce
+Spark-matching results, so a plain `query` cannot tell them apart. Use these modes on fixtures
+where the mechanism is the point of the test, typically an expression whose support depends on
+its argument type.
+
+```sql
+-- BinaryType has no native path and must route through the dispatcher
+query expect_dispatch(bit_length)
+SELECT bit_length(b) FROM test_bit_length_binary
+
+-- StringType must stay on the native path
+query expect_native(bit_length)
+SELECT bit_length(s) FROM test_bit_length
+```
+
+Names are comma-separated. A name is the expression's `prettyName` lowercased (`bit_length`,
+`octet_length`, `rlike`), which is not always the SQL alias used to invoke it. Naming an
+expression asserts both that it ran through the expected mechanism and that it did not run
+through the other one.
+
+A query carries one mode, so a query mixing a native and a dispatched expression has to be split
+into two queries, one per mode.
+
 #### `query ignore(<reason>)`
 
 Skips the query entirely. Use this for queries that hit known bugs. The reason should be a
@@ -243,6 +272,11 @@ SELECT array(1, 2, 3)[10]
    when you expect Comet to run the expression natively. Use `query spark_answer_only` when
    native execution is not yet expected.
 
+   If the expression's serde routes some input types to a native DataFusion expression and
+   others through the JVM codegen dispatcher, use `expect_native(...)` and `expect_dispatch(...)`
+   for those queries. A plain `query` cannot tell the two mechanisms apart, so the split is
+   otherwise untested.
+
 6. Run the tests to verify:
 
    ```shell
@@ -251,6 +285,17 @@ SELECT array(1, 2, 3)[10]
 
 ### Tips for writing thorough tests
 
+#### Pin the mechanism where the serde chooses one
+
+Reach for `expect_native(...)` / `expect_dispatch(...)` whenever the fixture's own comments
+explain which path an input takes. That comment is a claim about behavior, and these modes are
+what turn it into a test. Expressions worth annotating are the ones whose support level depends
+on argument type or on a config: `round` (float and double dispatch, decimal and integral stay
+native), `lower` / `upper` (dispatch by default), and anything mixing in `CodegenDispatchFallback`.
+
+A query carries a single mode, so a query that mixes both mechanisms has to be split. That split
+is usually worth doing on its own: it forces you to say which argument takes which path.
+
 #### Cover all combinations of literal and column arguments
 
 Comet often uses different code paths for literal values versus column references. Tests

spark/src/test/resources/sql-tests/expressions/math/round.sql

@@ -35,21 +35,21 @@ INSERT INTO test_round VALUES
  (cast('Infinity' as double), cast('Infinity' as float), 0.0, 0, 0),
  (cast('-Infinity' as double), cast('-Infinity' as float), 0.0, 0, 0)
 
-query
+query expect_dispatch(round)
 SELECT d, round(d), round(d, 0), round(d, 2), round(d, -1) FROM test_round
 
-query
+query expect_dispatch(round)
 SELECT f, round(f), round(f, 0), round(f, 2), round(f, -1) FROM test_round
 
 -- Null scale makes the whole result null without evaluating the child.
 query
 SELECT round(d, NULL), round(f, NULL) FROM test_round
 
 -- Decimal and integral inputs stay on the native path.
-query
+query expect_native(round)
 SELECT dec, round(dec), round(dec, 2), round(dec, -1) FROM test_round
 
-query
+query expect_native(round)
 SELECT i, round(i, 0), round(i, -1), round(l, -1) FROM test_round
 
 -- Doubles whose shortest decimal representation rounds differently than the exact binary value.
@@ -63,12 +63,18 @@ CREATE TABLE test_round_repr(d double) USING parquet
 statement
 INSERT INTO test_round_repr VALUES (-5.81855622136895E8), (6.1317116247283497E18)
 
-query
+query expect_dispatch(round)
 SELECT d, round(d, 5), round(d, -5) FROM test_round_repr
 
--- literal + literal
-query
-SELECT round(123.456, 2), round(2.5, 0), round(3.5, 0), round(-2.5, 0), round(NULL, 0)
+-- literal + literal. Unsuffixed decimal literals keep the native path; the D / F suffixed ones
+-- are float and double, so they dispatch.
+query expect_native(round)
+SELECT round(123.456, 2), round(2.5, 0), round(3.5, 0), round(-2.5, 0)
 
-query
+-- `round(NULL, 0)` is not a decimal: the untyped null is implicitly cast to double, so this one
+-- dispatches. Kept in its own query because a single query carries a single mode.
+query expect_dispatch(round)
+SELECT round(NULL, 0)
+
+query expect_dispatch(round)
 SELECT round(2.5D, 0), round(3.5D, 0), round(-2.5D, 0), round(2.5F, 0), round(-2.5F, 0)

spark/src/test/resources/sql-tests/expressions/string/lower.sql

@@ -25,11 +25,11 @@ CREATE TABLE test_lower(s string) USING parquet
 statement
 INSERT INTO test_lower VALUES ('HELLO'), ('hello'), ('Hello World'), (''), (NULL), ('123ABC')
 
-query
+query expect_dispatch(lower)
 SELECT lower(s) FROM test_lower
 
 -- literal arguments
-query
+query expect_dispatch(lower)
 SELECT lower('HELLO'), lower(''), lower(NULL)
 
 -- locale-sensitive characters: Greek sigma and Turkish dotted I
@@ -39,5 +39,5 @@ CREATE TABLE test_lower_unicode(s string) USING parquet
 statement
 INSERT INTO test_lower_unicode VALUES ('ΣIGMA'), ('İSTANBUL'), ('GROSSE'), ('CAFÉ')
 
-query
+query expect_dispatch(lower)
 SELECT lower(s) FROM test_lower_unicode

spark/src/test/resources/sql-tests/expressions/string/upper.sql

@@ -25,11 +25,11 @@ CREATE TABLE test_upper(s string) USING parquet
 statement
 INSERT INTO test_upper VALUES ('hello'), ('HELLO'), ('Hello World'), (''), (NULL), ('123abc')
 
-query
+query expect_dispatch(upper)
 SELECT upper(s) FROM test_upper
 
 -- literal arguments
-query
+query expect_dispatch(upper)
 SELECT upper('hello'), upper(''), upper(NULL)
 
 -- locale-sensitive characters: German sharp s and Turkish dotted/dotless I
@@ -39,5 +39,5 @@ CREATE TABLE test_upper_unicode(s string) USING parquet
 statement
 INSERT INTO test_upper_unicode VALUES ('straße'), ('istanbul'), ('İstanbul'), ('finish')
 
-query
+query expect_dispatch(upper)
 SELECT upper(s) FROM test_upper_unicode

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

@@ -21,6 +21,8 @@ package org.apache.comet
 
 import scala.util.Random
 
+import org.scalatest.exceptions.TestFailedException
+
 import org.apache.arrow.vector._
 import org.apache.spark.{SparkConf, SparkEnv, TaskContext}
 import org.apache.spark.sql.CometTestBase
@@ -210,6 +212,35 @@ class CometCodegenSuite
     }
   }
 
+  test("checkSparkAnswerAndImpl pins the mechanism and fails when the claim is wrong") {
+    // The assertion helper is only worth having if it fails. `abs` lowers to a native DataFusion
+    // expression and `hypot` is a `CometCodegenDispatch`, so this query exercises both buckets at
+    // once and each wrong claim below must be rejected.
+    withTable("t") {
+      sql("CREATE TABLE t (a DOUBLE, b DOUBLE) USING parquet")
+      sql("INSERT INTO t VALUES (3.0, 4.0)")
+      val query = "SELECT abs(a), hypot(a, b) FROM t"
+
+      checkSparkAnswerAndImpl(sql(query), native = Seq("abs"), dispatched = Seq("hypot"))
+
+      // Claiming the wrong mechanism fails, in both directions.
+      intercept[TestFailedException] {
+        checkSparkAnswerAndImpl(sql(query), native = Seq("hypot"))
+      }
+      intercept[TestFailedException] {
+        checkSparkAnswerAndImpl(sql(query), dispatched = Seq("abs"))
+      }
+      // So does naming an expression the query does not contain, which is what a typo in a
+      // fixture looks like.
+      intercept[TestFailedException] {
+        checkSparkAnswerAndImpl(sql(query), native = Seq("no_such_expression"))
+      }
+      intercept[TestFailedException] {
+        checkSparkAnswerAndImpl(sql(query), dispatched = Seq("no_such_expression"))
+      }
+    }
+  }
+
   test("expression coverage stats split native from codegen-dispatch expressions") {
     // `abs` and `sqrt` lower to native DataFusion expressions; `hypot` and `nanvl` are
     // `CometCodegenDispatch` and so run Spark's own codegen inside the Comet pipeline. The

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

@@ -155,6 +155,10 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper {
                     checkSparkAnswerAndOperatorWithTolerance(sql, tol)
                   case ExpectFallback(reason) =>
                     checkSparkAnswerAndFallbackReason(sql, reason)
+                  case ExpectDispatch(names) =>
+                    checkSparkAnswerAndImpl(sql, native = Seq.empty, dispatched = names)
+                  case ExpectNative(names) =>
+                    checkSparkAnswerAndImpl(sql, native = names, dispatched = Seq.empty)
                   case Ignore(reason) =>
                     logInfo(s"IGNORED query ($reason): $sql")
                   case ExpectError(pattern) =>

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

@@ -56,6 +56,23 @@ case class WithTolerance(tol: Double) extends QueryAssertionMode
 case class ExpectFallback(reason: String) extends QueryAssertionMode
 case class Ignore(reason: String) extends QueryAssertionMode
 
+/**
+ * Checks results and coverage like [[CheckCoverageAndAnswer]], and additionally asserts that
+ * Comet ran each named expression through the JVM codegen dispatcher rather than lowering it to a
+ * native DataFusion expression.
+ *
+ * Matching results cannot distinguish the two mechanisms, so without this a serde that gains a
+ * native path (or loses its dispatcher route) changes how the query executes while every other
+ * assertion in the fixture stays green.
+ */
+case class ExpectDispatch(names: Seq[String]) extends QueryAssertionMode
+
+/**
+ * The native counterpart of [[ExpectDispatch]]: asserts Comet lowered each named expression to a
+ * native DataFusion expression rather than routing it through the JVM codegen dispatcher.
+ */
+case class ExpectNative(names: Seq[String]) extends QueryAssertionMode
+
 /**
  * Asserts that both Spark and Comet raise an error whose message contains `pattern`.
  *
@@ -177,6 +194,8 @@ object SqlFileTestParser {
   private val FallbackPattern = """query\s+expect_fallback\((.+)\)""".r
   private val IgnorePattern = """query\s+ignore\((.+)\)""".r
   private val ErrorPattern = """query\s+expect_error\((.+)\)""".r
+  private val DispatchPattern = """query\s+expect_dispatch\((.+)\)""".r
+  private val NativePattern = """query\s+expect_native\((.+)\)""".r
 
   private def parseQueryAssertionMode(directive: String): QueryAssertionMode = {
     directive match {
@@ -186,6 +205,10 @@ object SqlFileTestParser {
         Ignore(reason.trim)
       case ErrorPattern(pattern) =>
         ExpectError(pattern.trim)
+      case DispatchPattern(names) =>
+        ExpectDispatch(splitNames(names))
+      case NativePattern(names) =>
+        ExpectNative(splitNames(names))
       case _ =>
         val parts = directive.split("\\s+")
         if (parts.length == 1) return CheckCoverageAndAnswer
@@ -198,6 +221,13 @@ object SqlFileTestParser {
     }
   }
 
+  /**
+   * Split a comma-separated expression-name list, dropping empties so `expect_dispatch(a, b,)`
+   * and stray whitespace do not produce a name that can never match.
+   */
+  private def splitNames(names: String): Seq[String] =
+    names.split(",").map(_.trim).filter(_.nonEmpty).toSeq
+
   /** Collect SQL lines until a blank line or end of file. */
   private def collectSql(lines: Seq[String], start: Int): (String, Int) = {
     val sb = new StringBuilder