refactor: move AttributeReference and KnownFloatingPointNormalized expression support checks to getSupportLevel

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)

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())
+    }
   }
 }

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) {