diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala index 801eb70412..2ed7e33c90 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala @@ -66,6 +66,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { classOf[TimeNanoVector], classOf[TimeStampMicroVector], classOf[TimeStampMicroTZVector], + classOf[IntervalYearVector], classOf[IntervalMonthDayNanoVector]) private val cometPlainVectorName: String = classOf[CometPlainVector].getName @@ -131,7 +132,8 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { } val intCases = withOrd.collect { case (ArrowColumnSpec(cls, _), ord) - if cls == classOf[IntVector] || cls == classOf[DateDayVector] => + if cls == classOf[IntVector] || cls == classOf[DateDayVector] || + cls == classOf[IntervalYearVector] => s" case $ord: return this.col$ord.getInt(this.rowIdx);" } val longCases = withOrd.collect { @@ -604,7 +606,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { case BooleanType => s"getBoolean($idx)" case ByteType => s"getByte($idx)" case ShortType => s"getShort($idx)" - case IntegerType | DateType => s"getInt($idx)" + case IntegerType | DateType | _: YearMonthIntervalType => s"getInt($idx)" case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType => s"getLong($idx)" case CalendarIntervalType => s"getInterval($idx)" @@ -704,7 +706,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { | public short getShort(int i) { | return $childField.getShort(startIndex + i); | }""".stripMargin - case IntegerType | DateType => + case IntegerType | DateType | _: YearMonthIntervalType => s""" @Override | public int getInt(int i) { | return $childField.getInt(startIndex + i); @@ -870,7 +872,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { s" case $fi: return ${path}_f$fi.getByte(this.rowIdx);" case ShortType => s" case $fi: return ${path}_f$fi.getShort(this.rowIdx);" - case IntegerType | DateType => + case IntegerType | DateType | _: YearMonthIntervalType => s" case $fi: return ${path}_f$fi.getInt(this.rowIdx);" case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType => s" case $fi: return ${path}_f$fi.getLong(this.rowIdx);" @@ -924,7 +926,9 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { fieldReadScalar(fi, ShortType, f.nullable) } val intCases = scalarOrd.collect { - case (f, fi) if f.sparkType == IntegerType || f.sparkType == DateType => + case (f, fi) + if f.sparkType == IntegerType || f.sparkType == DateType || + f.sparkType.isInstanceOf[YearMonthIntervalType] => fieldReadScalar(fi, IntegerType, f.nullable) } val longCases = scalarOrd.collect { diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala index 24c18fcb49..33e6c0c035 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -410,8 +410,9 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { case BooleanType => s"$target.getBoolean($idx)" case ByteType => s"$target.getByte($idx)" case ShortType => s"$target.getShort($idx)" - case IntegerType | DateType => s"$target.getInt($idx)" - case LongType | TimestampType | TimestampNTZType => s"$target.getLong($idx)" + case IntegerType | DateType | _: YearMonthIntervalType => s"$target.getInt($idx)" + case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType => + s"$target.getLong($idx)" case CalendarIntervalType => s"$target.getInterval($idx)" case dt if isTimeType(dt) => s"$target.getLong($idx)" case FloatType => s"$target.getFloat($idx)" diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala b/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala index 41353fb3d8..d258e9e6ef 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala @@ -40,8 +40,9 @@ private[codegen] object CometSpecializedGettersDispatch { case BooleanType => java.lang.Boolean.valueOf(g.getBoolean(ordinal)) case ByteType => java.lang.Byte.valueOf(g.getByte(ordinal)) case ShortType => java.lang.Short.valueOf(g.getShort(ordinal)) - case IntegerType | DateType => java.lang.Integer.valueOf(g.getInt(ordinal)) - case LongType | TimestampType | TimestampNTZType => + case IntegerType | DateType | _: YearMonthIntervalType => + java.lang.Integer.valueOf(g.getInt(ordinal)) + case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType => java.lang.Long.valueOf(g.getLong(ordinal)) case CalendarIntervalType => g.getInterval(ordinal) case FloatType => java.lang.Float.valueOf(g.getFloat(ordinal)) 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 ed69f82e2a..f2026013cb 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.DataType +import org.apache.spark.sql.types.{ArrayType, DataType, DayTimeIntervalType, MapType, StructType, YearMonthIntervalType} import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.withFallbackReason @@ -43,6 +43,16 @@ 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 @@ -54,8 +64,7 @@ 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 => supportedDataType(a.dataType, allowComplex = true)) + val supportedTypes = op.output.forall(a => supportedSinkDataType(a.dataType)) if (!supportedTypes) { withFallbackReason(op, "Unsupported data type") @@ -123,8 +132,7 @@ object CometExchangeSink extends CometSink[SparkPlan] { private def convertToShuffleScan( op: SparkPlan, builder: Operator.Builder): Option[OperatorOuterClass.Operator] = { - val supportedTypes = - op.output.forall(a => supportedDataType(a.dataType, allowComplex = true)) + val supportedTypes = op.output.forall(a => supportedSinkDataType(a.dataType)) if (!supportedTypes) { withFallbackReason(op, "Unsupported data type for shuffle direct read") diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala index 91cb1eb543..56746f5156 100644 --- a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala @@ -220,7 +220,7 @@ class CometScalaUDFCodegen extends CometUDF with Logging { case _: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector | _: BigIntVector | _: Float4Vector | _: Float8Vector | _: DecimalVector | _: VarCharVector | _: VarBinaryVector | _: DateDayVector | _: DurationVector | _: TimeStampMicroVector | - _: TimeStampMicroTZVector | _: IntervalMonthDayNanoVector => + _: TimeStampMicroTZVector | _: IntervalYearVector | _: IntervalMonthDayNanoVector => ScalarColumnSpec(v.getClass.asInstanceOf[Class[_ <: ValueVector]], nullable = true) case other => throw new UnsupportedOperationException( 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 28ff6c85c2..d5bd16bdae 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,7 +31,7 @@ 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, NullType, StructType} +import org.apache.spark.sql.types.{DataType, DayTimeIntervalType, NullType, StructType, YearMonthIntervalType} import com.google.common.base.Objects @@ -138,7 +138,7 @@ object CometLocalTableScanExec extends CometSink[LocalTableScanExec] with DataTy dt: DataType, name: String, fallbackReasons: ListBuffer[String]): Boolean = dt match { - case _: NullType => true + case _: NullType | _: YearMonthIntervalType | _: DayTimeIntervalType => true case _ => super.isTypeSupported(dt, name, fallbackReasons) } 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 ed6bdcbcc0..2a50bbd89b 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, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType} +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.vectorized.ColumnarBatch import org.apache.spark.util.MutablePair import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, RecordComparator} @@ -413,7 +413,7 @@ object CometShuffleExchangeExec case _: BooleanType | _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType | _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType | _: DecimalType | _: DateType | _: NullType | - CalendarIntervalType => + _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType => true case dt if isTimeType(dt) => true diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_dt_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_dt_interval.sql index 2c61040de9..371f01be43 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_dt_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_dt_interval.sql @@ -15,6 +15,9 @@ -- specific language governing permissions and limitations -- under the License. +-- Config: spark.comet.exec.localTableScan.enabled=true +-- Config: spark.comet.exec.shuffle.mode=native + -- Routes make_dt_interval through the codegen dispatcher; produces DayTimeIntervalType. statement @@ -34,6 +37,33 @@ SELECT make_dt_interval(1, 2, 3, 4.5), make_dt_interval(0, 0, 0, 0) query SELECT make_dt_interval(1), make_dt_interval(1, 2), make_dt_interval() +-- nested interval output through LocalTableScan and codegen dispatch +query +SELECT transform(a, x -> x) AS result +FROM VALUES + (array(make_dt_interval(1, 2, 3, 4.5), CAST(NULL AS INTERVAL DAY TO SECOND))) +AS t(a) + +-- top-level, struct, and map interval output through native shuffle +query +SELECT d, h, mi, s, + make_dt_interval(d, h, mi, s) AS i, + named_struct('i', make_dt_interval(d, h, mi, s)) AS st, + map('i', make_dt_interval(d, h, mi, s)) AS m +FROM test_mdi +DISTRIBUTE BY d + +-- null top-level, struct, and map interval output through native shuffle +query +SELECT i, st, m +FROM VALUES + (1, + CAST(NULL AS INTERVAL DAY TO SECOND), + named_struct('i', CAST(NULL AS INTERVAL DAY TO SECOND)), + map('i', CAST(NULL AS INTERVAL DAY TO SECOND))) +AS t(k, i, st, m) +DISTRIBUTE BY k + -- overflow: days * MICROS_PER_DAY exceeds the int64 microsecond range. makeDayTimeInterval throws -- unconditionally (not ANSI-gated); this confirms the dispatched codegen path propagates Spark's -- exception. The pattern is the lowercase word so it matches every version: Spark 4.x raises diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_ym_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_ym_interval.sql index a642b4925d..e2763cce7a 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_ym_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_ym_interval.sql @@ -15,6 +15,9 @@ -- specific language governing permissions and limitations -- under the License. +-- Config: spark.comet.exec.localTableScan.enabled=true +-- Config: spark.comet.exec.shuffle.mode=native + -- Routes make_ym_interval through the codegen dispatcher; produces YearMonthIntervalType. statement @@ -34,6 +37,33 @@ SELECT make_ym_interval(1, 2), make_ym_interval(0, 0), make_ym_interval(-5, 11) query SELECT make_ym_interval(3), make_ym_interval() +-- nested interval output through LocalTableScan and codegen dispatch +query +SELECT transform(a, x -> x) AS result +FROM VALUES + (array(make_ym_interval(1, 2), CAST(NULL AS INTERVAL YEAR TO MONTH))) +AS t(a) + +-- top-level, struct, and map interval output through native shuffle +query +SELECT y, m, + make_ym_interval(y, m) AS i, + named_struct('i', make_ym_interval(y, m)) AS s, + map('i', make_ym_interval(y, m)) AS m +FROM test_myi +DISTRIBUTE BY y + +-- null top-level, struct, and map interval output through native shuffle +query +SELECT i, s, m +FROM VALUES + (1, + CAST(NULL AS INTERVAL YEAR TO MONTH), + named_struct('i', CAST(NULL AS INTERVAL YEAR TO MONTH)), + map('i', CAST(NULL AS INTERVAL YEAR TO MONTH))) +AS t(k, i, s, m) +DISTRIBUTE BY k + -- overflow: years * 12 exceeds Int range. makeYearMonthInterval throws unconditionally (not -- ANSI-gated); this confirms the dispatched codegen path propagates Spark's exception. The -- pattern is the lowercase word so it matches every version: Spark 4.x raises