diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 8b15d829aad..85a91e951c8 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.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 }}] diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 20046ff5b10..a12a5c9b297 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.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 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 be4bc9c3412..11f624544f5 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/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) } /** 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 f2026013cbf..c5fc0e4858a 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 56e3a75fd3e..66f1a73ed60 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 1f5d06a53af..58b75531a0f 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} @@ -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 } 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 bf354356386..911062dd33e 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 00000000000..e973bbdee34 --- /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")) + } + } +}