diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index cbb73fab985..76761d01120 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -584,6 +584,7 @@ jobs: org.apache.comet.CometTemporalExpressionSuite org.apache.comet.CometArrayExpressionSuite org.apache.comet.CometNativeCastSuite + org.apache.comet.CometCastCollatedStringSuite org.apache.comet.CometDateTimeUtilsSuite org.apache.comet.CometMathExpressionSuite org.apache.comet.CometStringExpressionSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index bcf73b8a56e..948e375fc3b 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -232,6 +232,7 @@ jobs: org.apache.comet.CometTemporalExpressionSuite org.apache.comet.CometArrayExpressionSuite org.apache.comet.CometNativeCastSuite + org.apache.comet.CometCastCollatedStringSuite org.apache.comet.CometDateTimeUtilsSuite org.apache.comet.CometMathExpressionSuite org.apache.comet.CometStringExpressionSuite diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index 83fbca6b635..eadda24da70 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -87,6 +87,21 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come * single child and the generated writer NPEs on the missing ordinal-1 vector. * `CometCreateNamedStruct` declines them on the native path for the same reason, but a struct * nested inside a dispatcher-built value (a `CreateMap` value) never reaches that check. + * + * The `StringType` case admits non-default collations (Spark 4+), on purpose. Collation is + * carried by the expression, not by the value: the kernel runs Spark's own `doGenCode` against + * the bound tree, whose `collationId` survives closure serialization, and Arrow is only the + * byte store for the `UTF8String`s that code produces. So a dispatched expression over collated + * input answers exactly as Spark does. + * + * The proto type is a separate matter. `CometScalaUDF.emitJvmCodegenDispatch` declares the + * return type through `QueryPlanSerde.serializeDataType`, which flattens every `StringType` to + * one proto id, so the native plan describes a dispatched collated output as a plain string. + * Nothing on this route reads that back. The values are bytes and the kernel is what produced + * them. Whether a downstream operator may then treat the column collation-blind is decided per + * operator against the Catalyst `DataType`, which keeps its collation, and does not depend on + * what this predicate admits. Rejecting collated strings here would force a full Spark fallback + * for a route that is already correct. */ def isSupportedDataType(dt: DataType): Boolean = dt match { case BooleanType | ByteType | ShortType | IntegerType | LongType => true diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala index d29ef7cd3b7..b1799ed1f2c 100644 --- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala +++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala @@ -50,6 +50,16 @@ object CometCast private[comet] val legacyCastComplexTypesToStringReason: String = "spark.sql.legacy.castComplexTypesToString.enabled=true is not supported" + // The generic `Cast from $fromType to $toType is not supported` template is unhelpful for a + // collation rejection. An identity cast prints both sides identically, so it reads like a + // nonsensical refusal of a string-to-string cast unless the reader already knows collation is + // the cause. Named here, and shared with `CometCastCollatedStringSuite`, so the asserted reason + // cannot drift from production. Follows the phrasing the other collation reasons use, e.g. + // `CometReverse` in `collectionOperations.scala` and `ComparisonUtils` in `predicates.scala`. + private[comet] val nonDefaultCollationReason: String = + "Cast involving a non-default string collation is not supported " + + "(https://github.com/apache/datafusion-comet/issues/4489)" + private def legacyCastComplexTypesToString: Boolean = SQLConf.get .getConfString("spark.sql.legacy.castComplexTypesToString.enabled", "false") @@ -186,6 +196,24 @@ object CometCast return unsupported(fromType, toType) } + // Spark 4.0's collation metadata rides on `StringType`, but `serializeDataType` maps every + // `StringType` to the same proto id, so a non-default collation is dropped on the way into + // the native plan with no warning. Reject the cast outright rather than relying on the + // pattern matching below, which only misses collated types because `DataTypes.StringType` is + // the default-collation singleton and Scala pattern equality happens not to match. This runs + // above the `fromType == toType` shortcut so that an identity cast on a collated type is + // checked too, and `hasNonDefaultStringCollation` walks nested element, key, value, and field + // types. The version-shimmed helper returns false on Spark 3.x, where collation does not + // exist. See https://github.com/apache/datafusion-comet/issues/4489. + // + // `Unsupported` here means there is no native path, not that the plan falls back to Spark. + // `CodegenDispatchFallback` offers the cast to the JVM codegen dispatcher first, and that + // route is result-correct: see the note on `isSupportedDataType` in + // `CometBatchKernelCodegen` for why a collated string is safe to admit there. + if (hasNonDefaultStringCollation(fromType) || hasNonDefaultStringCollation(toType)) { + return Unsupported(Some(nonDefaultCollationReason)) + } + if (fromType == toType) { return Compatible() } diff --git a/spark/src/test/spark-4.x/org/apache/comet/CometCastCollatedStringSuite.scala b/spark/src/test/spark-4.x/org/apache/comet/CometCastCollatedStringSuite.scala new file mode 100644 index 00000000000..f547836a7b7 --- /dev/null +++ b/spark/src/test/spark-4.x/org/apache/comet/CometCastCollatedStringSuite.scala @@ -0,0 +1,273 @@ +/* + * 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.spark.sql.CometTestBase +import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, IntegerType, MapType, StringType, StructField, StructType} + +import org.apache.comet.expressions.{CometCast, CometEvalMode} +import org.apache.comet.serde.{Compatible, Unsupported} + +/** + * https://github.com/apache/datafusion-comet/issues/4489 + * + * Spark 4.0 carries collation metadata on `StringType`, but `serializeDataType` maps every + * `StringType` to one proto type id, so the collation is dropped on the way into the native plan + * with no warning. Nothing used to stop that from happening. `CometCast.isSupported` matched + * string casts through `case (DataTypes.StringType, _)` and `case (_, DataTypes.StringType)`, and + * those only failed to match a collated `StringType` because `DataTypes.StringType` is the + * default-collation singleton and Scala pattern equality compares the whole instance. The right + * answer fell out of an accident of pattern matching. The `fromType == toType` shortcut let + * identity casts on collated types through regardless, including nested ones such as + * `ARRAY`, which were byte-safe only because the cast was a no-op. + * + * `isSupported` now rejects any source or target type carrying a non-default collation before + * either of those paths runs. This suite pins the resulting matrix down so the behaviour no + * longer depends on pattern-match semantics staying the way they are. + * + * A note on what `Unsupported` means here. It does not mean the query falls back to Spark. + * `CometCast` mixes in `CodegenDispatchFallback`, so `QueryPlanSerde.exprToProtoInternal` offers + * the expression to the JVM codegen dispatcher first, which runs Spark's own `doGenCode` inside + * the Comet pipeline. `spark.comet.exec.scalaUDF.codegen.enabled` defaults to true and + * `CometBatchKernelCodegen.isSupportedDataType` admits every `StringType` regardless of + * collation, so under default config a collated cast usually stays inside Comet. That route is + * result-correct, because the kernel evaluates Spark's own generated code and the `collationId` + * rides along on the expression. `Unsupported` means there is no native path. The end-to-end + * tests at the bottom run both settings of that config on both the scalar and the struct case. + * + * These are Scala tests rather than `CometSqlFileTestSuite` fixtures because most of the file + * asserts on `CometCast.isSupported` over type pairs no SQL can construct, such as + * `ArrayType(NullType) -> ArrayType(STRING COLLATE UTF8_LCASE)`. The end-to-end tests stay here + * with them because `--Config` is file scoped, so the query that has to assert a fallback reason + * with the dispatcher off and the query that has to assert native execution with it on could not + * share a fixture file. + * + * This lives under `spark-4.x`, shared by every 4.x profile, rather than `spark-4.1+`, because + * collation is a Spark 4.0 feature and `StringType(collationName)` already resolves there. + */ +class CometCastCollatedStringSuite extends CometTestBase { + + private val lcase = StringType("UTF8_LCASE") + private val unicode = StringType("UNICODE") + + private val evalModes = Seq(CometEvalMode.LEGACY, CometEvalMode.TRY, CometEvalMode.ANSI) + + private def structWith(dt: DataType): StructType = StructType(Seq(StructField("s", dt))) + + /** Asserts that Comet reports no native path for this cast under every eval mode. */ + private def assertNoNativePath(fromType: DataType, toType: DataType): Unit = { + evalModes.foreach { evalMode => + CometCast.isSupported(fromType, toType, None, evalMode) match { + case _: Unsupported => // expected + case other => + fail(s"expected Unsupported for $fromType -> $toType under $evalMode, got $other") + } + } + } + + /** Asserts that the collation guard leaves an uncollated cast alone. */ + private def assertCompatible(fromType: DataType, toType: DataType): Unit = { + evalModes.foreach { evalMode => + CometCast.isSupported(fromType, toType, None, evalMode) match { + case _: Compatible => // expected + case other => + fail(s"expected Compatible for $fromType -> $toType under $evalMode, got $other") + } + } + } + + // ---- scalar collated strings ---------------------------------------------------- + + test("cast collated string to IntegerType has no native path") { + assertNoNativePath(lcase, IntegerType) + } + + test("cast IntegerType to collated string has no native path") { + // This pair leaves through `canCastFromInt`'s catch-all rather than `canCastToString`, + // because `case (_, DataTypes.StringType)` does not match a collated target. The guard now + // answers ahead of both. + assertNoNativePath(IntegerType, lcase) + } + + test("cast collated string to default-collation StringType has no native path") { + assertNoNativePath(lcase, DataTypes.StringType) + } + + test("cast default-collation StringType to collated string has no native path") { + assertNoNativePath(DataTypes.StringType, lcase) + } + + test("cast between two different collations has no native path") { + assertNoNativePath(lcase, unicode) + } + + test("cast collated string to the same collation has no native path") { + // The `fromType == toType` shortcut answered `Compatible()` here before the guard existed. + // The result would be right, since the cast is a byte-level no-op, but `isSupported` was + // clearing a collated type for a proto that cannot carry the collation, and nothing recorded + // that. This is the implicit behaviour #4489 names, so the guard sits above the shortcut. + assertNoNativePath(lcase, lcase) + } + + // ---- nested collated strings ---------------------------------------------------- + + test("cast array of collated strings to another collation has no native path") { + assertNoNativePath(ArrayType(lcase), ArrayType(unicode)) + } + + test("cast array of collated strings to the same collation has no native path") { + // Same identity shortcut as the scalar case, one level down. + assertNoNativePath(ArrayType(lcase), ArrayType(lcase)) + } + + test("cast array of collated strings to StringType has no native path") { + // `case (dt: ArrayType, DataTypes.StringType)` recurses on the element type, so the reason + // used to describe the element rather than the array. The guard answers for the whole type. + assertNoNativePath(ArrayType(lcase), DataTypes.StringType) + } + + test("cast struct with a collated field has no native path") { + assertNoNativePath(structWith(lcase), structWith(unicode)) + assertNoNativePath(structWith(lcase), structWith(lcase)) + } + + test("cast map with a collated key has no native path") { + assertNoNativePath(MapType(lcase, IntegerType), MapType(unicode, IntegerType)) + assertNoNativePath(MapType(lcase, IntegerType), MapType(lcase, IntegerType)) + } + + test("cast map with a collated value has no native path") { + assertNoNativePath(MapType(IntegerType, lcase), MapType(IntegerType, unicode)) + assertNoNativePath(MapType(IntegerType, lcase), MapType(IntegerType, lcase)) + } + + test("cast struct whose collated field is unchanged while a sibling field is cast") { + // The field zip used to answer per field, so the collated field hit the identity shortcut + // and reported Compatible while the sibling carried the cast. That cleared a collated field + // for the native plan on another field's back. The guard answers for the whole struct. + val from = StructType(Seq(StructField("a", IntegerType), StructField("s", lcase))) + val to = StructType(Seq(StructField("a", DataTypes.StringType), StructField("s", lcase))) + assertNoNativePath(from, to) + } + + test("cast map whose collated key is unchanged while the value type is cast") { + assertNoNativePath(MapType(lcase, IntegerType), MapType(lcase, DataTypes.LongType)) + } + + test("cast array of nulls to array of collated strings has no native path") { + // `case (dt: ArrayType, _: ArrayType) if dt.elementType == NullType` returns Compatible + // without looking at the target element type, so this was one more pair where + // `isSupported` cleared a collated target type. + assertNoNativePath(ArrayType(DataTypes.NullType), ArrayType(lcase)) + } + + // ---- the guard must not over-block ---------------------------------------------- + + test("default-collation string casts are untouched by the collation guard") { + assertCompatible(DataTypes.StringType, DataTypes.StringType) + assertCompatible(DataTypes.StringType, IntegerType) + assertCompatible(IntegerType, DataTypes.StringType) + } + + test("nested default-collation string casts are untouched by the collation guard") { + assertCompatible(ArrayType(DataTypes.StringType), ArrayType(DataTypes.StringType)) + assertCompatible(structWith(DataTypes.StringType), structWith(DataTypes.StringType)) + assertCompatible( + MapType(DataTypes.StringType, IntegerType), + MapType(DataTypes.StringType, IntegerType)) + } + + // ---- end to end ----------------------------------------------------------------- + // + // The matrix above only exercises `isSupported`. These four run a query and pin down what the + // planner actually does with the answer, which is what #4489 asked for. A plain-string Parquet + // column with `COLLATE` applied on top reaches the cast as a collated child, the same shape + // the datetime tests in `CometCollationSuite` rely on. The cast child has to be a column + // rather than a literal, since `getSupportLevel` folds literal children before `isSupported` + // is consulted. + + private def withCollatedTable(f: => Unit): Unit = + withParquetTable(Seq(("123", 1), ("456", 2)), "collated_cast_tbl")(f) + + // Read from production rather than retyped, so the two cannot drift apart. The guard is the + // only thing that produces this reason: without it `StringType(UTF8_LCASE) -> IntegerType` + // exits through `isSupported`'s `case _` catch-all and reports the generic + // `Cast from ... to ... is not supported` template instead. + private val collationReason = CometCast.nonDefaultCollationReason + + // A scalar identity pair such as `lcase -> lcase` cannot be reached from SQL. Spark's + // `SimplifyCasts` drops a cast whose child already carries the target type, so + // `CAST(_1 COLLATE utf8_lcase AS STRING COLLATE UTF8_LCASE)` arrives at the planner as a bare + // `Collate` and the only fallback reason on the plan is "collate is not supported", which comes + // from a different serde. Those pairs stay pinned at the `isSupported` level above, in the same + // spirit as the join tests in `CometCollationSuite` that no query can reach. A struct is the + // way in: when a sibling field changes type the cast survives, and the collated field rides + // along inside it. That case is covered end to end below. + + test("cast from a collated string falls back to Spark when codegen dispatch is off") { + withCollatedTable { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + checkSparkAnswerAndFallbackReason( + "SELECT CAST(_1 COLLATE utf8_lcase AS INT) FROM collated_cast_tbl", + collationReason) + } + } + } + + test("cast from a collated string routes through the codegen dispatcher when it is on") { + withCollatedTable { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + checkSparkAnswerAndOperator( + "SELECT CAST(_1 COLLATE utf8_lcase AS INT) FROM collated_cast_tbl") + } + } + } + + test("cast of a struct carrying a collated field has no native path end to end") { + // A pair the guard changes the outcome for, not just the reason string. The sibling field + // changes type so the cast survives `SimplifyCasts`. Without the guard the field zip + // answered `Compatible`, since the collated field matched the identity shortcut, but the + // query still fell back to Spark: `Collate` has no serde, so serializing the struct stopped + // at the collated field with "collate is not supported". With the guard the cast itself + // is `Unsupported`: with the dispatcher off the plan carries the collation reason, and with + // it on (next test) the query moves out of that Spark fallback and into the dispatcher. + withCollatedTable { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + checkSparkAnswerAndFallbackReason( + "SELECT CAST(struct(_2 AS a, _1 COLLATE utf8_lcase AS s) AS " + + "STRUCT) FROM collated_cast_tbl", + collationReason) + } + } + } + + test("cast of a struct carrying a collated field routes through the codegen dispatcher") { + // The struct is the one case the guard redirects rather than merely relabels, so it needs the + // dispatcher-on half too. `isSupportedDataType` recurses into the fields and accepts them, so + // the cast stays in the Comet pipeline and the kernel runs Spark's own `Cast.doGenCode`. + withCollatedTable { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + checkSparkAnswerAndOperator( + "SELECT CAST(struct(_2 AS a, _1 COLLATE utf8_lcase AS s) AS " + + "STRUCT) FROM collated_cast_tbl") + } + } + } +}