You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/25 16:25:45 UTC
(datafusion-comet) 01/01: refactor: centralize data type support predicates (#5025)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5025-aaa33eb203422fa2c1015740e5d889825dc48985
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 4ab4dfa0a8348f0942a06d1f6d27494a42c155d4
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Fri Sep 25 14:37:52 2026 +0000
refactor: centralize data type support predicates (#5025)
* feat: support interval codegen dispatch
* add struct and map with nested interval type
* Centralize data type support predicates
* fix: keep columnar shuffle rejecting CalendarIntervalType
The refactor placed CalendarIntervalType in the unconditional base arm of
supportedDataType, which made columnar shuffle newly accept it. The native
unsafe-row-to-Arrow converter (spark_unsafe/row.rs) has no interval support,
so this would panic at runtime instead of falling back to Spark.
Add an allowCalendarInterval capability (default true, matching every other
pre-refactor call site) and pass false at the columnar shuffle boundary.
Also document that the local-scan super call only records fallback reasons.
Co-Authored-By: Claude Fable 5 <no...@anthropic.com>
* test: cover data type boundary policies
* ci: register QueryPlanSerdeSuite
---------
Co-authored-by: Claude Fable 5 <no...@anthropic.com>
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
.../org/apache/comet/serde/QueryPlanSerde.scala | 79 +++++++++++++---
.../apache/comet/serde/operator/CometSink.scala | 18 +---
.../spark/sql/comet/CometLocalTableScanExec.scala | 24 +++--
.../shuffle/CometShuffleExchangeExec.scala | 73 ++++-----------
.../comet/exec/CometColumnarShuffleSuite.scala | 25 ++++-
.../apache/comet/serde/QueryPlanSerdeSuite.scala | 104 +++++++++++++++++++++
8 files changed, 233 insertions(+), 92 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index d711c90170..2317106337 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -622,6 +622,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 }}]
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 6bd5d8742f..85ff13dfca 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -270,6 +270,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
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 3944b13f49..1d87294c20 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -39,6 +39,7 @@ import org.apache.spark.sql.types._
import org.apache.comet.CometConf
import org.apache.comet.CometExplainInfo
import org.apache.comet.CometSparkSessionExtensions.{appendTagValues, withFallbackReason, withFallbackReasons, withInfo, withNativeExpr}
+import org.apache.comet.DataTypeSupport
import org.apache.comet.expressions._
import org.apache.comet.parquet.CometParquetUtils
import org.apache.comet.serde.ExprOuterClass.{AggExpr, Expr, ScalarFunc}
@@ -576,21 +577,69 @@ 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 || !DataTypeSupport.hasDuplicateFieldNames(s.fields)) &&
+ 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)
}
/**
diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala
index f2026013cb..c5fc0e4858 100644
--- a/spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala
+++ b/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")
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTableScanExec.scala
index 56e3a75fd3..66f1a73ed6 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTableScanExec.scala
+++ b/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(
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index c2dd19ebcd..3ad5527d8c 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++ b/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}
@@ -49,7 +49,7 @@ import org.apache.spark.util.random.XORShiftRandom
import com.google.common.base.Objects
-import org.apache.comet.{CometConf, CometExplainInfo, DataTypeSupport}
+import org.apache.comet.{CometConf, CometExplainInfo}
import org.apache.comet.CometConf.{COMET_SHUFFLE_ENABLED, COMET_SHUFFLE_MODE}
import org.apache.comet.CometSparkSessionExtensions.{cometCelebornShuffleFallbackReason, hasFallbackReason, isCometCelebornShuffleManagerEnabled, isCometShuffleManagerEnabled, isSpark40Plus, withFallbackReasons}
import org.apache.comet.serde.{Compatible, OperatorOuterClass, QueryPlanSerde, SupportLevel, Unsupported}
@@ -583,33 +583,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
- !DataTypeSupport.hasDuplicateFieldNames(fields)
- 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)) {
@@ -620,7 +593,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
}
@@ -704,31 +683,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
- !DataTypeSupport.hasDuplicateFieldNames(fields)
- 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)) {
@@ -749,7 +703,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
}
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
index a44aac42b3..157cf972d6 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
+++ b/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") {
diff --git a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala b/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala
new file mode 100644
index 0000000000..e973bbdee3
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala
@@ -0,0 +1,104 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.serde
+
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.spark.sql.types._
+
+import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus}
+import org.apache.comet.serde.QueryPlanSerde.supportedDataType
+
+class QueryPlanSerdeSuite extends AnyFunSuite {
+
+ test("supportedDataType matches each caller boundary") {
+ val complex = ArrayType(IntegerType)
+ val nestedInterval = StructType(
+ Seq(
+ StructField(
+ "i",
+ ArrayType(
+ MapType(StringType, DayTimeIntervalType(), valueContainsNull = true),
+ containsNull = true))))
+ val nestedCalendarInterval =
+ StructType(Seq(StructField("i", ArrayType(CalendarIntervalType, containsNull = true))))
+ val duplicateFields =
+ ArrayType(StructType(Seq(StructField("i", IntegerType), StructField("i", IntegerType))))
+ val emptyStruct = StructType(Nil)
+ val timeTypes = if (isSpark41Plus) Seq(DataType.fromDDL("TIME")) else Seq.empty
+ val collatedStrings =
+ if (isSpark40Plus) Seq(DataType.fromDDL("STRING COLLATE UTF8_LCASE")) else Seq.empty
+
+ val boundaries: Seq[(String, DataType => Boolean, Seq[DataType], Seq[DataType])] = Seq(
+ (
+ "expression serde defaults",
+ supportedDataType(_),
+ Seq(IntegerType, StringType, CalendarIntervalType) ++ timeTypes ++ collatedStrings,
+ Seq(complex, YearMonthIntervalType(), DayTimeIntervalType(), emptyStruct)),
+ (
+ "CometSink",
+ supportedDataType(_, allowComplex = true, allowIntervals = true),
+ Seq(complex, nestedInterval, nestedCalendarInterval, duplicateFields) ++
+ timeTypes ++ collatedStrings,
+ Seq(emptyStruct)),
+ (
+ "CometLocalTableScanExec",
+ supportedDataType(
+ _,
+ allowComplex = true,
+ allowIntervals = true,
+ allowTimeType = false,
+ allowAnyStringType = false),
+ Seq(
+ IntegerType,
+ StringType,
+ complex,
+ nestedInterval,
+ nestedCalendarInterval,
+ duplicateFields),
+ Seq(emptyStruct) ++ timeTypes ++ collatedStrings),
+ (
+ "native shuffle",
+ supportedDataType(_, allowComplex = true, allowIntervals = true),
+ Seq(complex, nestedInterval, nestedCalendarInterval, duplicateFields) ++
+ timeTypes ++ collatedStrings,
+ Seq(emptyStruct)),
+ (
+ "JVM columnar shuffle",
+ supportedDataType(
+ _,
+ allowComplex = true,
+ allowCalendarInterval = false,
+ allowDuplicateStructFieldNames = false),
+ Seq(IntegerType, StringType, complex) ++ timeTypes ++ collatedStrings,
+ Seq(
+ YearMonthIntervalType(),
+ DayTimeIntervalType(),
+ CalendarIntervalType,
+ nestedCalendarInterval,
+ duplicateFields,
+ emptyStruct)))
+
+ boundaries.foreach { case (name, supports, accepted, rejected) =>
+ accepted.foreach(dt => assert(supports(dt), s"$name should accept $dt"))
+ rejected.foreach(dt => assert(!supports(dt), s"$name should reject $dt"))
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org