test: cover JSON and cast expression routing

spark/src/test/resources/sql-tests/expressions/cast/routing_cast_disabled.sql

@@ -0,0 +1,47 @@
+-- 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.
+
+-- Config: spark.comet.exec.scalaUDF.codegen.enabled=false
+-- Config: spark.sql.legacy.castComplexTypesToString.enabled=true
+-- ConfigMatrix: spark.comet.expression.Cast.allowIncompatible=false,true
+
+statement
+CREATE TABLE routing_cast(i INT, b BOOLEAN, a ARRAY<INT>, s STRUCT<v: INT>, m MAP<STRING, INT>) USING parquet
+
+statement
+INSERT INTO routing_cast VALUES
+  (1, true, array(1, null, 3), named_struct('v', 1), map('a', 1, 'b', null)),
+  (0, false, array(), named_struct('v', null), map()),
+  (NULL, NULL, NULL, NULL, NULL)
+
+-- Compatible casts stay native regardless of dispatcher and opt-in settings.
+query expect_native(cast)
+SELECT CAST(i AS BIGINT) FROM routing_cast
+
+-- These casts have no native path, even with allowIncompatible enabled.
+query expect_fallback(cast: spark.comet.exec.scalaUDF.codegen.enabled=false)
+SELECT CAST(b AS DECIMAL(10, 2)) FROM routing_cast
+
+-- Legacy complex-to-string formatting requires Spark's implementation.
+query expect_fallback(cast: spark.comet.exec.scalaUDF.codegen.enabled=false)
+SELECT CAST(a AS STRING) FROM routing_cast
+
+query expect_fallback(cast: spark.comet.exec.scalaUDF.codegen.enabled=false)
+SELECT CAST(s AS STRING) FROM routing_cast
+
+query expect_fallback(cast: spark.comet.exec.scalaUDF.codegen.enabled=false)
+SELECT CAST(m AS STRING) FROM routing_cast

spark/src/test/resources/sql-tests/expressions/cast/routing_cast_enabled.sql

@@ -0,0 +1,47 @@
+-- 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.
+
+-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true
+-- Config: spark.sql.legacy.castComplexTypesToString.enabled=true
+-- ConfigMatrix: spark.comet.expression.Cast.allowIncompatible=false,true
+
+statement
+CREATE TABLE routing_cast(i INT, b BOOLEAN, a ARRAY<INT>, s STRUCT<v: INT>, m MAP<STRING, INT>) USING parquet
+
+statement
+INSERT INTO routing_cast VALUES
+  (1, true, array(1, null, 3), named_struct('v', 1), map('a', 1, 'b', null)),
+  (0, false, array(), named_struct('v', null), map()),
+  (NULL, NULL, NULL, NULL, NULL)
+
+-- Compatible casts stay native regardless of dispatcher and opt-in settings.
+query expect_native(cast)
+SELECT CAST(i AS BIGINT) FROM routing_cast
+
+-- These casts have no native path, even with allowIncompatible enabled.
+query expect_dispatch(cast)
+SELECT CAST(b AS DECIMAL(10, 2)) FROM routing_cast
+
+-- Legacy complex-to-string formatting requires Spark's implementation.
+query expect_dispatch(cast)
+SELECT CAST(a AS STRING) FROM routing_cast
+
+query expect_dispatch(cast)
+SELECT CAST(s AS STRING) FROM routing_cast
+
+query expect_dispatch(cast)
+SELECT CAST(m AS STRING) FROM routing_cast

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

@@ -39,6 +39,7 @@ import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus
 import org.apache.comet.codegen.CometBatchKernelCodegen
 import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec
 import org.apache.comet.serde.{CometScalaUDF, QueryPlanSerde}
+import org.apache.comet.serde.ExprOuterClass.Expr.ExprStructCase
 import org.apache.comet.udf.codegen.CometScalaUDFCodegen
 import org.apache.comet.vector.CometVector
 
@@ -124,6 +125,75 @@ class CometCodegenSuite
     }
   }
 
+  for {
+    (name, configName, nativeKind, expressions) <- Seq(
+      (
+        "from_json",
+        "JsonToStructs",
+        ExprStructCase.FROM_JSON,
+        Seq(
+          "from_json(j, 'a INT, b STRING')" -> true,
+          "from_json(j, 'a INT, arr ARRAY<INT>')" -> false)),
+      (
+        "to_json",
+        "StructsToJson",
+        ExprStructCase.TO_JSON,
+        Seq(
+          "to_json(s)" -> true,
+          "to_json(a)" -> false,
+          "to_json(s, map('ignoreNullFields', 'false'))" -> false)))
+  } {
+    test(s"$name routing follows native opt-in and dispatcher settings") {
+      withTable("json_routing") {
+        sql("""CREATE TABLE json_routing(j STRING, s STRUCT<a: INT, b: STRING>, a ARRAY<INT>)
+              |USING parquet""".stripMargin)
+        sql("""INSERT INTO json_routing VALUES
+              |('{"a":1,"b":"x","arr":[1,null,3]}', named_struct('a', 1, 'b', 'x'), array(1, null, 3)),
+              |('{"a":null,"b":"","arr":[]}', named_struct('a', null, 'b', ''), array()),
+              |('{}', named_struct('a', null, 'b', null), array()),
+              |(NULL, NULL, NULL)""".stripMargin)
+        for {
+          allowIncompatible <- Seq(false, true)
+          codegenEnabled <- Seq(false, true)
+        } {
+          withSQLConf(
+            s"spark.comet.expression.$configName.allowIncompatible" ->
+              allowIncompatible.toString,
+            CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> codegenEnabled.toString) {
+            expressions.foreach { case (expression, nativeSupported) =>
+              withClue(
+                s"allowIncompatible=$allowIncompatible, codegen=$codegenEnabled: $expression") {
+                val query = s"SELECT $expression FROM json_routing"
+                val expectNative = allowIncompatible && nativeSupported
+                if (!expectNative && !codegenEnabled) {
+                  checkSparkAnswerAndFallbackReason(
+                    query,
+                    s"$name: spark.comet.exec.scalaUDF.codegen.enabled=false")
+                } else {
+                  val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query))
+                  // Spark 4 rewrites to_json to Invoke without preserving implementation tags.
+                  // Inspect the executable expression to distinguish native from dispatch.
+                  val expr = stripAQEPlan(cometPlan)
+                    .collectFirst { case project: CometProjectExec =>
+                      project.nativeOp.getProjection.getProjectList(0)
+                    }
+                    .getOrElse(fail("Expected a Comet projection"))
+                  if (expectNative) {
+                    assert(expr.getExprStructCase === nativeKind)
+                  } else {
+                    assert(expr.hasJvmScalarUdf)
+                    assert(
+                      expr.getJvmScalarUdf.getClassName === classOf[CometScalaUDFCodegen].getName)
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
   private def withTwoStringCols(rows: (String, String)*)(f: => Unit): Unit = {
     withTable("t") {
       sql("CREATE TABLE t (c1 STRING, c2 STRING) USING parquet")