diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala index 329c91f49ea..1a82d789645 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala @@ -26,28 +26,35 @@ import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.serde.QueryPlanSerde.{serializeDataType, supportedDataType} object CometScalarSubquery extends CometExpressionSerde[ScalarSubquery] { + + override def getUnsupportedReasons(): Seq[String] = Seq( + "Not all data types are supported for scalar subquery results") + + override def getSupportLevel(expr: ScalarSubquery): SupportLevel = + if (supportedDataType(expr.dataType)) { + Compatible() + } else { + Unsupported(Some(s"Unsupported data type: ${expr.dataType}")) + } + override def convert( expr: ScalarSubquery, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - if (supportedDataType(expr.dataType)) { - val dataType = serializeDataType(expr.dataType) - if (dataType.isEmpty) { - withFallbackReason( - expr, - s"Failed to serialize datatype ${expr.dataType} for scalar subquery") - return None - } - - val builder = ExprOuterClass.Subquery - .newBuilder() - .setId(expr.exprId.id) - .setDatatype(dataType.get) - Some(ExprOuterClass.Expr.newBuilder().setSubquery(builder).build()) - } else { - withFallbackReason(expr, s"Unsupported data type: ${expr.dataType}") - None + // getSupportLevel has already screened the data type with `supportedDataType`. That is a + // different predicate from `serializeDataType`, which can still decline, so keep this check. + val dataType = serializeDataType(expr.dataType) + if (dataType.isEmpty) { + withFallbackReason( + expr, + s"Failed to serialize datatype ${expr.dataType} for scalar subquery") + return None } + val builder = ExprOuterClass.Subquery + .newBuilder() + .setId(expr.exprId.id) + .setDatatype(dataType.get) + Some(ExprOuterClass.Expr.newBuilder().setSubquery(builder).build()) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/collectionOperations.scala b/spark/src/main/scala/org/apache/comet/serde/collectionOperations.scala index 79714696f39..58ef3339365 100644 --- a/spark/src/main/scala/org/apache/comet/serde/collectionOperations.scala +++ b/spark/src/main/scala/org/apache/comet/serde/collectionOperations.scala @@ -22,7 +22,6 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, Reverse, Shuffle} import org.apache.spark.sql.types.ArrayType -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal import org.apache.comet.shims.CometTypeShim @@ -63,28 +62,35 @@ object CometReverse object CometShuffle extends CometExpressionSerde[Shuffle] with ArraysBase { + private val unresolvedSeedReason = "shuffle requires a resolved random seed" + + override def getUnsupportedReasons(): Seq[String] = Seq(unresolvedSeedReason) + // Comet reproduces Spark's random permutation exactly: the resolved seed is combined with the // partition index and drives the same MersenneTwister-based inside-out Fisher-Yates shuffle as // org.apache.spark.sql.catalyst.util.RandomIndicesGenerator, so results match Spark bit for bit. - override def getSupportLevel(expr: Shuffle): SupportLevel = childTypesSupportLevel(expr) + override def getSupportLevel(expr: Shuffle): SupportLevel = { + // In a resolved plan `randomSeed` is always defined (resolution requires it). Guard anyway, + // here rather than in `convert`, so the decline goes through the normal support-level path. + if (expr.randomSeed.isEmpty) { + Unsupported(Some(unresolvedSeedReason)) + } else { + childTypesSupportLevel(expr) + } + } override def convert(expr: Shuffle, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { - // In a resolved plan `randomSeed` is always defined (resolution requires it). Guard anyway. - expr.randomSeed match { - case Some(seed) => - exprToProtoInternal(expr.child, inputs, binding).map { child => - ExprOuterClass.Expr + // getSupportLevel has already verified the seed is resolved. + val seed = expr.randomSeed.get + exprToProtoInternal(expr.child, inputs, binding).map { child => + ExprOuterClass.Expr + .newBuilder() + .setShuffle( + ExprOuterClass.Shuffle .newBuilder() - .setShuffle( - ExprOuterClass.Shuffle - .newBuilder() - .setChild(child) - .setSeed(seed)) - .build() - } - case None => - withFallbackReason(expr, "shuffle requires a resolved random seed") - None + .setChild(child) + .setSeed(seed)) + .build() } } } 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 9ea1f6e27fe..940ee597b6e 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -27,7 +27,6 @@ import org.apache.spark.sql.types.{CalendarIntervalType, DataType, DateType, Dou import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.expressions.{CometCast, CometEvalMode} import org.apache.comet.serde.CometGetDateField.CometGetDateField import org.apache.comet.serde.ExprOuterClass.Expr @@ -310,13 +309,18 @@ object CometUnixTimestamp extends CometExpressionSerde[UnixTimestamp] { } override def getSupportLevel(expr: UnixTimestamp): SupportLevel = { - if (DatetimeCollation.hasNonDefaultCollation(expr)) { - Incompatible(Some(collationReason)) - } else if (isSupportedInputType(expr)) { - Compatible() - } else { + // The input type is screened ahead of the collation check on purpose. A non-date/timestamp + // input has no native path at all, so it must report `Unsupported` rather than + // `Incompatible`: the latter is waved straight through to `convert` when + // `spark.comet.expression.UnixTimestamp.allowIncompatible=true`, and the native kernel then + // raises an execution error on the string child instead of falling back to Spark. + if (!isSupportedInputType(expr)) { val inputType = expr.children.head.dataType Unsupported(Some(s"unix_timestamp does not support input type: $inputType")) + } else if (DatetimeCollation.hasNonDefaultCollation(expr)) { + Incompatible(Some(collationReason)) + } else { + Compatible() } } @@ -324,12 +328,7 @@ object CometUnixTimestamp extends CometExpressionSerde[UnixTimestamp] { expr: UnixTimestamp, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - if (!isSupportedInputType(expr)) { - val inputType = expr.children.head.dataType - withFallbackReason(expr, s"unix_timestamp does not support input type: $inputType") - return None - } - + // getSupportLevel reports an unsupported input type before reaching here, so no re-check. val childExpr = exprToProtoInternal(expr.children.head, inputs, binding) if (childExpr.isDefined) { diff --git a/spark/src/main/scala/org/apache/comet/serde/nondetermenistic.scala b/spark/src/main/scala/org/apache/comet/serde/nondetermenistic.scala index a88ce92c9de..b777e02d713 100644 --- a/spark/src/main/scala/org/apache/comet/serde/nondetermenistic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/nondetermenistic.scala @@ -21,8 +21,6 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Literal, MonotonicallyIncreasingID, Rand, Randn, SparkPartitionID, Uuid} -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason - object CometSparkPartitionId extends CometExpressionSerde[SparkPartitionID] { override def convert( expr: SparkPartitionID, @@ -72,6 +70,15 @@ sealed abstract class CometRandCommonSerde[T <: Expression] extends CometExpress object CometUuid extends CometExpressionSerde[Uuid] { + private val unresolvedSeedReason = "uuid requires a resolved random seed" + + override def getUnsupportedReasons(): Seq[String] = Seq(unresolvedSeedReason) + + // In a resolved plan `randomSeed` is always defined (resolution requires it). Guard anyway, + // here rather than in `convert`, so the decline goes through the normal support-level path. + override def getSupportLevel(expr: Uuid): SupportLevel = + if (expr.randomSeed.isEmpty) Unsupported(Some(unresolvedSeedReason)) else Compatible() + // Comet reproduces Spark's UUIDs exactly: the resolved seed is combined with the partition index // and seeds the same Commons Math3 MersenneTwister that drives // org.apache.spark.sql.catalyst.util.RandomUUIDGenerator, so results match Spark bit for bit. @@ -79,18 +86,12 @@ object CometUuid extends CometExpressionSerde[Uuid] { expr: Uuid, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - // In a resolved plan `randomSeed` is always defined (resolution requires it). Guard anyway. - expr.randomSeed match { - case Some(seed) => - Some( - ExprOuterClass.Expr - .newBuilder() - .setUuid(ExprOuterClass.Uuid.newBuilder().setSeed(seed)) - .build()) - case None => - withFallbackReason(expr, "uuid requires a resolved random seed") - None - } + // getSupportLevel has already verified the seed is resolved. + Some( + ExprOuterClass.Expr + .newBuilder() + .setUuid(ExprOuterClass.Uuid.newBuilder().setSeed(expr.randomSeed.get)) + .build()) } } 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 08285c2cf12..861799fcb9e 100644 --- a/spark/src/main/scala/org/apache/comet/serde/unixtime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/unixtime.scala @@ -22,7 +22,6 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, FromUnixTime, Literal} import org.apache.spark.sql.catalyst.util.TimestampFormatter -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProto} // TODO: DataFusion supports only -8334601211038 <= sec <= 8210266876799 @@ -65,6 +64,8 @@ object CometFromUnixTime extends CometExpressionSerde[FromUnixTime] with Codegen expr: FromUnixTime, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { + // getSupportLevel reports a non-default format as `Unsupported`, so it is routed to the + // codegen dispatcher (or falls back) before reaching here; only the default pattern arrives. val secExpr = exprToProtoInternal(expr.sec, inputs, binding) // TODO: DataFusion toChar does not support Spark datetime pattern format // https://github.com/apache/datafusion/issues/16577 @@ -73,10 +74,7 @@ object CometFromUnixTime extends CometExpressionSerde[FromUnixTime] with Codegen val formatExpr = exprToProtoInternal(Literal("%Y-%m-%d %H:%M:%S"), inputs, binding) val timeZone = exprToProtoInternal(Literal(expr.timeZoneId.orNull), inputs, binding) - if (expr.format != Literal(TimestampFormatter.defaultPattern())) { - withFallbackReason(expr, formatReason) - None - } else if (secExpr.isDefined && formatExpr.isDefined) { + if (secExpr.isDefined && formatExpr.isDefined) { val timestampExpr = scalarFunctionExprToProto("from_unixtime", Seq(secExpr, timeZone): _*) val optExpr = scalarFunctionExprToProto("to_char", Seq(timestampExpr, formatExpr): _*) diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index 1ce65f0660c..fce36496aa4 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -31,6 +31,7 @@ import org.apache.spark.sql.functions.col import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus import org.apache.comet.serde.{CometDateFormat, CometTruncDate, CometTruncTimestamp} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} @@ -405,6 +406,28 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH } } + test("unix_timestamp - string input falls back even when a collated format opts into native") { + assume(isSpark40Plus, "string collation requires Spark 4.0+") + withTempView("string_tbl") { + val schema = StructType(Seq(StructField("ts_str", DataTypes.StringType, true))) + val data = Seq(Row("2020-01-01 00:00:00"), Row("2021-06-15 12:30:45"), Row(null)) + spark + .createDataFrame(spark.sparkContext.parallelize(data), schema) + .createOrReplaceTempView("string_tbl") + + // A collated format argument makes the expression `Incompatible`, which + // `allowIncompatible=true` waves straight through to `convert`. The input type has to be + // rejected ahead of that opt-in: the native kernel accepts only date/timestamp input, so + // serializing a string child raises an execution error instead of falling back to Spark. + withSQLConf(CometConf.getExprAllowIncompatConfigKey("UnixTimestamp") -> "true") { + checkSparkAnswerAndFallbackReason( + "SELECT ts_str, unix_timestamp(ts_str, 'yyyy-MM-dd HH:mm:ss' COLLATE UTF8_LCASE) " + + "from string_tbl order by ts_str", + "unix_timestamp does not support input type: StringType") + } + } + } + // Time-window grouping (PreciseTimestampConversion / KnownNullable) is covered by the SQL file // test spark/src/test/resources/sql-tests/expressions/datetime/window.sql.