From 9caeb8d2e1b98bf2b1309d73a8a55f991daefa13 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Fri, 18 Sep 2026 13:16:33 +0800 Subject: [PATCH] test: cover JSON and cast expression routing --- .../cast/routing_cast_disabled.sql | 47 +++++++++++++ .../expressions/cast/routing_cast_enabled.sql | 47 +++++++++++++ .../org/apache/comet/CometCodegenSuite.scala | 70 +++++++++++++++++++ 3 files changed, 164 insertions(+) create mode 100644 spark/src/test/resources/sql-tests/expressions/cast/routing_cast_disabled.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/cast/routing_cast_enabled.sql diff --git a/spark/src/test/resources/sql-tests/expressions/cast/routing_cast_disabled.sql b/spark/src/test/resources/sql-tests/expressions/cast/routing_cast_disabled.sql new file mode 100644 index 00000000000..daee1ff1361 --- /dev/null +++ b/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, s STRUCT, m MAP) 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 diff --git a/spark/src/test/resources/sql-tests/expressions/cast/routing_cast_enabled.sql b/spark/src/test/resources/sql-tests/expressions/cast/routing_cast_enabled.sql new file mode 100644 index 00000000000..b5713ddc581 --- /dev/null +++ b/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, s STRUCT, m MAP) 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 diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 149ee7a1454..7bb8fc0a9d8 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/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')" -> 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 ARRAY) + |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")