diff --git a/spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala b/spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala index 419a5be0c1..6eeba5e615 100644 --- a/spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala +++ b/spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala @@ -22,7 +22,6 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, KnownFloatingPointNormalized, KnownNullable} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithFallbackReason, serializeDataType} object CometKnownFloatingPointNormalized @@ -33,7 +32,13 @@ object CometKnownFloatingPointNormalized override def getSupportLevel(expr: KnownFloatingPointNormalized): SupportLevel = { expr.child match { - case _: NormalizeNaNAndZero => Compatible() + case normalized: NormalizeNaNAndZero => + val wrapped = normalized.child + if (serializeDataType(wrapped.dataType).isDefined) { + Compatible() + } else { + Unsupported(Some(s"Unsupported datatype ${wrapped.dataType}")) + } case _ => Unsupported( Some( @@ -48,17 +53,13 @@ object CometKnownFloatingPointNormalized val wrapped = expr.child.asInstanceOf[NormalizeNaNAndZero].child - val dataType = serializeDataType(wrapped.dataType) - if (dataType.isEmpty) { - withFallbackReason(wrapped, s"Unsupported datatype ${wrapped.dataType}") - return None - } + val dataType = serializeDataType(wrapped.dataType).get val ex = exprToProtoInternal(wrapped, inputs, binding) val optExpr = ex.map { child => val builder = ExprOuterClass.NormalizeNaNAndZero .newBuilder() .setChild(child) - .setDatatype(dataType.get) + .setDatatype(dataType) ExprOuterClass.Expr.newBuilder().setNormalizeNanAndZero(builder).build() } optExprWithFallbackReason(optExpr, expr, wrapped) diff --git a/spark/src/main/scala/org/apache/comet/serde/namedExpressions.scala b/spark/src/main/scala/org/apache/comet/serde/namedExpressions.scala index b778f2ea87..fb910ec3ac 100644 --- a/spark/src/main/scala/org/apache/comet/serde/namedExpressions.scala +++ b/spark/src/main/scala/org/apache/comet/serde/namedExpressions.scala @@ -38,53 +38,55 @@ object CometAlias extends CometExpressionSerde[Alias] { } object CometAttributeReference extends CometExpressionSerde[AttributeReference] { + override def getSupportLevel(attr: AttributeReference): SupportLevel = { + if (serializeDataType(attr.dataType).isDefined) { + Compatible() + } else { + Unsupported(Some(s"unsupported datatype: ${attr.dataType}")) + } + } + override def convert( attr: AttributeReference, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val dataType = serializeDataType(attr.dataType) + val dataType = serializeDataType(attr.dataType).get - if (dataType.isDefined) { - if (binding) { - // Spark may produce unresolvable attributes in some cases, - // for example https://github.com/apache/datafusion-comet/issues/925. - // So, we allow the binding to fail. - val boundRef: Any = BindReferences - .bindReference(attr, inputs, allowFailures = true) + if (binding) { + // Spark may produce unresolvable attributes in some cases, + // for example https://github.com/apache/datafusion-comet/issues/925. + // So, we allow the binding to fail. + val boundRef: Any = BindReferences + .bindReference(attr, inputs, allowFailures = true) - if (boundRef.isInstanceOf[AttributeReference]) { - withFallbackReason(attr, s"cannot resolve $attr among ${inputs.mkString(", ")}") - return None - } + if (boundRef.isInstanceOf[AttributeReference]) { + withFallbackReason(attr, s"cannot resolve $attr among ${inputs.mkString(", ")}") + return None + } - val boundExpr = ExprOuterClass.BoundReference - .newBuilder() - .setIndex(boundRef.asInstanceOf[BoundReference].ordinal) - .setDatatype(dataType.get) - .build() + val boundExpr = ExprOuterClass.BoundReference + .newBuilder() + .setIndex(boundRef.asInstanceOf[BoundReference].ordinal) + .setDatatype(dataType) + .build() - Some( - ExprOuterClass.Expr - .newBuilder() - .setBound(boundExpr) - .build()) - } else { - val unboundRef = ExprOuterClass.UnboundReference + Some( + ExprOuterClass.Expr .newBuilder() - .setName(attr.name) - .setDatatype(dataType.get) - .build() - - Some( - ExprOuterClass.Expr - .newBuilder() - .setUnbound(unboundRef) - .build()) - } + .setBound(boundExpr) + .build()) } else { - withFallbackReason(attr, s"unsupported datatype: ${attr.dataType}") - None - } + val unboundRef = ExprOuterClass.UnboundReference + .newBuilder() + .setName(attr.name) + .setDatatype(dataType) + .build() + Some( + ExprOuterClass.Expr + .newBuilder() + .setUnbound(unboundRef) + .build()) + } } } diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index d6955fb78b..f3c442ad29 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -25,8 +25,8 @@ import scala.util.Random import org.apache.hadoop.fs.Path import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row} -import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, FromUnixTime, Literal, StructsToJson, TruncDate, TruncTimestamp} -import org.apache.spark.sql.catalyst.optimizer.SimplifyExtractValueOps +import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, Cast, FromUnixTime, KnownFloatingPointNormalized, Literal, StructsToJson, TruncDate, TruncTimestamp} +import org.apache.spark.sql.catalyst.optimizer.{NormalizeNaNAndZero, SimplifyExtractValueOps} import org.apache.spark.sql.comet.CometProjectExec import org.apache.spark.sql.execution.{ProjectExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper @@ -36,6 +36,7 @@ import org.apache.spark.sql.internal.SQLConf.SESSION_LOCAL_TIMEZONE import org.apache.spark.sql.types._ import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus, isSpark42Plus} +import org.apache.comet.serde.{CometAttributeReference, CometKnownFloatingPointNormalized, Unsupported} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -1050,6 +1051,18 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("misc scalar serdes report static unsupported cases via getSupportLevel") { + val intervalAttr = AttributeReference("interval_attr", CalendarIntervalType)() + + assert( + CometAttributeReference.getSupportLevel(intervalAttr) == + Unsupported(Some("unsupported datatype: CalendarIntervalType"))) + assert( + CometKnownFloatingPointNormalized.getSupportLevel( + KnownFloatingPointNormalized(NormalizeNaNAndZero(intervalAttr))) == + Unsupported(Some("Unsupported datatype CalendarIntervalType"))) + } + test("rlike with non-scalar pattern runs via codegen dispatcher") { val table = "rlike_non_scalar" withTable(table) {