From 3921c0049eefc92cd149b2d74a78953820cb20a4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 1 Sep 2026 08:12:12 -0600 Subject: [PATCH] test: add helpers to assert whether an expression ran natively or via codegen dispatch Comet evaluates a scalar expression natively, through the JVM codegen dispatcher, or not at all. Only the third is visible to our tests: checkSparkAnswerAndOperator and the default sql-file query mode assert "no fallback", but native and dispatched execution both produce Spark-matching results by construction, so a serde that swaps one for the other passes every existing assertion unchanged. Add checkSparkAnswerAndImpl to CometTestBase and expect_native / expect_dispatch query modes to the sql-file harness, both backed by the same assertion over ExtendedExplainInfo's native and codegen-dispatch expression sets. Naming an expression asserts both that it ran through the expected mechanism and that it did not run through the other one. Annotate round.sql, lower.sql and upper.sql as worked examples. Doing so split one round query: round(NULL, 0) implicitly casts the untyped null to double, so it dispatches while the decimal literals beside it stay native. Closes #5609. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../adding_a_new_expression.md | 29 +++++++ .../contributor-guide/sql-file-tests.md | 45 ++++++++++ .../sql-tests/expressions/math/round.sql | 24 ++++-- .../sql-tests/expressions/string/lower.sql | 6 +- .../sql-tests/expressions/string/upper.sql | 6 +- .../org/apache/comet/CometCodegenSuite.scala | 31 +++++++ .../apache/comet/CometSqlFileTestSuite.scala | 4 + .../org/apache/comet/SqlFileTestParser.scala | 30 +++++++ .../apache/comet/SqlFileTestParserSuite.scala | 84 +++++++++++++++++++ .../org/apache/spark/sql/CometTestBase.scala | 78 +++++++++++++++++ 12 files changed, 324 insertions(+), 15 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 31cc3c34da9..7ab2825c457 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.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 diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 7b7fae8fbea..877cc649129 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.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 diff --git a/docs/source/contributor-guide/adding_a_new_expression.md b/docs/source/contributor-guide/adding_a_new_expression.md index 2a92139830d..d9e7707bab4 100644 --- a/docs/source/contributor-guide/adding_a_new_expression.md +++ b/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 diff --git a/docs/source/contributor-guide/sql-file-tests.md b/docs/source/contributor-guide/sql-file-tests.md index a92e0898d71..c7f0171200f 100644 --- a/docs/source/contributor-guide/sql-file-tests.md +++ b/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()` / `query expect_native()` + +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()` 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 diff --git a/spark/src/test/resources/sql-tests/expressions/math/round.sql b/spark/src/test/resources/sql-tests/expressions/math/round.sql index 9974d872a82..0ea7301ab1f 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/round.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/round.sql @@ -35,10 +35,10 @@ 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. @@ -46,10 +46,10 @@ 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) diff --git a/spark/src/test/resources/sql-tests/expressions/string/lower.sql b/spark/src/test/resources/sql-tests/expressions/string/lower.sql index e21bf7710c3..f49765695c5 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/lower.sql +++ b/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 diff --git a/spark/src/test/resources/sql-tests/expressions/string/upper.sql b/spark/src/test/resources/sql-tests/expressions/string/upper.sql index c4f64376eb0..ded03942c72 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/upper.sql +++ b/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 diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 5806cb35015..53e9beda03e 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/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 diff --git a/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala index 01c4c972b3f..0b0a7615a9e 100644 --- a/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala +++ b/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) => diff --git a/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala b/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala index d6e42b8833d..2b7c573bb19 100644 --- a/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala +++ b/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 diff --git a/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala b/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala new file mode 100644 index 00000000000..29e9fdc0b63 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala @@ -0,0 +1,84 @@ +/* + * 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.scalatest.funsuite.AnyFunSuite + +/** + * Unit tests for [[SqlFileTestParser]]. Pure text parsing, so no Spark session is needed. The + * end-to-end behaviour of each query mode is covered by `CometSqlFileTestSuite` running the + * fixtures themselves. + */ +class SqlFileTestParserSuite extends AnyFunSuite { + + private def parseQueries(lines: String*): Seq[SqlQuery] = + SqlFileTestParser.parse(lines).records.collect { case q: SqlQuery => q } + + private def modeOf(directive: String): QueryAssertionMode = + parseQueries(directive, "SELECT 1").head.mode + + test("bare query directive defaults to checking coverage and answer") { + assert(modeOf("query") === CheckCoverageAndAnswer) + } + + test("expect_dispatch parses a single expression name") { + assert(modeOf("query expect_dispatch(bit_length)") === ExpectDispatch(Seq("bit_length"))) + } + + test("expect_native parses a single expression name") { + assert(modeOf("query expect_native(length)") === ExpectNative(Seq("length"))) + } + + test("expect_dispatch parses a comma-separated list and trims whitespace") { + assert( + modeOf("query expect_dispatch(rlike, regexp_replace ,split)") === + ExpectDispatch(Seq("rlike", "regexp_replace", "split"))) + } + + test("expect_native tolerates extra whitespace around the directive") { + assert(modeOf("query expect_native( round , abs )") === ExpectNative(Seq("round", "abs"))) + } + + test("empty names are dropped rather than becoming unmatchable entries") { + // A name that is the empty string could never appear in the plan's expression set, so it + // would fail the assertion for a reason that has nothing to do with the query. + assert(modeOf("query expect_dispatch(lower,,)") === ExpectDispatch(Seq("lower"))) + } + + test("the new modes do not shadow the existing ones") { + assert(modeOf("query expect_fallback(some reason)") === ExpectFallback("some reason")) + assert(modeOf("query expect_error(DIVIDE_BY_ZERO)") === ExpectError("DIVIDE_BY_ZERO")) + assert(modeOf("query spark_answer_only") === SparkAnswerOnly) + assert(modeOf("query tolerance=0.001") === WithTolerance(0.001)) + assert( + modeOf("query ignore(https://example.com/issue)") === Ignore("https://example.com/issue")) + } + + test("query mode and SQL text are associated with the right record") { + val queries = parseQueries( + "query expect_native(abs)", + "SELECT abs(a) FROM t", + "", + "query expect_dispatch(hypot)", + "SELECT hypot(a, b) FROM t") + assert(queries.map(_.mode) === Seq(ExpectNative(Seq("abs")), ExpectDispatch(Seq("hypot")))) + assert(queries.map(_.sql) === Seq("SELECT abs(a) FROM t", "SELECT hypot(a, b) FROM t")) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 975ba509ccc..4b9079abd11 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -342,6 +342,84 @@ abstract class CometTestBase (sparkPlan, cometPlan) } + /** + * Check for the correct results, that Comet replaced all possible operators, and that the named + * expressions ran through the mechanism the caller expects. + * + * Comet evaluates an expression one of three ways: natively (a DataFusion expression), through + * the JVM codegen dispatcher (Spark's own `doGenCode` compiled into an Arrow batch kernel), or + * not at all (the operator falls back to Spark). Only the third is visible to + * [[checkSparkAnswerAndOperator]]; the first two produce Spark-matching results by + * construction, so a serde that quietly widens from native to dispatch (losing the native + * kernel) or narrows from dispatch to native (losing Spark-exact semantics) passes every other + * assertion here. Use this to pin which one actually ran. + * + * Names are the expression's `prettyName` lowercased, as [[ExtendedExplainInfo]] reports them + * (`bit_length`, `octet_length`, `rlike`), not necessarily the SQL alias used to invoke it: a + * function registered with `setAlias` reports the invoked alias, everything else reports its + * own `prettyName`. + * + * For fallback assertions use [[checkSparkAnswerAndFallbackReason]] instead. + */ + protected def checkSparkAnswerAndImpl( + df: => DataFrame, + native: Seq[String] = Seq.empty, + dispatched: Seq[String] = Seq.empty): (SparkPlan, SparkPlan) = { + val (sparkPlan, cometPlan) = checkSparkAnswerAndOperator(df) + assertExpressionImpl(cometPlan, native, dispatched) + (sparkPlan, cometPlan) + } + + /** Check for the correct results and the expected per-expression implementation. */ + protected def checkSparkAnswerAndImpl( + query: String, + native: Seq[String], + dispatched: Seq[String]): (SparkPlan, SparkPlan) = { + checkSparkAnswerAndImpl(sql(query), native, dispatched) + } + + /** + * Assert how Comet evaluated the named expressions in an already-executed Comet plan. Split out + * from [[checkSparkAnswerAndImpl]] so callers holding a plan can reuse it, and so the assertion + * itself is testable. + * + * Each name must appear in its expected set and must be absent from the other, so naming an + * expression is a claim about which mechanism ran it rather than a claim that it ran somehow. + */ + protected def assertExpressionImpl( + cometPlan: SparkPlan, + native: Seq[String], + dispatched: Seq[String]): Unit = { + val explainInfo = new ExtendedExplainInfo() + val actualNative = explainInfo.getNativeExpressions(cometPlan) + val actualDispatched = explainInfo.getCodegenDispatchExpressions(cometPlan) + def detail: String = + s"native=[${actualNative.mkString(", ")}] " + + s"codegen-dispatched=[${actualDispatched.mkString(", ")}]" + native.foreach { name => + if (actualDispatched.contains(name)) { + fail( + s"Expected `$name` to run as a native expression but it ran through the JVM " + + s"codegen dispatcher. Actual: $detail") + } + if (!actualNative.contains(name)) { + fail(s"Expected `$name` to run as a native expression but it did not. Actual: $detail") + } + } + dispatched.foreach { name => + if (actualNative.contains(name)) { + fail( + s"Expected `$name` to run through the JVM codegen dispatcher but it ran as a " + + s"native expression. Actual: $detail") + } + if (!actualDispatched.contains(name)) { + fail( + s"Expected `$name` to run through the JVM codegen dispatcher but it did not. " + + s"Actual: $detail") + } + } + } + /** * Try executing the query against Spark and Comet and return the results or the exception. *