refactor: centralize data type support predicates

.github/workflows/pr_build_linux.yml

@@ -546,6 +546,7 @@ jobs:
               org.apache.comet.CometUuidExpressionSuite
               org.apache.comet.serde.CometScalarFunctionSuite
               org.apache.comet.serde.CometLiteralSuite
+              org.apache.comet.serde.QueryPlanSerdeSuite
               org.apache.comet.CometFallbackInvarianceSuite
       fail-fast: false
     name: ${{ matrix.profile.name }} [${{ matrix.suite.name }}]

.github/workflows/pr_build_macos.yml

@@ -252,6 +252,7 @@ jobs:
               org.apache.comet.CometUuidExpressionSuite
               org.apache.comet.serde.CometScalarFunctionSuite
               org.apache.comet.serde.CometLiteralSuite
+              org.apache.comet.serde.QueryPlanSerdeSuite
               org.apache.comet.CometFallbackInvarianceSuite
 
       fail-fast: false

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -557,21 +557,70 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     builder.build()
   }
 
-  def supportedDataType(dt: DataType, allowComplex: Boolean = false): Boolean = dt match {
-    case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
-        _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType |
-        _: DecimalType | _: DateType | _: BooleanType | _: NullType | CalendarIntervalType =>
-      true
-    case dt if isTimeType(dt) =>
-      true
-    case s: StructType if allowComplex =>
-      s.fields.nonEmpty && s.fields.map(_.dataType).forall(supportedDataType(_, allowComplex))
-    case a: ArrayType if allowComplex =>
-      supportedDataType(a.elementType, allowComplex)
-    case m: MapType if allowComplex =>
-      supportedDataType(m.keyType, allowComplex) && supportedDataType(m.valueType, allowComplex)
-    case _ =>
-      false
+  /**
+   * Returns whether `dt` is supported at a caller's data-type boundary.
+   *
+   * The defaults preserve expression-serde behavior: primitive types, `CalendarIntervalType`,
+   * `TimeType`, and all `StringType` variants are accepted, while complex and ANSI interval types
+   * are rejected. Sinks and native shuffle enable complex and ANSI interval types because their
+   * Arrow IPC paths support them. Local scans additionally reject `TimeType` and non-default
+   * strings, while JVM columnar shuffle rejects ANSI intervals, calendar intervals, and duplicate
+   * struct field names because its unsafe-row-to-Arrow path cannot handle them.
+   *
+   * Note that the option polarity is mixed: `allowComplex` and `allowIntervals` are restrictive
+   * by default; the other four options are permissive by default.
+   *
+   * @param dt
+   *   data type to check
+   * @param allowComplex
+   *   recursively allow non-empty structs, arrays, and maps
+   * @param allowIntervals
+   *   allow year-month and day-time interval types
+   * @param allowCalendarInterval
+   *   allow calendar interval types
+   * @param allowTimeType
+   *   allow Spark `TimeType`
+   * @param allowAnyStringType
+   *   allow non-default `StringType` variants such as collated strings; when false, only the
+   *   default `StringType` is accepted
+   * @param allowDuplicateStructFieldNames
+   *   allow duplicate field names in nested structs
+   */
+  def supportedDataType(
+      dt: DataType,
+      allowComplex: Boolean = false,
+      allowIntervals: Boolean = false,
+      allowCalendarInterval: Boolean = true,
+      allowTimeType: Boolean = true,
+      allowAnyStringType: Boolean = true,
+      allowDuplicateStructFieldNames: Boolean = true): Boolean = {
+    def supported(dt: DataType): Boolean = dt match {
+      case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
+          _: DoubleType | _: BinaryType | _: TimestampType | _: TimestampNTZType |
+          _: DecimalType | _: DateType | _: BooleanType | _: NullType =>
+        true
+      case CalendarIntervalType if allowCalendarInterval =>
+        true
+      case st: StringType if allowAnyStringType || st == StringType =>
+        true
+      case _: YearMonthIntervalType | _: DayTimeIntervalType if allowIntervals =>
+        true
+      case dt if allowTimeType && isTimeType(dt) =>
+        true
+      case s: StructType if allowComplex =>
+        s.fields.nonEmpty &&
+        (allowDuplicateStructFieldNames ||
+          s.fields.map(_.name).distinct.length == s.fields.length) &&
+        s.fields.forall(f => supported(f.dataType))
+      case a: ArrayType if allowComplex =>
+        supported(a.elementType)
+      case m: MapType if allowComplex =>
+        supported(m.keyType) && supported(m.valueType)
+      case _ =>
+        false
+    }
+
+    supported(dt)
   }
 
   /**

spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala

@@ -26,7 +26,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec
 import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
-import org.apache.spark.sql.types.{ArrayType, DataType, DayTimeIntervalType, MapType, StructType, YearMonthIntervalType}
+import org.apache.spark.sql.types.DataType
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
@@ -43,16 +43,6 @@ abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] {
 
   override def enabledConfig: Option[ConfigEntry[Boolean]] = None
 
-  protected final def supportedSinkDataType(dt: DataType): Boolean = dt match {
-    case _: YearMonthIntervalType | _: DayTimeIntervalType => true
-    case StructType(fields) =>
-      fields.nonEmpty && fields.forall(f => supportedSinkDataType(f.dataType))
-    case ArrayType(elementType, _) => supportedSinkDataType(elementType)
-    case MapType(keyType, valueType, _) =>
-      supportedSinkDataType(keyType) && supportedSinkDataType(valueType)
-    case _ => supportedDataType(dt)
-  }
-
   /**
    * The data type to declare for a scan output field. Overridden by sinks whose source carries
    * non-null nested child fields that must be widened to match the planned kernel output types
@@ -64,7 +54,8 @@ abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] {
       op: T,
       builder: Operator.Builder,
       childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = {
-    val supportedTypes = op.output.forall(a => supportedSinkDataType(a.dataType))
+    val supportedTypes = op.output.forall(a =>
+      supportedDataType(a.dataType, allowComplex = true, allowIntervals = true))
 
     if (!supportedTypes) {
       withFallbackReason(op, "Unsupported data type")
@@ -132,7 +123,8 @@ object CometExchangeSink extends CometSink[SparkPlan] {
   private def convertToShuffleScan(
       op: SparkPlan,
       builder: Operator.Builder): Option[OperatorOuterClass.Operator] = {
-    val supportedTypes = op.output.forall(a => supportedSinkDataType(a.dataType))
+    val supportedTypes = op.output.forall(a =>
+      supportedDataType(a.dataType, allowComplex = true, allowIntervals = true))
 
     if (!supportedTypes) {
       withFallbackReason(op, "Unsupported data type for shuffle direct read")

spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTableScanExec.scala

@@ -31,13 +31,14 @@ import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNative
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution.{LeafExecNode, LocalTableScanExec}
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
-import org.apache.spark.sql.types.{DataType, DayTimeIntervalType, NullType, StructType, YearMonthIntervalType}
+import org.apache.spark.sql.types.{DataType, StructType}
 
 import com.google.common.base.Objects
 
 import org.apache.comet.{CometConf, ConfigEntry, DataTypeSupport}
 import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
 import org.apache.comet.serde.OperatorOuterClass.Operator
+import org.apache.comet.serde.QueryPlanSerde.supportedDataType
 import org.apache.comet.serde.operator.CometSink
 
 case class CometLocalTableScanExec(
@@ -131,15 +132,24 @@ object CometLocalTableScanExec extends CometSink[LocalTableScanExec] with DataTy
   // downstream expression serdes (issue #4789).
   override protected def scanFieldType(dt: DataType): DataType = dt.asNullable
 
-  // ArrowWriter (used by RowArrowReader) handles NullType via Utils.toArrowType + NullWriter;
-  // other types off DataTypeSupport's allow list (TimeType, intervals, ...) have no ArrowWriter
-  // coverage and must fall back to Spark.
+  // RowArrowReader handles NullType and intervals, but not TimeType. Non-default string collations
+  // remain unsupported here, matching DataTypeSupport's existing local-scan boundary.
   override def isTypeSupported(
       dt: DataType,
       name: String,
-      fallbackReasons: ListBuffer[String]): Boolean = dt match {
-    case _: NullType | _: YearMonthIntervalType | _: DayTimeIntervalType => true
-    case _ => super.isTypeSupported(dt, name, fallbackReasons)
+      fallbackReasons: ListBuffer[String]): Boolean = {
+    val supported = supportedDataType(
+      dt,
+      allowComplex = true,
+      allowIntervals = true,
+      allowTimeType = false,
+      allowAnyStringType = false,
+      // Java Arrow keys struct children by name; declining here hands the struct to
+      // DataTypeSupport, which records the duplicate-field-name fallback reason.
+      allowDuplicateStructFieldNames = false)
+    // DataTypeSupport accepts no type that supportedDataType rejects here, so the super call
+    // cannot widen the accepted set; it only records the fallback reason for rejected types.
+    if (supported) true else super.isTypeSupported(dt, name, fallbackReasons)
   }
 
   override def convert(

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala

@@ -41,7 +41,7 @@ import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec
 import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin}
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter}
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.util.MutablePair
 import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, RecordComparator}
@@ -462,33 +462,6 @@ object CometShuffleExchangeExec
         false
     }
 
-    /**
-     * Determine which data types are supported as data columns in native shuffle.
-     *
-     * Native shuffle relies on the Arrow IPC writer to serialize batches to disk, so it should
-     * support all types that Comet supports.
-     */
-    def supportedSerializableDataType(dt: DataType): Boolean = dt match {
-      case _: BooleanType | _: ByteType | _: ShortType | _: IntegerType | _: LongType |
-          _: FloatType | _: DoubleType | _: StringType | _: BinaryType | _: TimestampType |
-          _: TimestampNTZType | _: DecimalType | _: DateType | _: NullType |
-          _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType =>
-        true
-      case dt if isTimeType(dt) =>
-        true
-      case StructType(fields) =>
-        fields.nonEmpty && fields.forall(f => supportedSerializableDataType(f.dataType)) &&
-        // Java Arrow keys struct children by name, so the FFI import of a decoded batch
-        // fails on duplicate field names
-        fields.map(f => f.name).distinct.length == fields.length
-      case ArrayType(elementType, _) =>
-        supportedSerializableDataType(elementType)
-      case MapType(keyType, valueType, _) =>
-        supportedSerializableDataType(keyType) && supportedSerializableDataType(valueType)
-      case _ =>
-        false
-    }
-
     val reasons = scala.collection.mutable.ListBuffer.empty[String]
 
     if (!isCometNativeShuffleMode(s.conf)) {
@@ -499,7 +472,13 @@ object CometShuffleExchangeExec
     val inputs = s.child.output
 
     for (input <- inputs) {
-      if (!supportedSerializableDataType(input.dataType)) {
+      if (!QueryPlanSerde.supportedDataType(
+          input.dataType,
+          allowComplex = true,
+          allowIntervals = true,
+          // Java Arrow keys struct children by name, so the FFI import of a decoded batch
+          // fails on duplicate field names.
+          allowDuplicateStructFieldNames = false)) {
         reasons += s"unsupported shuffle data type ${input.dataType} for input $input"
         return reasons.toSeq
       }
@@ -591,31 +570,6 @@ object CometShuffleExchangeExec
    */
   private def columnarShuffleFailureReasons(s: ShuffleExchangeExec): Seq[String] = {
 
-    /**
-     * Determine which data types are supported as data columns in columnar shuffle.
-     *
-     * Comet columnar shuffle used native code to convert Spark unsafe rows to Arrow batches, see
-     * shuffle/row.rs
-     */
-    def supportedSerializableDataType(dt: DataType): Boolean = dt match {
-      case _: BooleanType | _: ByteType | _: ShortType | _: IntegerType | _: LongType |
-          _: FloatType | _: DoubleType | _: StringType | _: BinaryType | _: TimestampType |
-          _: TimestampNTZType | _: DecimalType | _: DateType | _: NullType =>
-        true
-      case dt if isTimeType(dt) =>
-        true
-      case StructType(fields) =>
-        fields.nonEmpty && fields.forall(f => supportedSerializableDataType(f.dataType)) &&
-        // Java Arrow stream reader cannot work on duplicate field name
-        fields.map(f => f.name).distinct.length == fields.length
-      case ArrayType(elementType, _) =>
-        supportedSerializableDataType(elementType)
-      case MapType(keyType, valueType, _) =>
-        supportedSerializableDataType(keyType) && supportedSerializableDataType(valueType)
-      case _ =>
-        false
-    }
-
     val reasons = scala.collection.mutable.ListBuffer.empty[String]
 
     if (!isCometJVMShuffleMode(s.conf)) {
@@ -636,7 +590,14 @@ object CometShuffleExchangeExec
     val inputs = s.child.output
 
     for (input <- inputs) {
-      if (!supportedSerializableDataType(input.dataType)) {
+      if (!QueryPlanSerde.supportedDataType(
+          input.dataType,
+          allowComplex = true,
+          // The native row-to-Arrow converter (spark_unsafe/row.rs) has no CalendarInterval
+          // support, so calendar intervals must fall back to Spark shuffle.
+          allowCalendarInterval = false,
+          // Java Arrow stream reader cannot work on duplicate field names.
+          allowDuplicateStructFieldNames = false)) {
         reasons += s"unsupported shuffle data type ${input.dataType} for input $input"
         return reasons.toSeq
       }

spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala

@@ -31,7 +31,7 @@ import org.apache.spark.{Partitioner, SparkConf}
 import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
 import org.apache.spark.sql.comet.execution.shuffle.{CometShuffleDependency, CometShuffleExchangeExec, CometShuffleManager}
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEShuffleReadExec, ShuffleQueryStageExec}
-import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.execution.joins.SortMergeJoinExec
 import org.apache.spark.sql.functions.col
 import org.apache.spark.sql.internal.SQLConf
@@ -75,6 +75,29 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
       """.stripMargin).select($"r.*")
 
     checkSparkAnswer(df)
+    checkCometExchange(df, 0, false)
+  }
+
+  test("Fallback to Spark when shuffling CalendarIntervalType data") {
+    val df = spark
+      .sql("select id, make_interval(1,2,3,4,5,6,7) as i from range(100)")
+      .repartition(4)
+
+    assert(df.collect().length == 100)
+
+    val plan = df.queryExecution.executedPlan
+    assert(
+      find(plan) {
+        case _: ShuffleExchangeExec => true
+        case _ => false
+      }.nonEmpty,
+      plan)
+    assert(
+      find(plan) {
+        case _: CometShuffleExchangeExec => true
+        case _ => false
+      }.isEmpty,
+      plan)
   }
 
   test("Unsupported types for SinglePartition should fallback to Spark") {