This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 7123be7b52 feat: support multiply_dt_interval with codegen dispatch
(#4900)
7123be7b52 is described below
commit 7123be7b52251739c85259a8f92fff7bf961892c
Author: Peter Lee <[email protected]>
AuthorDate: Fri Jul 24 22:33:13 2026 +0800
feat: support multiply_dt_interval with codegen dispatch (#4900)
---
native/core/src/execution/planner.rs | 8 ++-
.../codegen/CometBatchKernelCodegenInput.scala | 15 +++--
.../org/apache/comet/serde/QueryPlanSerde.scala | 1 +
.../scala/org/apache/comet/serde/datetime.scala | 4 +-
.../scala/org/apache/comet/serde/literals.scala | 9 ++-
.../comet/udf/codegen/CometScalaUDFCodegen.scala | 2 +-
.../expressions/datetime/multiply_dt_interval.sql | 67 ++++++++++++++++++++++
7 files changed, 95 insertions(+), 11 deletions(-)
diff --git a/native/core/src/execution/planner.rs
b/native/core/src/execution/planner.rs
index a4aedfc991..835dfcfde4 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -365,6 +365,9 @@ impl PhysicalPlanner {
DataType::Time64(TimeUnit::Nanosecond) => {
ScalarValue::Time64Nanosecond(None)
}
+ DataType::Duration(TimeUnit::Microsecond) => {
+ ScalarValue::DurationMicrosecond(None)
+ }
dt => {
return Err(GeneralError(format!("{dt:?} is not
supported in Comet")))
}
@@ -391,9 +394,12 @@ impl PhysicalPlanner {
DataType::Timestamp(TimeUnit::Microsecond,
Some(tz)) => {
ScalarValue::TimestampMicrosecond(Some(*value), Some(tz))
}
+ DataType::Duration(TimeUnit::Microsecond) => {
+ ScalarValue::DurationMicrosecond(Some(*value))
+ }
dt => {
return Err(GeneralError(format!(
- "Expected either 'Int64' or 'Timestamp'
for LongVal, but found {dt:?}"
+ "Expected 'Int64', 'Timestamp', or
'Duration(Microsecond)' for LongVal, but found {dt:?}"
)))
}
},
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 f31c712c62..1d94bdc8dc 100644
---
a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala
+++
b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala
@@ -62,6 +62,7 @@ private[codegen] object CometBatchKernelCodegenInput extends
CometTypeShim {
classOf[Float4Vector],
classOf[Float8Vector],
classOf[DateDayVector],
+ classOf[DurationVector],
classOf[TimeNanoVector],
classOf[TimeStampMicroVector],
classOf[TimeStampMicroTZVector])
@@ -135,6 +136,7 @@ private[codegen] object CometBatchKernelCodegenInput
extends CometTypeShim {
val longCases = withOrd.collect {
case (ArrowColumnSpec(cls, _), ord)
if cls == classOf[BigIntVector] ||
+ cls == classOf[DurationVector] ||
cls == classOf[TimeNanoVector] ||
cls == classOf[TimeStampMicroVector] ||
cls == classOf[TimeStampMicroTZVector] =>
@@ -594,7 +596,8 @@ private[codegen] object CometBatchKernelCodegenInput
extends CometTypeShim {
case ByteType => s"getByte($idx)"
case ShortType => s"getShort($idx)"
case IntegerType | DateType => s"getInt($idx)"
- case LongType | TimestampType | TimestampNTZType => s"getLong($idx)"
+ case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType
=>
+ s"getLong($idx)"
case dt if isTimeType(dt) => s"getLong($idx)"
case FloatType => s"getFloat($idx)"
case DoubleType => s"getDouble($idx)"
@@ -696,7 +699,7 @@ private[codegen] object CometBatchKernelCodegenInput
extends CometTypeShim {
| public int getInt(int i) {
| return $childField.getInt(startIndex + i);
| }""".stripMargin
- case LongType | TimestampType | TimestampNTZType =>
+ case LongType | TimestampType | TimestampNTZType | _:
DayTimeIntervalType =>
s""" @Override
| public long getLong(int i) {
| return $childField.getLong(startIndex + i);
@@ -854,7 +857,7 @@ private[codegen] object CometBatchKernelCodegenInput
extends CometTypeShim {
s" case $fi: return ${path}_f$fi.getShort(this.rowIdx);"
case IntegerType | DateType =>
s" case $fi: return ${path}_f$fi.getInt(this.rowIdx);"
- case LongType | TimestampType | TimestampNTZType =>
+ case LongType | TimestampType | TimestampNTZType | _:
DayTimeIntervalType =>
s" case $fi: return ${path}_f$fi.getLong(this.rowIdx);"
case dt if isTimeType(dt) =>
s" case $fi: return ${path}_f$fi.getLong(this.rowIdx);"
@@ -908,8 +911,10 @@ private[codegen] object CometBatchKernelCodegenInput
extends CometTypeShim {
val longCases = scalarOrd.collect {
case (f, fi)
if f.sparkType == LongType || f.sparkType == TimestampType ||
- f.sparkType == TimestampNTZType || isTimeType(f.sparkType) =>
- fieldReadScalar(fi, LongType, f.nullable)
+ f.sparkType == TimestampNTZType ||
+ f.sparkType.isInstanceOf[DayTimeIntervalType] ||
+ isTimeType(f.sparkType) =>
+ fieldReadScalar(fi, f.sparkType, f.nullable)
}
val floatCases =
scalarOrd.collect {
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 f62db397cb..19ac36b040 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -301,6 +301,7 @@ object QueryPlanSerde extends Logging with CometExprShim
with CometTypeShim {
classOf[MakeTimestamp] -> CometMakeTimestamp,
classOf[MakeYMInterval] -> CometMakeYMInterval,
classOf[MakeDTInterval] -> CometMakeDTInterval,
+ classOf[MultiplyDTInterval] -> CometMultiplyDTInterval,
classOf[MicrosToTimestamp] -> CometMicrosToTimestamp,
classOf[MillisToTimestamp] -> CometMillisToTimestamp,
classOf[MonthsBetween] -> CometMonthsBetween,
diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala
b/spark/src/main/scala/org/apache/comet/serde/datetime.scala
index 9f2533d85f..d968c00cd1 100644
--- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala
@@ -21,7 +21,7 @@ package org.apache.comet.serde
import java.util.Locale
-import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute,
ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub,
DayOfMonth, DayOfWeek, DayOfYear, Days, Expression, FromUTCTimestamp,
GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate,
MakeDTInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp,
MillisToTimestamp, Minute, Month, MonthsBetween, NextDay,
PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, ToUnixTimest
[...]
+import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute,
ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub,
DayOfMonth, DayOfWeek, DayOfYear, Days, Expression, FromUTCTimestamp,
GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate,
MakeDTInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp,
MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, NextDay,
PreciseTimestampConversion, Quarter, Second, SecondsToTim [...]
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{DataType, DateType, DoubleType, FloatType,
IntegerType, LongType, StringType, TimestampNTZType, TimestampType}
import org.apache.spark.unsafe.types.UTF8String
@@ -954,6 +954,8 @@ object CometMakeYMInterval extends
CometCodegenDispatch[MakeYMInterval]
object CometMakeDTInterval extends CometCodegenDispatch[MakeDTInterval]
+object CometMultiplyDTInterval extends CometCodegenDispatch[MultiplyDTInterval]
+
/**
* Spark's internal `PreciseTimestampConversion` reinterprets a value between
the timestamp types
* (`TimestampType` / `TimestampNTZType`) and `LongType` without losing
microsecond precision. It
diff --git a/spark/src/main/scala/org/apache/comet/serde/literals.scala
b/spark/src/main/scala/org/apache/comet/serde/literals.scala
index 4f2a5dfa5e..bac3631bab 100644
--- a/spark/src/main/scala/org/apache/comet/serde/literals.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala
@@ -24,7 +24,7 @@ import java.lang
import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.expressions.{Attribute, Literal}
import org.apache.spark.sql.catalyst.util.ArrayData
-import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType,
ByteType, DateType, Decimal, DecimalType, DoubleType, FloatType, IntegerType,
LongType, NullType, ShortType, StringType, TimestampNTZType, TimestampType}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType,
ByteType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType,
FloatType, IntegerType, LongType, NullType, ShortType, StringType,
TimestampNTZType, TimestampType}
import org.apache.spark.unsafe.types.UTF8String
import com.google.protobuf.ByteString
@@ -56,7 +56,10 @@ object CometLiteral extends CometExpressionSerde[Literal]
with Logging {
.isInstanceOf[ArrayType])))) {
Compatible(None)
} else {
- Unsupported(Some(s"Unsupported data type ${expr.dataType}"))
+ expr.dataType match {
+ case _: DayTimeIntervalType => Compatible(None)
+ case _ => Unsupported(Some(s"Unsupported data type ${expr.dataType}"))
+ }
}
}
@@ -78,7 +81,7 @@ object CometLiteral extends CometExpressionSerde[Literal]
with Logging {
case _: ByteType => exprBuilder.setByteVal(value.asInstanceOf[Byte])
case _: ShortType => exprBuilder.setShortVal(value.asInstanceOf[Short])
case _: IntegerType | _: DateType =>
exprBuilder.setIntVal(value.asInstanceOf[Int])
- case _: LongType | _: TimestampType | _: TimestampNTZType =>
+ case _: LongType | _: TimestampType | _: TimestampNTZType | _:
DayTimeIntervalType =>
exprBuilder.setLongVal(value.asInstanceOf[Long])
case dt if isTimeType(dt) =>
exprBuilder.setLongVal(value.asInstanceOf[Long])
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 f575dd5b53..145698102e 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
@@ -219,7 +219,7 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
StructColumnSpec(nullable = true, fieldSpecs)
case _: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector |
_: BigIntVector |
_: Float4Vector | _: Float8Vector | _: DecimalVector | _:
VarCharVector |
- _: VarBinaryVector | _: DateDayVector | _: TimeStampMicroVector |
+ _: VarBinaryVector | _: DateDayVector | _: DurationVector | _:
TimeStampMicroVector |
_: TimeStampMicroTZVector =>
ScalarColumnSpec(v.getClass.asInstanceOf[Class[_ <: ValueVector]],
nullable = true)
case other =>
diff --git
a/spark/src/test/resources/sql-tests/expressions/datetime/multiply_dt_interval.sql
b/spark/src/test/resources/sql-tests/expressions/datetime/multiply_dt_interval.sql
new file mode 100644
index 0000000000..65f7535a6d
--- /dev/null
+++
b/spark/src/test/resources/sql-tests/expressions/datetime/multiply_dt_interval.sql
@@ -0,0 +1,67 @@
+-- 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.
+
+-- Routes multiply_dt_interval through the codegen dispatcher; produces
DayTimeIntervalType.
+-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true
+
+statement
+CREATE TABLE test_multiply_dt_interval(days int, hours int, minutes int,
seconds decimal(18,6), b tinyint, s smallint, i int, l long, f float, d double,
dec decimal(10,2)) USING parquet
+
+statement
+INSERT INTO test_multiply_dt_interval VALUES
+ (1, 2, 3, 4.500000, CAST(2 AS TINYINT), CAST(3 AS SMALLINT), 2, CAST(3 AS
BIGINT), CAST(1.5 AS FLOAT), CAST(2.5 AS DOUBLE), CAST(2.50 AS DECIMAL(10, 2))),
+ (-1, 0, 30, 15.250000, CAST(-2 AS TINYINT), CAST(-3 AS SMALLINT), -2,
CAST(-3 AS BIGINT), CAST(-1.5 AS FLOAT), CAST(-2.5 AS DOUBLE), CAST(-2.50 AS
DECIMAL(10, 2))),
+ (0, 0, 0, 0.000001, CAST(0 AS TINYINT), CAST(0 AS SMALLINT), 0, CAST(0 AS
BIGINT), CAST(0.5 AS FLOAT), CAST(0.5 AS DOUBLE), CAST(0.50 AS DECIMAL(10, 2))),
+ (2, -6, 0, 0.000000, NULL, NULL, NULL, NULL, NULL, NULL, NULL)
+
+query
+SELECT
+ make_dt_interval(days, hours, minutes, seconds) * b,
+ make_dt_interval(days, hours, minutes, seconds) * s,
+ make_dt_interval(days, hours, minutes, seconds) * i,
+ make_dt_interval(days, hours, minutes, seconds) * l,
+ make_dt_interval(days, hours, minutes, seconds) * f,
+ make_dt_interval(days, hours, minutes, seconds) * d,
+ make_dt_interval(days, hours, minutes, seconds) * dec
+FROM test_multiply_dt_interval
+
+-- literal interval input
+query
+SELECT INTERVAL '1 02:03:04.500000' DAY TO SECOND * i FROM
test_multiply_dt_interval
+
+-- numeric on the left is normalized by Spark to multiply_dt_interval.
+query
+SELECT i * make_dt_interval(days, hours, minutes, seconds),
+ 2 * INTERVAL '1 02:03:04.500000' DAY TO SECOND
+FROM test_multiply_dt_interval
+
+-- literal multipliers, including half-up rounding to the nearest microsecond.
+query
+SELECT
+ make_dt_interval(1, 2, 3, 4.5) * 2,
+ INTERVAL '0.000001' SECOND * 0.5D,
+ INTERVAL '0.000001' SECOND * CAST(0.50 AS DECIMAL(10, 2)),
+ make_dt_interval(-1, 0, 30, 15.25) * 1.5D
+
+-- null interval input
+query
+SELECT make_dt_interval(NULL, hours, minutes, seconds) * 2
+FROM test_multiply_dt_interval
+
+-- 106751991 days is valid, but multiplying it by two exceeds the int64
microsecond range.
+query expect_error(overflow)
+SELECT make_dt_interval(106751991) * 2
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
