diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index 4b24abc140..cfdf519665 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -21,7 +21,7 @@ package org.apache.comet.serde import java.util.Locale -import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeTimestamp, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, NextDay, Quarter, Second, SecondsToTimestamp, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year} +import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeTimestamp, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, NextDay, Quarter, Second, SecondsToTimestamp, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataType, DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String @@ -32,6 +32,7 @@ import org.apache.comet.expressions.{CometCast, CometEvalMode} import org.apache.comet.serde.CometGetDateField.CometGetDateField import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.QueryPlanSerde._ +import org.apache.comet.shims.CometTypeShim private object CometGetDateField extends Enumeration { type CometGetDateField = Value @@ -310,11 +311,28 @@ object CometSecond extends CometExpressionSerde[Second] with CodegenDispatchFall } } +private[serde] object DatetimeCollation extends CometTypeShim { + def reason(functionName: String): String = + s"$functionName does not support non-UTF8_BINARY collations " + + "(https://github.com/apache/datafusion-comet/issues/4646)" + + def incompatibleReasons(functionName: String): Seq[String] = + if (hasCollationSupport) Seq(reason(functionName)) else Seq.empty + + def hasNonDefaultCollation(expr: Expression): Boolean = + expr.children.exists(c => hasNonDefaultStringCollation(c.dataType)) +} + object CometUnixTimestamp extends CometExpressionSerde[UnixTimestamp] { + private val collationReason = DatetimeCollation.reason("unix_timestamp") + override def getUnsupportedReasons(): Seq[String] = Seq( "Only `DateType`, `TimestampType`, and `TimestampNTZType` inputs are supported.") + override def getIncompatibleReasons(): Seq[String] = + DatetimeCollation.incompatibleReasons("unix_timestamp") + private def isSupportedInputType(expr: UnixTimestamp): Boolean = { expr.children.head.dataType match { case TimestampType | DateType => true @@ -324,7 +342,9 @@ object CometUnixTimestamp extends CometExpressionSerde[UnixTimestamp] { } override def getSupportLevel(expr: UnixTimestamp): SupportLevel = { - if (isSupportedInputType(expr)) { + if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) + } else if (isSupportedInputType(expr)) { Compatible() } else { val inputType = expr.children.head.dataType @@ -420,11 +440,19 @@ object CometConvertTimezone extends CometExpressionSerde[ConvertTimezone] with CodegenDispatchFallback { - override def getSupportLevel(expr: ConvertTimezone): SupportLevel = - Incompatible(Some(UTCTimestampSerde.tzParseIncompatReason)) + private val collationReason = DatetimeCollation.reason("convert_timezone") + + override def getSupportLevel(expr: ConvertTimezone): SupportLevel = { + if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) + } else { + Incompatible(Some(UTCTimestampSerde.tzParseIncompatReason)) + } + } override def getIncompatibleReasons(): Seq[String] = - Seq(UTCTimestampSerde.tzParseIncompatReason) + Seq(UTCTimestampSerde.tzParseIncompatReason) ++ + DatetimeCollation.incompatibleReasons("convert_timezone") override def convert( expr: ConvertTimezone, @@ -446,6 +474,18 @@ object CometNextDay extends CometExpressionSerde[NextDay] { * `dayOfWeek` rather than returning NULL. The resolved flag is passed to native via the * `ScalarFunc.fail_on_error` field. */ + private val collationReason = DatetimeCollation.reason("next_day") + + override def getIncompatibleReasons(): Seq[String] = + DatetimeCollation.incompatibleReasons("next_day") + + override def getSupportLevel(expr: NextDay): SupportLevel = { + if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) + } else { + Compatible() + } + } override def convert(expr: NextDay, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { val childExpr = expr.children.map(exprToProtoInternal(_, inputs, binding)) val optExpr = scalarFunctionExprToProtoWithReturnType( @@ -527,6 +567,8 @@ object CometTruncDate extends CometExpressionSerde[TruncDate] with CodegenDispat val supportedFormats: Seq[String] = Seq("year", "yyyy", "yy", "quarter", "mon", "month", "mm", "week") + private val collationReason = DatetimeCollation.reason("trunc") + private val nonLiteralFormatIncompatReason: String = "Non-literal format strings will throw an exception instead of returning NULL" @@ -534,21 +576,26 @@ object CometTruncDate extends CometExpressionSerde[TruncDate] with CodegenDispat s"Format $fmt is not supported. Only the following formats are supported: " + supportedFormats.mkString(", ") - override def getIncompatibleReasons(): Seq[String] = Seq(nonLiteralFormatIncompatReason) + override def getIncompatibleReasons(): Seq[String] = + Seq(nonLiteralFormatIncompatReason) ++ DatetimeCollation.incompatibleReasons("trunc") override def getUnsupportedReasons(): Seq[String] = Seq( "Only the following formats are supported: " + supportedFormats.mkString(", ")) override def getSupportLevel(expr: TruncDate): SupportLevel = { - expr.format match { - case Literal(fmt: UTF8String, _) => - if (supportedFormats.contains(fmt.toString.toLowerCase(Locale.ROOT))) { - Compatible() - } else { - Unsupported(Some(unsupportedFormatReason(fmt))) - } - case _ => - Incompatible(Some(nonLiteralFormatIncompatReason)) + if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) + } else { + expr.format match { + case Literal(fmt: UTF8String, _) => + if (supportedFormats.contains(fmt.toString.toLowerCase(Locale.ROOT))) { + Compatible() + } else { + Unsupported(Some(unsupportedFormatReason(fmt))) + } + case _ => + Incompatible(Some(nonLiteralFormatIncompatReason)) + } } } @@ -591,6 +638,8 @@ object CometTruncTimestamp "millisecond", "microsecond") + private val collationReason = DatetimeCollation.reason("date_trunc") + private val nonUtcIncompatReason: String = "Produces incorrect results when used with non-UTC timezones. Compatible when timezone is" + " UTC. (https://github.com/apache/datafusion-comet/issues/2649)" @@ -603,27 +652,32 @@ object CometTruncTimestamp supportedFormats.mkString(", ") override def getIncompatibleReasons(): Seq[String] = - Seq(nonUtcIncompatReason, nonLiteralFormatIncompatReason) + Seq(nonUtcIncompatReason, nonLiteralFormatIncompatReason) ++ + DatetimeCollation.incompatibleReasons("date_trunc") override def getUnsupportedReasons(): Seq[String] = Seq( "Only the following formats are supported: " + supportedFormats.mkString(", ")) override def getSupportLevel(expr: TruncTimestamp): SupportLevel = { - val timezone = expr.timeZoneId.getOrElse("UTC") - val isUtc = timezone == "UTC" || timezone == "Etc/UTC" - expr.format match { - case Literal(fmt: UTF8String, _) => - if (supportedFormats.contains(fmt.toString.toLowerCase(Locale.ROOT))) { - if (isUtc) { - Compatible() + if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) + } else { + val timezone = expr.timeZoneId.getOrElse("UTC") + val isUtc = timezone == "UTC" || timezone == "Etc/UTC" + expr.format match { + case Literal(fmt: UTF8String, _) => + if (supportedFormats.contains(fmt.toString.toLowerCase(Locale.ROOT))) { + if (isUtc) { + Compatible() + } else { + Incompatible(Some(nonUtcIncompatReason)) + } } else { - Incompatible(Some(nonUtcIncompatReason)) + Unsupported(Some(unsupportedFormatReason(fmt))) } - } else { - Unsupported(Some(unsupportedFormatReason(fmt))) - } - case _ => - Incompatible(Some(nonLiteralFormatIncompatReason)) + case _ => + Incompatible(Some(nonLiteralFormatIncompatReason)) + } } } @@ -667,10 +721,9 @@ object CometTruncTimestamp * by [[CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED]]. When that flag is disabled the operator * falls back to Spark. */ -object CometDateFormat extends CometExpressionSerde[DateFormatClass] with NativeOptInAvailable { - - override def getIncompatibleReasons(): Seq[String] = - Seq("Non-UTC timezones may produce different results than Spark") +object CometDateFormat + extends CometExpressionSerde[DateFormatClass] + with CodegenDispatchFallback { /** * Mapping from Spark SimpleDateFormat patterns to strftime patterns. Only formats in this map @@ -708,6 +761,12 @@ object CometDateFormat extends CometExpressionSerde[DateFormatClass] with Native // ISO formats "yyyy-MM-dd'T'HH:mm:ss" -> "%Y-%m-%dT%H:%M:%S") + private val collationReason = DatetimeCollation.reason("date_format") + + override def getIncompatibleReasons(): Seq[String] = + Seq("Non-UTC timezones may produce different results than Spark") ++ + DatetimeCollation.incompatibleReasons("date_format") + // Returns true when the format literal is in the native-format whitelist. private def nativeApplicable(expr: DateFormatClass): Boolean = expr.right match { case Literal(fmt: UTF8String, _) => supportedFormats.contains(fmt.toString) @@ -720,15 +779,19 @@ object CometDateFormat extends CometExpressionSerde[DateFormatClass] with Native } override def getSupportLevel(expr: DateFormatClass): SupportLevel = { - // Show the opt-in hint only when: native is applicable, the config is OFF, and native is not - // already running due to UTC timezone. When isUtc is true, native already runs regardless of - // the config, so the hint would be misleading. - val isExprAllowIncompat = CometConf.isExprAllowIncompat(getExprConfigName(expr)) - if (nativeApplicable(expr) && !isUtc(expr) && !isExprAllowIncompat) { - Compatible(nativeOptIn = - Some(NativeOptIn(CometConf.getExprAllowIncompatConfigKey(getExprConfigName(expr))))) + if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) } else { - Compatible() + // Show the opt-in hint only when: native is applicable, the config is OFF, and native is not + // already running due to UTC timezone. When isUtc is true, native already runs regardless of + // the config, so the hint would be misleading. + val isExprAllowIncompat = CometConf.isExprAllowIncompat(getExprConfigName(expr)) + if (nativeApplicable(expr) && !isUtc(expr) && !isExprAllowIncompat) { + Compatible(nativeOptIn = + Some(NativeOptIn(CometConf.getExprAllowIncompatConfigKey(getExprConfigName(expr))))) + } else { + Compatible() + } } } @@ -868,7 +931,23 @@ object CometAddMonths extends CometCodegenDispatch[AddMonths] object CometMonthsBetween extends CometCodegenDispatch[MonthsBetween] -object CometMakeTimestamp extends CometCodegenDispatch[MakeTimestamp] +object CometMakeTimestamp + extends CometCodegenDispatch[MakeTimestamp] + with CodegenDispatchFallback { + + private val collationReason = DatetimeCollation.reason("make_timestamp") + + override def getIncompatibleReasons(): Seq[String] = + DatetimeCollation.incompatibleReasons("make_timestamp") + + override def getSupportLevel(expr: MakeTimestamp): SupportLevel = { + if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) + } else { + Compatible() + } + } +} object CometMicrosToTimestamp extends CometCodegenDispatch[MicrosToTimestamp] @@ -880,6 +959,22 @@ object CometUnixMillis extends CometCodegenDispatch[UnixMillis] object CometUnixMicros extends CometCodegenDispatch[UnixMicros] -object CometToUnixTimestamp extends CometCodegenDispatch[ToUnixTimestamp] +object CometToUnixTimestamp + extends CometCodegenDispatch[ToUnixTimestamp] + with CodegenDispatchFallback { + + private val collationReason = DatetimeCollation.reason("to_unix_timestamp") + + override def getIncompatibleReasons(): Seq[String] = + DatetimeCollation.incompatibleReasons("to_unix_timestamp") + + override def getSupportLevel(expr: ToUnixTimestamp): SupportLevel = { + if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) + } else { + Compatible() + } + } +} object CometGetTimestamp extends CometCodegenDispatch[GetTimestamp] diff --git a/spark/src/main/scala/org/apache/comet/serde/unixtime.scala b/spark/src/main/scala/org/apache/comet/serde/unixtime.scala index 1a3a4ca677..467058ccea 100644 --- a/spark/src/main/scala/org/apache/comet/serde/unixtime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/unixtime.scala @@ -29,12 +29,23 @@ import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithFa // https://github.com/apache/datafusion/issues/16594 object CometFromUnixTime extends CometExpressionSerde[FromUnixTime] with CodegenDispatchFallback { - override def getIncompatibleReasons(): Seq[String] = Seq( + private val collationReason = DatetimeCollation.reason("from_unixtime") + + private val formatReason = "Only supports the default datetime format pattern `yyyy-MM-dd HH:mm:ss`." + " DataFusion's valid timestamp range differs from Spark" + - " (https://github.com/apache/datafusion/issues/16594)") + " (https://github.com/apache/datafusion/issues/16594)" + + override def getIncompatibleReasons(): Seq[String] = + Seq(formatReason) ++ DatetimeCollation.incompatibleReasons("from_unixtime") - override def getSupportLevel(expr: FromUnixTime): SupportLevel = Incompatible(None) + override def getSupportLevel(expr: FromUnixTime): SupportLevel = { + if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) + } else { + Incompatible(Some(formatReason)) + } + } override def convert( expr: FromUnixTime, diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala index 041eb3ce67..97320be9e7 100644 --- a/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala @@ -30,6 +30,9 @@ trait CometTypeShim { @nowarn // Spark 4 feature; stubbed to false in Spark 3.x for compatibility. def hasNonDefaultStringCollation(dt: DataType): Boolean = false + @nowarn // Spark 4 feature; collation does not exist in Spark 3.x. + def hasCollationSupport: Boolean = false + @nowarn // Spark 4 feature; Variant shredding doesn't exist in Spark 3.x. def isVariantStruct(s: StructType): Boolean = false diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala index 17b2738b35..1d4a9f601e 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala @@ -56,4 +56,6 @@ trait CometTypeShim { def isTimeType(dt: DataType): Boolean = dt.getClass.getSimpleName.startsWith("TimeType") + + def hasCollationSupport: Boolean = true } diff --git a/spark/src/test/spark-4.0/org/apache/spark/sql/CometCollationSuite.scala b/spark/src/test/spark-4.0/org/apache/spark/sql/CometCollationSuite.scala index 2c8451fc36..4ac077d240 100644 --- a/spark/src/test/spark-4.0/org/apache/spark/sql/CometCollationSuite.scala +++ b/spark/src/test/spark-4.0/org/apache/spark/sql/CometCollationSuite.scala @@ -247,4 +247,147 @@ class CometCollationSuite extends CometTestBase { "string keys; the collation guard for #4051 must not over-block.") } } + + // ---- datetime expression collation guards (issue #4646) -------------------------------- + // + // Comet's native datetime functions use string arguments (format patterns, timezones, + // day-of-week) as raw bytes, so non-default collations on those arguments must not reach + // the native path silently. Expressions without a codegen-dispatcher fallback (next_day, + // unix_timestamp) fall back to Spark entirely. Expressions with CodegenDispatchFallback + // (trunc, date_trunc, date_format, from_unixtime, make_timestamp, to_unix_timestamp, + // convert_timezone) fall back to Spark when COMET_SCALA_UDF_CODEGEN_ENABLED is false + // or route through Spark codegen inside the Comet pipeline when it is true. + + private def withDatetimeCollationTable(f: => Unit): Unit = { + withParquetTable( + Seq( + ( + "2024-01-01", + "2024-01-01 00:00:00", + "2024-06-15", + "2024-06-15 10:00:00", + "MON", + "yyyy-MM-dd", + "yyyy-MM-dd HH:mm:ss", + "YEAR", + "UTC", + 1718451045L)), + "datetime_collation_tbl")(f) + } + + private def checkDatetimeFallback(query: String, fallbackReason: String): Unit = { + withDatetimeCollationTable { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + checkSparkAnswerAndFallbackReason(query, fallbackReason) + } + } + } + + private def checkDatetimeDispatcher(query: String): Unit = { + withDatetimeCollationTable { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + checkSparkAnswerAndOperator(query) + } + } + } + + test("next_day rejects non-UTF8_BINARY collated dayOfWeek (issue #4646)") { + checkDatetimeFallback( + "SELECT next_day(CAST(_1 AS DATE), _5 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl", + "next_day does not support non-UTF8_BINARY collations") + } + + test("unix_timestamp rejects non-UTF8_BINARY collated format (issue #4646)") { + checkDatetimeFallback( + "SELECT unix_timestamp(CAST(_2 AS TIMESTAMP), _7 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl", + "unix_timestamp does not support non-UTF8_BINARY collations") + } + + test("from_unixtime rejects non-UTF8_BINARY collated format (issue #4646)") { + checkDatetimeFallback( + "SELECT from_unixtime(_10, _6 COLLATE utf8_lcase) FROM datetime_collation_tbl", + "from_unixtime does not support non-UTF8_BINARY collations") + } + + test("make_timestamp rejects non-UTF8_BINARY collated timezone (issue #4646)") { + checkDatetimeFallback( + "SELECT make_timestamp(2024, 6, 15, 10, 30, 45.0, _9 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl", + "make_timestamp does not support non-UTF8_BINARY collations") + } + + test("to_unix_timestamp rejects non-UTF8_BINARY collated format (issue #4646)") { + checkDatetimeFallback( + "SELECT to_unix_timestamp(_3, _6 COLLATE utf8_lcase) FROM datetime_collation_tbl", + "to_unix_timestamp does not support non-UTF8_BINARY collations") + } + + test("convert_timezone rejects non-UTF8_BINARY collated timezone (issue #4646)") { + checkDatetimeFallback( + "SELECT convert_timezone(_9 COLLATE utf8_lcase, 'America/Los_Angeles', " + + "TIMESTAMP_NTZ '2024-01-01 00:00:00') FROM datetime_collation_tbl", + "convert_timezone does not support non-UTF8_BINARY collations") + } + + test("trunc falls back with collated format when codegen is disabled (issue #4646)") { + checkDatetimeFallback( + "SELECT trunc(CAST(_3 AS DATE), _8 COLLATE utf8_lcase) FROM datetime_collation_tbl", + "trunc does not support non-UTF8_BINARY collations") + } + + test("date_trunc falls back with collated format when codegen is disabled (issue #4646)") { + checkDatetimeFallback( + "SELECT date_trunc(_8 COLLATE utf8_lcase, CAST(_4 AS TIMESTAMP)) " + + "FROM datetime_collation_tbl", + "date_trunc does not support non-UTF8_BINARY collations") + } + + test("date_format falls back with collated format when codegen is disabled (issue #4646)") { + checkDatetimeFallback( + "SELECT date_format(CAST(_2 AS TIMESTAMP), _6 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl", + "date_format does not support non-UTF8_BINARY collations") + } + + test("date_format uses native path with collated format when allowIncompatible is enabled") { + withDatetimeCollationTable { + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false", + "spark.comet.expression.DateFormatClass.allowIncompatible" -> "true") { + checkSparkAnswerAndOperator( + "SELECT date_format(CAST(_2 AS TIMESTAMP), 'yyyy-MM-dd' COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl") + } + } + } + + test("trunc routes collated format through codegen dispatcher (issue #4646)") { + checkDatetimeDispatcher( + "SELECT trunc(CAST(_3 AS DATE), _8 COLLATE utf8_lcase) FROM datetime_collation_tbl") + } + + test("date_trunc routes collated format through codegen dispatcher (issue #4646)") { + checkDatetimeDispatcher( + "SELECT date_trunc(_8 COLLATE utf8_lcase, CAST(_4 AS TIMESTAMP)) " + + "FROM datetime_collation_tbl") + } + + test("date_format routes collated format through codegen dispatcher (issue #4646)") { + checkDatetimeDispatcher( + "SELECT date_format(CAST(_2 AS TIMESTAMP), _6 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl") + } + + test("datetime expressions still run with default UTF8_BINARY collation (issue #4646)") { + withDatetimeCollationTable { + checkSparkAnswerAndOperator( + "SELECT next_day(CAST(_1 AS DATE), _5) FROM datetime_collation_tbl") + checkSparkAnswerAndOperator( + "SELECT trunc(CAST(_3 AS DATE), _8) FROM datetime_collation_tbl") + checkSparkAnswerAndOperator( + "SELECT unix_timestamp(CAST(_2 AS TIMESTAMP), _7) FROM datetime_collation_tbl") + } + } } diff --git a/spark/src/test/spark-4.1/org/apache/spark/sql/CometCollationSuite.scala b/spark/src/test/spark-4.1/org/apache/spark/sql/CometCollationSuite.scala index 463e169b66..e1b9713d75 100644 --- a/spark/src/test/spark-4.1/org/apache/spark/sql/CometCollationSuite.scala +++ b/spark/src/test/spark-4.1/org/apache/spark/sql/CometCollationSuite.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql +import org.apache.comet.CometConf + class CometCollationSuite extends CometTestBase { // Queries that group, sort, or shuffle on a non-default collated string must fall back to @@ -66,4 +68,147 @@ class CometCollationSuite extends CometTestBase { checkSparkAnswerAndOperator("SELECT DISTINCT _1 FROM tbl ORDER BY _1") } } + + // ---- datetime expression collation guards (issue #4646) -------------------------------- + // + // Comet's native datetime functions use string arguments (format patterns, timezones, + // day-of-week) as raw bytes, so non-default collations on those arguments must not reach + // the native path silently. Expressions without a codegen-dispatcher fallback (next_day, + // unix_timestamp) fall back to Spark entirely. Expressions with CodegenDispatchFallback + // (trunc, date_trunc, date_format, from_unixtime, make_timestamp, to_unix_timestamp, + // convert_timezone) fall back to Spark when COMET_SCALA_UDF_CODEGEN_ENABLED is false + // or route through Spark codegen inside the Comet pipeline when it is true. + + private def withDatetimeCollationTable(f: => Unit): Unit = { + withParquetTable( + Seq( + ( + "2024-01-01", + "2024-01-01 00:00:00", + "2024-06-15", + "2024-06-15 10:00:00", + "MON", + "yyyy-MM-dd", + "yyyy-MM-dd HH:mm:ss", + "YEAR", + "UTC", + 1718451045L)), + "datetime_collation_tbl")(f) + } + + private def checkDatetimeFallback(query: String, fallbackReason: String): Unit = { + withDatetimeCollationTable { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + checkSparkAnswerAndFallbackReason(query, fallbackReason) + } + } + } + + private def checkDatetimeDispatcher(query: String): Unit = { + withDatetimeCollationTable { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + checkSparkAnswerAndOperator(query) + } + } + } + + test("next_day rejects non-UTF8_BINARY collated dayOfWeek (issue #4646)") { + checkDatetimeFallback( + "SELECT next_day(CAST(_1 AS DATE), _5 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl", + "next_day does not support non-UTF8_BINARY collations") + } + + test("unix_timestamp rejects non-UTF8_BINARY collated format (issue #4646)") { + checkDatetimeFallback( + "SELECT unix_timestamp(CAST(_2 AS TIMESTAMP), _7 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl", + "unix_timestamp does not support non-UTF8_BINARY collations") + } + + test("from_unixtime rejects non-UTF8_BINARY collated format (issue #4646)") { + checkDatetimeFallback( + "SELECT from_unixtime(_10, _6 COLLATE utf8_lcase) FROM datetime_collation_tbl", + "from_unixtime does not support non-UTF8_BINARY collations") + } + + test("make_timestamp rejects non-UTF8_BINARY collated timezone (issue #4646)") { + checkDatetimeFallback( + "SELECT make_timestamp(2024, 6, 15, 10, 30, 45.0, _9 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl", + "make_timestamp does not support non-UTF8_BINARY collations") + } + + test("to_unix_timestamp rejects non-UTF8_BINARY collated format (issue #4646)") { + checkDatetimeFallback( + "SELECT to_unix_timestamp(_3, _6 COLLATE utf8_lcase) FROM datetime_collation_tbl", + "to_unix_timestamp does not support non-UTF8_BINARY collations") + } + + test("convert_timezone rejects non-UTF8_BINARY collated timezone (issue #4646)") { + checkDatetimeFallback( + "SELECT convert_timezone(_9 COLLATE utf8_lcase, 'America/Los_Angeles', " + + "TIMESTAMP_NTZ '2024-01-01 00:00:00') FROM datetime_collation_tbl", + "convert_timezone does not support non-UTF8_BINARY collations") + } + + test("trunc falls back with collated format when codegen is disabled (issue #4646)") { + checkDatetimeFallback( + "SELECT trunc(CAST(_3 AS DATE), _8 COLLATE utf8_lcase) FROM datetime_collation_tbl", + "trunc does not support non-UTF8_BINARY collations") + } + + test("date_trunc falls back with collated format when codegen is disabled (issue #4646)") { + checkDatetimeFallback( + "SELECT date_trunc(_8 COLLATE utf8_lcase, CAST(_4 AS TIMESTAMP)) " + + "FROM datetime_collation_tbl", + "date_trunc does not support non-UTF8_BINARY collations") + } + + test("date_format falls back with collated format when codegen is disabled (issue #4646)") { + checkDatetimeFallback( + "SELECT date_format(CAST(_2 AS TIMESTAMP), _6 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl", + "date_format does not support non-UTF8_BINARY collations") + } + + test("date_format uses native path with collated format when allowIncompatible is enabled") { + withDatetimeCollationTable { + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false", + "spark.comet.expression.DateFormatClass.allowIncompatible" -> "true") { + checkSparkAnswerAndOperator( + "SELECT date_format(CAST(_2 AS TIMESTAMP), 'yyyy-MM-dd' COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl") + } + } + } + + test("trunc routes collated format through codegen dispatcher (issue #4646)") { + checkDatetimeDispatcher( + "SELECT trunc(CAST(_3 AS DATE), _8 COLLATE utf8_lcase) FROM datetime_collation_tbl") + } + + test("date_trunc routes collated format through codegen dispatcher (issue #4646)") { + checkDatetimeDispatcher( + "SELECT date_trunc(_8 COLLATE utf8_lcase, CAST(_4 AS TIMESTAMP)) " + + "FROM datetime_collation_tbl") + } + + test("date_format routes collated format through codegen dispatcher (issue #4646)") { + checkDatetimeDispatcher( + "SELECT date_format(CAST(_2 AS TIMESTAMP), _6 COLLATE utf8_lcase) " + + "FROM datetime_collation_tbl") + } + + test("datetime expressions still run with default UTF8_BINARY collation (issue #4646)") { + withDatetimeCollationTable { + checkSparkAnswerAndOperator( + "SELECT next_day(CAST(_1 AS DATE), _5) FROM datetime_collation_tbl") + checkSparkAnswerAndOperator( + "SELECT trunc(CAST(_3 AS DATE), _8) FROM datetime_collation_tbl") + checkSparkAnswerAndOperator( + "SELECT unix_timestamp(CAST(_2 AS TIMESTAMP), _7) FROM datetime_collation_tbl") + } + } }