From 72997e9e8f6d234d6811174e174f9cf96c0dd96e Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 7 Aug 2026 22:13:45 +0800 Subject: [PATCH 01/20] fix: preserve CalendarInterval microseconds --- docs/source/user-guide/latest/expressions.md | 2 +- native/Cargo.lock | 1 - native/core/src/execution/serde.rs | 5 +- native/spark-expr/Cargo.toml | 3 +- .../src/datetime_funcs/make_interval.rs | 181 ++++++++++++++---- native/spark-expr/src/datetime_funcs/mod.rs | 2 +- native/spark-expr/src/lib.rs | 8 +- .../comet/vector/CometStructVector.java | 10 + .../CometBatchKernelCodegenInput.scala | 20 +- .../CometBatchKernelCodegenOutput.scala | 19 +- .../org/apache/comet/serde/datetime.scala | 24 +-- .../udf/codegen/CometScalaUDFCodegen.scala | 2 + .../comet/execution/arrow/ArrowWriters.scala | 23 +++ .../apache/spark/sql/comet/util/Utils.scala | 40 +++- .../expressions/datetime/make_interval.sql | 43 ++++- .../datetime/make_interval_ansi.sql | 43 ++++- .../datetime/make_interval_dispatch.sql | 4 +- .../datetime/make_interval_dispatch_ansi.sql | 2 +- .../datetime/try_make_interval.sql | 12 ++ .../org/apache/comet/CometCodegenSuite.scala | 16 +- .../CometDatetimeExpressionBenchmark.scala | 9 +- .../arrow/CometArrowStreamSuite.scala | 17 +- 22 files changed, 367 insertions(+), 119 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 192eae807eb..9515efaae72 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -277,7 +277,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | `localtimestamp` | ✅ | — | | | `make_date` | ✅ | Native | | | `make_dt_interval` | ✅ | Codegen dispatch | | -| `make_interval` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; intervals outside Arrow's nanosecond range are tracked by [#5279](https://github.com/apache/datafusion-comet/issues/5279); the native path is opt-in via allowIncompatible ([details](compatibility/expressions/datetime.md)) | +| `make_interval` | ✅ | Native | | | `make_time` | 🔜 | — | Spark 4.1 TIME type; tracked by [#4288](https://github.com/apache/datafusion-comet/issues/4288) | | `make_timestamp` | ✅ | Hybrid | | | `make_timestamp_ltz` | ✅ | — | 2-arg TIME form falls back | diff --git a/native/Cargo.lock b/native/Cargo.lock index ae8a2f7be7f..eeaa9da5ce2 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2066,7 +2066,6 @@ dependencies = [ "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", - "datafusion-spark", "futures", "jni 0.22.4", "num", diff --git a/native/core/src/execution/serde.rs b/native/core/src/execution/serde.rs index f31cb9cd35a..31f3553e474 100644 --- a/native/core/src/execution/serde.rs +++ b/native/core/src/execution/serde.rs @@ -31,6 +31,7 @@ use datafusion_comet_proto::{ spark_expression::DataType, spark_operator, }; +use datafusion_comet_spark_expr::calendar_interval_type; use prost::Message; use std::{io::Cursor, sync::Arc}; @@ -102,9 +103,7 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType { // Spark's DayTimeIntervalType stores microseconds in an int64, which matches Arrow // Duration(Microsecond) rather than the lossy Interval(DayTime) {days, millis} layout. DataTypeId::DayTimeInterval => ArrowDataType::Duration(TimeUnit::Microsecond), - // Spark's CalendarIntervalType stores months, days, and microseconds. Arrow stores the - // same components with nanosecond precision. - DataTypeId::CalendarInterval => ArrowDataType::Interval(IntervalUnit::MonthDayNano), + DataTypeId::CalendarInterval => calendar_interval_type(), DataTypeId::Null => ArrowDataType::Null, DataTypeId::List => match dt_value .type_info diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6faa9fec4ec..a1de8616fe9 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -30,7 +30,6 @@ edition = { workspace = true } arrow = { workspace = true } chrono = { workspace = true } datafusion = { workspace = true } -datafusion-spark = { workspace = true } chrono-tz = { workspace = true } num = { workspace = true } regex = { workspace = true } @@ -222,4 +221,4 @@ harness = false [[bench]] name = "cast_int_to_decimal" -harness = false \ No newline at end of file +harness = false diff --git a/native/spark-expr/src/datetime_funcs/make_interval.rs b/native/spark-expr/src/datetime_funcs/make_interval.rs index 9f7b21ca2a7..bb721442b6f 100644 --- a/native/spark-expr/src/datetime_funcs/make_interval.rs +++ b/native/spark-expr/src/datetime_funcs/make_interval.rs @@ -16,72 +16,175 @@ // under the License. use crate::arithmetic_overflow_error; -use arrow::array::Array; -use arrow::datatypes::DataType; +use arrow::array::{Array, ArrayRef, Decimal128Array, Int32Array, Int64Array, StructArray}; +use arrow::buffer::NullBuffer; +use arrow::datatypes::{DataType, Field, Fields}; use datafusion::common::Result; -use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature}; -use datafusion_spark::function::datetime::make_interval::SparkMakeInterval as DataFusionMakeInterval; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use std::collections::HashMap; +use std::sync::Arc; + +const CALENDAR_INTERVAL_STRUCT_KEY: &str = "SPARK::calendarInterval::struct"; +const MICROS_PER_HOUR: i64 = 3_600_000_000; +const MICROS_PER_MINUTE: i64 = 60_000_000; + +pub fn calendar_interval_type() -> DataType { + let months = Field::new("months", DataType::Int32, false).with_metadata(HashMap::from([( + CALENDAR_INTERVAL_STRUCT_KEY.to_string(), + "true".to_string(), + )])); + DataType::Struct(Fields::from(vec![ + months, + Field::new("days", DataType::Int32, false), + Field::new("microseconds", DataType::Int64, false), + ])) +} #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkMakeInterval { - inner: DataFusionMakeInterval, + signature: Signature, fail_on_error: bool, } impl SparkMakeInterval { pub fn new(fail_on_error: bool) -> Self { Self { - inner: DataFusionMakeInterval::new(), + signature: Signature::exact( + vec![ + DataType::Int32, + DataType::Int32, + DataType::Int32, + DataType::Int32, + DataType::Int32, + DataType::Int32, + DataType::Decimal128(18, 6), + ], + Volatility::Immutable, + ), fail_on_error, } } } +fn make_interval( + years: i32, + months: i32, + weeks: i32, + days: i32, + hours: i32, + minutes: i32, + seconds_micros: i128, +) -> Option<(i32, i32, i64)> { + let months = years.checked_mul(12)?.checked_add(months)?; + let days = weeks.checked_mul(7)?.checked_add(days)?; + let micros = i64::try_from(seconds_micros) + .ok()? + .checked_add(i64::from(hours).checked_mul(MICROS_PER_HOUR)?)? + .checked_add(i64::from(minutes).checked_mul(MICROS_PER_MINUTE)?)?; + Some((months, days, micros)) +} + impl ScalarUDFImpl for SparkMakeInterval { fn name(&self) -> &str { - self.inner.name() + "make_interval" } fn signature(&self) -> &Signature { - self.inner.signature() + &self.signature } - fn return_type(&self, arg_types: &[DataType]) -> Result { - self.inner.return_type(arg_types) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(calendar_interval_type()) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let inputs = if self.fail_on_error { - Some(args.args.clone()) - } else { - None - }; - let result = self.inner.invoke_with_args(args)?; - - if let Some(inputs) = inputs { - let inputs_are_valid = |i| { - inputs.iter().all(|input| match input { - ColumnarValue::Array(values) => values.is_valid(i), - ColumnarValue::Scalar(value) => !value.is_null(), - }) - }; - let overflow = match &result { - ColumnarValue::Array(values) => values.nulls().is_some_and(|nulls| { - nulls.null_count() != 0 - && nulls - .iter() - .enumerate() - .any(|(i, is_valid)| !is_valid && inputs_are_valid(i)) - }), - ColumnarValue::Scalar(value) => value.is_null() && inputs_are_valid(0), - }; - if overflow { - // Spark identifies the integer or long operation that overflowed. The native - // wrapper only sees the result null mask, so it can only report interval overflow. - return Err(arithmetic_overflow_error("interval").into()); + let number_rows = args.number_rows; + let arrays = args + .args + .into_iter() + .map(|arg| arg.into_array(number_rows)) + .collect::>>()?; + let years = arrays[0].as_any().downcast_ref::().unwrap(); + let months = arrays[1].as_any().downcast_ref::().unwrap(); + let weeks = arrays[2].as_any().downcast_ref::().unwrap(); + let days = arrays[3].as_any().downcast_ref::().unwrap(); + let hours = arrays[4].as_any().downcast_ref::().unwrap(); + let minutes = arrays[5].as_any().downcast_ref::().unwrap(); + let seconds = arrays[6] + .as_any() + .downcast_ref::() + .unwrap(); + + let mut result_months = Vec::with_capacity(years.len()); + let mut result_days = Vec::with_capacity(years.len()); + let mut result_micros = Vec::with_capacity(years.len()); + let mut valid = Vec::with_capacity(years.len()); + + for i in 0..years.len() { + if arrays.iter().any(|array| array.is_null(i)) { + result_months.push(0); + result_days.push(0); + result_micros.push(0); + valid.push(false); + continue; + } + + match make_interval( + years.value(i), + months.value(i), + weeks.value(i), + days.value(i), + hours.value(i), + minutes.value(i), + seconds.value(i), + ) { + Some((months, days, micros)) => { + result_months.push(months); + result_days.push(days); + result_micros.push(micros); + valid.push(true); + } + None if self.fail_on_error => { + return Err(arithmetic_overflow_error("interval").into()); + } + None => { + result_months.push(0); + result_days.push(0); + result_micros.push(0); + valid.push(false); + } } } - Ok(result) + let columns: Vec = vec![ + Arc::new(Int32Array::from(result_months)), + Arc::new(Int32Array::from(result_days)), + Arc::new(Int64Array::from(result_micros)), + ]; + let DataType::Struct(fields) = calendar_interval_type() else { + unreachable!() + }; + Ok(ColumnarValue::Array(Arc::new(StructArray::new( + fields, + columns, + Some(NullBuffer::from(valid)), + )))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn preserves_spark_microsecond_range_and_overflow() { + assert_eq!( + make_interval(1, 2, 3, 4, 2_562_048, 0, 123_456_789_012_123_456), + Some((14, 25, 132_680_161_812_123_456)) + ); + assert!(make_interval(i32::MAX, 0, 0, 0, 0, 0, 0).is_none()); + assert!(make_interval(0, 0, 0, 0, i32::MAX, i32::MAX, i128::MAX).is_none()); } } diff --git a/native/spark-expr/src/datetime_funcs/mod.rs b/native/spark-expr/src/datetime_funcs/mod.rs index 05530f29c24..74f4f4ea12b 100644 --- a/native/spark-expr/src/datetime_funcs/mod.rs +++ b/native/spark-expr/src/datetime_funcs/mod.rs @@ -39,7 +39,7 @@ pub use extract_date_part::SparkMinute; pub use extract_date_part::SparkSecond; pub use hours::SparkHoursTransform; pub use make_date::SparkMakeDate; -pub use make_interval::SparkMakeInterval; +pub use make_interval::{calendar_interval_type, SparkMakeInterval}; pub use make_time::SparkMakeTime; pub use next_day::SparkNextDay; pub use seconds_to_timestamp::SparkSecondsToTimestamp; diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 2b5c29befc6..20294d90e3d 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -76,10 +76,10 @@ pub use comet_scalar_funcs::{ }; pub use csv_funcs::*; pub use datetime_funcs::{ - spark_day_name, spark_month_name, spark_to_time, SparkDateDiff, SparkDateFromUnixDate, - SparkDateTrunc, SparkHour, SparkHoursTransform, SparkMakeDate, SparkMakeInterval, - SparkMakeTime, SparkMinute, SparkNextDay, SparkSecond, SparkSecondsToTimestamp, - SparkUnixTimestamp, TimestampTruncExpr, + calendar_interval_type, spark_day_name, spark_month_name, spark_to_time, SparkDateDiff, + SparkDateFromUnixDate, SparkDateTrunc, SparkHour, SparkHoursTransform, SparkMakeDate, + SparkMakeInterval, SparkMakeTime, SparkMinute, SparkNextDay, SparkSecond, + SparkSecondsToTimestamp, SparkUnixTimestamp, TimestampTruncExpr, }; pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult}; pub use hash_funcs::*; diff --git a/spark/src/main/java/org/apache/comet/vector/CometStructVector.java b/spark/src/main/java/org/apache/comet/vector/CometStructVector.java index 259793b831a..25514f07f33 100644 --- a/spark/src/main/java/org/apache/comet/vector/CometStructVector.java +++ b/spark/src/main/java/org/apache/comet/vector/CometStructVector.java @@ -27,6 +27,7 @@ import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.util.TransferPair; import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.unsafe.types.CalendarInterval; /** * A {@link CometDecodedVector} for Spark struct columns, wrapping an Arrow {@link StructVector} and @@ -58,6 +59,15 @@ public ColumnVector getChild(int i) { return children.get(i); } + @Override + public CalendarInterval getInterval(int rowId) { + if (isNullAt(rowId)) return null; + return new CalendarInterval( + children.get(0).getInt(rowId), + children.get(1).getInt(rowId), + children.get(2).getLong(rowId)); + } + @Override public CometVector slice(int offset, int length) { TransferPair tp = this.valueVector.getTransferPair(this.valueVector.getAllocator()); 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 2ed7e33c904..69586898da4 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.types._ import org.apache.comet.codegen.CometBatchKernelCodegen.{ArrayColumnSpec, ArrowColumnSpec, MapColumnSpec, ScalarColumnSpec, StructColumnSpec} import org.apache.comet.shims.CometTypeShim -import org.apache.comet.vector.CometPlainVector +import org.apache.comet.vector.{CometPlainVector, CometStructVector} /** * Input-side emitters for the codegen kernel: typed field declarations, per-batch input casts, @@ -69,6 +69,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { classOf[IntervalYearVector], classOf[IntervalMonthDayNanoVector]) private val cometPlainVectorName: String = classOf[CometPlainVector].getName + private val cometStructVectorName: String = classOf[CometStructVector].getName /** Emit kernel typed-vector field declarations for every level of every input column. */ def emitInputFieldDecls(inputSchema: Seq[ArrowColumnSpec]): String = { @@ -112,6 +113,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { // CometPlainVector exposes `isNullAt`; Arrow-typed fields expose `isNull`. Same semantics. val method = spec.vectorClass match { case cls if wrapsInCometPlainVector(cls) => "isNullAt" + case cls if cls == classOf[StructVector] => "isNullAt" case _ => "isNull" } s" case $ord: return this.col$ord.$method(this.rowIdx);" @@ -146,7 +148,8 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { s" case $ord: return this.col$ord.getLong(this.rowIdx);" } val intervalCases = withOrd.collect { - case (ArrowColumnSpec(cls, _), ord) if cls == classOf[IntervalMonthDayNanoVector] => + case (ArrowColumnSpec(cls, _), ord) + if cls == classOf[IntervalMonthDayNanoVector] || cls == classOf[StructVector] => s" case $ord: return this.col$ord.getInterval(this.rowIdx);" } val floatCases = withOrd.collect { @@ -425,6 +428,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { */ def nullCheckMethod(spec: ArrowColumnSpec): String = spec match { case sc: ScalarColumnSpec if wrapsInCometPlainVector(sc.vectorClass) => "isNullAt" + case sc: ScalarColumnSpec if sc.vectorClass == classOf[StructVector] => "isNullAt" case _ => "isNull" } @@ -436,9 +440,13 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { // Primitive scalars wrap in CometPlainVector for JIT-inlined Platform.get* against a // cached buffer address. Decimal/VarChar/VarBinary stay on the Arrow typed field with // cached data- (and offset-) buffer addresses for inline unsafe reads. - val fieldClass = - if (wrapsInCometPlainVector(sc.vectorClass)) cometPlainVectorName - else sc.vectorClass.getName + val fieldClass = if (wrapsInCometPlainVector(sc.vectorClass)) { + cometPlainVectorName + } else if (sc.vectorClass == classOf[StructVector]) { + cometStructVectorName + } else { + sc.vectorClass.getName + } out += s"private $fieldClass $path;" if (needsValueAddrField(sc.vectorClass)) { out += s"private long ${path}_valueAddr;" @@ -471,6 +479,8 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { case sc: ScalarColumnSpec => if (wrapsInCometPlainVector(sc.vectorClass)) { out += s"this.$path = new $cometPlainVectorName($source);" + } else if (sc.vectorClass == classOf[StructVector]) { + out += s"this.$path = new $cometStructVectorName($source, null);" } else { out += s"this.$path = (${sc.vectorClass.getName}) $source;" } 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 33e6c0c0355..996b037c406 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -175,7 +175,7 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { case TimestampNTZType => classOf[TimeStampMicroVector].getName case _: YearMonthIntervalType => classOf[IntervalYearVector].getName case _: DayTimeIntervalType => classOf[DurationVector].getName - case CalendarIntervalType => classOf[IntervalMonthDayNanoVector].getName + case CalendarIntervalType => classOf[StructVector].getName case _: ArrayType => classOf[ListVector].getName case _: StructType => classOf[StructVector].getName case _: MapType => classOf[MapVector].getName @@ -222,11 +222,22 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { case CalendarIntervalType => val set = if (nested) "setSafe" else "set" val interval = ctx.freshName("interval") + val months = ctx.freshName("intervalMonths") + val days = ctx.freshName("intervalDays") + val micros = ctx.freshName("intervalMicros") OutputEmit( - "", + s"""${classOf[IntVector].getName} $months = + | (${classOf[IntVector].getName}) $targetVec.getChildByOrdinal(0); + |${classOf[IntVector].getName} $days = + | (${classOf[IntVector].getName}) $targetVec.getChildByOrdinal(1); + |${classOf[BigIntVector].getName} $micros = + | (${classOf[ + BigIntVector].getName}) $targetVec.getChildByOrdinal(2);""".stripMargin, s"""org.apache.spark.unsafe.types.CalendarInterval $interval = $source; - |$targetVec.$set($idx, $interval.months, $interval.days, - | java.lang.Math.multiplyExact($interval.microseconds, 1000L));""".stripMargin) + |$targetVec.setIndexDefined($idx); + |$months.$set($idx, $interval.months); + |$days.$set($idx, $interval.days); + |$micros.$set($idx, $interval.microseconds);""".stripMargin) case dt if isTimeType(dt) => val set = if (nested) "setSafe" else "set" OutputEmit("", s"$targetVec.$set($idx, $source);") 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 a74a158e979..0c50613298b 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, Cast, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, TimestampAdd, TimestampDiff, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year} +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, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, TimestampAdd, TimestampDiff, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{CalendarIntervalType, DataType, DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String @@ -963,30 +963,14 @@ object CometMakeYMInterval extends CometCodegenDispatch[MakeYMInterval] object CometMakeDTInterval extends CometCodegenDispatch[MakeDTInterval] -object CometMakeInterval extends CometExpressionSerde[MakeInterval] with CodegenDispatchFallback { - private val incompatReason = - "The native implementation converts seconds to `Float64`, which can lose microsecond" + - " precision, and stores time in nanoseconds, which overflows for large time components" + - " (hours, minutes, seconds) that Spark can represent." - - override def getCompatibleNotes(): Seq[String] = Seq( - "Both the default JVM codegen-dispatch path and the native path currently limit the" + - " elapsed-time component to about 292 years in either direction. This only affects" + - " extreme intervals and is tracked in" + - " [#5279](https://github.com/apache/datafusion-comet/issues/5279).") - - override def getIncompatibleReasons(): Seq[String] = Seq(incompatReason) - - override def getSupportLevel(expr: MakeInterval): SupportLevel = - Incompatible(Some(incompatReason)) +object CometMakeInterval extends CometExpressionSerde[MakeInterval] { + override def getSupportLevel(expr: MakeInterval): SupportLevel = Compatible() override def convert( expr: MakeInterval, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { - // The explicit return type skips DataFusion's registry coercion, but its kernel needs Float64. - val children = expr.children.updated(6, Cast(expr.secs, DoubleType)) - val childExprs = children.map(exprToProtoInternal(_, inputs, binding)) + val childExprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) val optExpr = scalarFunctionExprToProtoWithReturnType( "make_interval", CalendarIntervalType, 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 ef541fb3e2d..791ae0a19e7 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 @@ -211,6 +211,8 @@ class CometScalaUDFCodegen extends CometUDF with Logging { case list: ListVector => val child = list.getDataVector ArrayColumnSpec(nullable = true, Utils.fromArrowField(child.getField), specFor(child)) + case struct: StructVector if Utils.isCalendarIntervalStructField(struct.getField) => + ScalarColumnSpec(classOf[StructVector], nullable = true) case struct: StructVector => val fieldSpecs = (0 until struct.size()).map { fi => val childVec = struct.getChildByOrdinal(fi).asInstanceOf[ValueVector] diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala index 3f47130e5ab..5104e92b43d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala @@ -25,6 +25,7 @@ import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector._ import org.apache.arrow.vector.complex._ import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow import org.apache.spark.sql.catalyst.expressions.SpecializedGetters import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.errors.QueryExecutionErrors @@ -87,6 +88,11 @@ private[arrow] object ArrowWriter { case (_: DayTimeIntervalType, vector: DurationVector) => new DurationWriter(vector) case (CalendarIntervalType, vector: IntervalMonthDayNanoVector) => new IntervalMonthDayNanoWriter(vector) + case (CalendarIntervalType, vector: StructVector) => + val children = (0 until vector.size()).map { ordinal => + createFieldWriter(vector.getChildByOrdinal(ordinal)) + } + new CalendarIntervalStructWriter(vector, children.toArray) case (dt, _) => throw QueryExecutionErrors.notSupportTypeError(dt) } @@ -470,6 +476,23 @@ private[arrow] class StructWriter( } } +private[arrow] class CalendarIntervalStructWriter( + valueVector: StructVector, + children: Array[ArrowFieldWriter]) + extends StructWriter(valueVector, children) { + + private val row = new GenericInternalRow(3) + + override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { + valueVector.setIndexDefined(count) + val interval = input.getInterval(ordinal) + row.update(0, interval.months) + row.update(1, interval.days) + row.update(2, interval.microseconds) + children.indices.foreach(i => children(i).write(row, i)) + } +} + private[arrow] class MapWriter( val valueVector: MapVector, val structVector: StructVector, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 9d4b0bce881..9ab0ac2f6d1 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -48,6 +48,8 @@ import org.apache.comet.shims.CometTypeShim import org.apache.comet.vector.CometVector object Utils extends CometTypeShim with Logging { + private val calendarIntervalStructKey = "SPARK::calendarInterval::struct" + def getConfPath(confFileName: String): String = { sys.env .get(COMET_CONF_DIR_ENV) @@ -77,6 +79,8 @@ object Utils extends CometTypeShim with Logging { val elementField = field.getChildren().get(0) val elementType = fromArrowField(elementField) ArrayType(elementType, containsNull = elementField.isNullable) + case ArrowType.Struct.INSTANCE if isCalendarIntervalStructField(field) => + CalendarIntervalType case ArrowType.Struct.INSTANCE => val fields = field.getChildren().asScala.map { child => val dt = fromArrowField(child) @@ -165,7 +169,7 @@ object Utils extends CometTypeShim with Logging { // Spark stores DayTimeIntervalType as microseconds in an int64, matching Arrow // Duration(Microsecond) rather than the lossy Interval(DayTime) {days, millis} layout. case _: DayTimeIntervalType => new ArrowType.Duration(TimeUnit.MICROSECOND) - case CalendarIntervalType => new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO) + case CalendarIntervalType => ArrowType.Struct.INSTANCE case _ => throw new UnsupportedOperationException( s"Unsupported data type: [${dt.getClass.getName}] ${dt.catalogString}") @@ -205,12 +209,46 @@ object Utils extends CometTypeShim with Logging { .add(MapVector.VALUE_NAME, valueType, nullable = valueContainsNull), nullable = false, timeZoneId)).asJava) + case CalendarIntervalType => + val fieldType = new FieldType(nullable, ArrowType.Struct.INSTANCE, null) + val monthsType = new FieldType( + false, + new ArrowType.Int(32, true), + null, + Map(calendarIntervalStructKey -> "true").asJava) + new Field( + name, + fieldType, + Seq( + new Field("months", monthsType, Seq.empty[Field].asJava), + new Field( + "days", + new FieldType(false, new ArrowType.Int(32, true), null), + Seq.empty[Field].asJava), + new Field( + "microseconds", + new FieldType(false, new ArrowType.Int(64, true), null), + Seq.empty[Field].asJava)).asJava) case dataType => val fieldType = new FieldType(nullable, toArrowType(dataType, timeZoneId), null) new Field(name, fieldType, Seq.empty[Field].asJava) } } + def isCalendarIntervalStructField(field: Field): Boolean = { + val children = field.getChildren + def child(index: Int, name: String, bits: Int): Boolean = { + val f = children.get(index) + f.getName == name && f.getType == new ArrowType.Int(bits, true) && !f.isNullable + } + field.getType == ArrowType.Struct.INSTANCE && + children.size == 3 && + child(0, "months", 32) && + child(1, "days", 32) && + child(2, "microseconds", 64) && + children.get(0).getMetadata.getOrDefault(calendarIntervalStructKey, "false") == "true" + } + /** * Maps schema from Spark to Arrow. NOTE: timeZoneId required for TimestampType in StructType */ diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql index 952a9085ee9..9226098e7a7 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql @@ -15,8 +15,6 @@ -- specific language governing permissions and limitations -- under the License. --- Config: spark.comet.expression.MakeInterval.allowIncompatible=true - statement CREATE TABLE test_make_interval( years int, @@ -29,6 +27,18 @@ CREATE TABLE test_make_interval( statement INSERT INTO test_make_interval VALUES + -- Adapted from Spark's MakeInterval expression tests: + -- https://github.com/apache/spark/blob/v4.2.0/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/IntervalExpressionsSuite.scala#L195-L231 + (0, 0, 0, 0, 0, 0, 0.000000), + (-123, 0, 0, 0, 0, 0, 0.000000), + (0, 0, 123, 0, 0, 0, 0.000000), + (0, 0, 0, 0, 0, 0, -0.123000), + (9999, 11, 0, 31, 23, 59, 59.999999), + (10000, 0, 0, 0, 0, 0, -0.000001), + (-9999, -11, 0, -31, -23, -59, -59.999999), + (-10000, 0, 0, 0, 0, 0, 0.000001), + (0, 0, 0, 0, 2147483647, 2147483647, 2149633277.790647), + (100, 11, 1, 1, 12, 30, 1.001001), (1, 2, 3, 4, 5, 6, 7.123456), (0, 1, 0, 1, 0, 0, 100.000001), (-1, -2, -1, -1, -1, -1, -1.500000), @@ -43,7 +53,28 @@ FROM test_make_interval ORDER BY years query -SELECT make_interval(1, 2), make_interval(3), make_interval() +-- Adapted from Spark's SQL and DataFrame API default-argument tests: +-- https://github.com/apache/spark/blob/v4.2.0/sql/core/src/test/resources/sql-tests/inputs/interval.sql#L81-L90 +-- https://github.com/apache/spark/blob/v4.2.0/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala#L1284-L1323 +SELECT make_interval(), + make_interval(1), + make_interval(1, 2), + make_interval(1, 2, 3), + make_interval(1, 2, 3, 4), + make_interval(1, 2, 3, 4, 5), + make_interval(1, 2, 3, 4, 5, 6), + make_interval(1, 2, 3, 4, 5, 6, 7.008009) + +query +SELECT make_interval(years), + make_interval(years, months), + make_interval(years, months, weeks), + make_interval(years, months, weeks, days), + make_interval(years, months, weeks, days, hours), + make_interval(years, months, weeks, days, hours, mins), + make_interval(years, months, weeks, days, hours, mins, secs) +FROM test_make_interval +WHERE years = 100 query SELECT make_interval(0, 1, 0, 1, 0, 0, 100.000001) @@ -51,16 +82,16 @@ SELECT make_interval(0, 1, 0, 1, 0, 0, 100.000001) query SELECT make_interval(2147483647) -query ignore(https://github.com/apache/datafusion-comet/issues/5131) +query SELECT make_interval(1, 2, 3, 4, 0, 0, 123456789012.123456) query SELECT make_interval(0, 0, 0, 0, 0, 0, 999999999.999999) -query ignore(https://github.com/apache/datafusion-comet/issues/5131) +query SELECT make_interval(0, 0, 0, 0, 0, 0, 999999999.000001) -query ignore(https://github.com/apache/datafusion-comet/issues/5131) +query SELECT make_interval(0, 0, 0, 0, 2562048) query diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql index 175b9ecd089..f3e2f8ca32a 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql @@ -15,27 +15,54 @@ -- specific language governing permissions and limitations -- under the License. --- Native ANSI execution must preserve Spark's overflow exception. -- Config: spark.sql.ansi.enabled=true --- Config: spark.comet.expression.MakeInterval.allowIncompatible=true statement -CREATE TABLE test_make_interval_ansi(years int) USING parquet +CREATE TABLE test_make_interval_ansi( + id int, + years int, + months int, + weeks int, + days int, + hours int, + mins int, + secs decimal(18, 6)) USING parquet statement -INSERT INTO test_make_interval_ansi VALUES (NULL) +INSERT INTO test_make_interval_ansi VALUES + -- Adapted from Spark's ANSI MakeInterval expression tests: + -- https://github.com/apache/spark/blob/v4.2.0/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/IntervalExpressionsSuite.scala#L233-L279 + (0, NULL, 0, 0, 0, 0, 0, 0.000000), + (1, 0, 0, 0, 0, 0, 0, 0.000000), + (2, -123, 0, 0, 0, 0, 0, 0.000000), + (3, 0, 0, 123, 0, 0, 0, 0.000000), + (4, 0, 0, 0, 0, 0, 0, -0.123000), + (5, 9999, 11, 0, 31, 23, 59, 59.999999), + (6, 10000, 0, 0, 0, 0, 0, -0.000001), + (7, -9999, -11, 0, -31, -23, -59, -59.999999), + (8, -10000, 0, 0, 0, 0, 0, 0.000001), + (9, 0, 0, 0, 0, 2147483647, 2147483647, 2149633277.790647), + (10, 2147483647, 0, 0, 0, 0, 0, 0.000000), + (11, 0, 0, 2147483647, 0, 0, 0, 0.000000) query SELECT make_interval(1, 2, 3, 4, 5, 6, 7.123456) query -SELECT make_interval(years) FROM test_make_interval_ansi +SELECT make_interval(years, months, weeks, days, hours, mins, secs) +FROM test_make_interval_ansi +WHERE id BETWEEN 0 AND 9 +ORDER BY id query expect_error(overflow. If necessary set) -SELECT make_interval(2147483647) +SELECT make_interval(years) +FROM test_make_interval_ansi +WHERE id = 10 query expect_error(overflow. If necessary set) -SELECT make_interval(0, 0, 2147483647) +SELECT make_interval(0, 0, weeks) +FROM test_make_interval_ansi +WHERE id = 11 -query ignore(https://github.com/apache/datafusion-comet/issues/5131) +query SELECT make_interval(0, 0, 0, 0, 2562048) diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch.sql index ee2d5e81605..57747625f29 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch.sql @@ -15,8 +15,6 @@ -- specific language governing permissions and limitations -- under the License. --- With allowIncompatible unset, MakeInterval uses Spark's JVM codegen dispatcher. - statement CREATE TABLE test_make_interval_dispatch( years int, @@ -43,7 +41,7 @@ FROM test_make_interval_dispatch WHERE hours != 2562048 ORDER BY years -query ignore(https://github.com/apache/datafusion-comet/issues/5279) +query SELECT make_interval(0, 0, 0, 0, hours) FROM test_make_interval_dispatch WHERE hours = 2562048 diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch_ansi.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch_ansi.sql index f199a8fb66a..1d9d885cad2 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch_ansi.sql @@ -42,7 +42,7 @@ SELECT make_interval(0, 0, weeks) FROM test_make_interval_dispatch_ansi WHERE weeks = 2147483647 -query ignore(https://github.com/apache/datafusion-comet/issues/5279) +query SELECT make_interval(0, 0, 0, 0, hours) FROM test_make_interval_dispatch_ansi WHERE hours = 2562048 diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/try_make_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/try_make_interval.sql index 5e481cb9ff0..2951a375bf6 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/try_make_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/try_make_interval.sql @@ -31,9 +31,21 @@ CREATE TABLE test_try_make_interval( statement INSERT INTO test_try_make_interval VALUES (1, 2, 3, 4, 5, 6, 7.123456), + (0, 0, 0, 0, 2562048, 0, 999999999.000001), (2147483647, 0, 0, 0, 0, 0, 0.000000) query SELECT try_make_interval(years, months, weeks, days, hours, mins, secs) FROM test_try_make_interval ORDER BY years + +query +-- Adapted from Spark's try_make_interval default-argument API tests: +-- https://github.com/apache/spark/blob/v4.2.0/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala#L2037-L2076 +SELECT try_make_interval(1), + try_make_interval(1, 2), + try_make_interval(1, 2, 3), + try_make_interval(1, 2, 3, 4), + try_make_interval(1, 2, 3, 4, 5), + try_make_interval(1, 2, 3, 4, 5, 6), + try_make_interval(1, 2, 3, 4, 5, 6, 7.008009) diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 5806cb35015..bc108faf192 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -26,6 +26,7 @@ import org.apache.spark.{SparkConf, SparkEnv, TaskContext} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.api.java.UDF1 import org.apache.spark.sql.catalyst.expressions.{BoundReference, CreateArray, CreateMap, CreateNamedStruct, Expression, Literal, MapConcat} +import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -82,18 +83,25 @@ class CometCodegenSuite } test("codegen kernel round-trips CalendarIntervalType") { - val input = new IntervalMonthDayNanoVector("in", CometArrowAllocator) + val input = Utils + .toArrowField("in", CalendarIntervalType, nullable = true, "UTC") + .createVector(CometArrowAllocator) + .asInstanceOf[org.apache.arrow.vector.complex.StructVector] val field = CometBatchKernelCodegen.toFfiArrowField("out", CalendarIntervalType, nullable = true) val output = CometBatchKernelCodegen.allocateOutput(field, 2, 0) try { input.allocateNew() - input.setSafe(0, 14, -3, 1234567000L) + input.setIndexDefined(0) + input.getChild("months").asInstanceOf[IntVector].setSafe(0, 14) + input.getChild("days").asInstanceOf[IntVector].setSafe(0, -3) + input.getChild("microseconds").asInstanceOf[BigIntVector].setSafe(0, Long.MaxValue) input.setNull(1) input.setValueCount(2) val expr = BoundReference(0, CalendarIntervalType, nullable = true) - val spec = ArrowColumnSpec(classOf[IntervalMonthDayNanoVector], nullable = true) + val spec = + ArrowColumnSpec(classOf[org.apache.arrow.vector.complex.StructVector], nullable = true) val kernel = CometBatchKernelCodegen.compile(expr, IndexedSeq(spec)).newInstance() kernel.init(0) kernel.process(Array(input), output, 2) @@ -103,7 +111,7 @@ class CometCodegenSuite val actual = comet.getInterval(0) assert(actual.months === 14) assert(actual.days === -3) - assert(actual.microseconds === 1234567L) + assert(actual.microseconds === Long.MaxValue) assert(comet.getInterval(1) == null) } finally { output.close() diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDatetimeExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDatetimeExpressionBenchmark.scala index a9d5b35e17a..ec558f9a7bf 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDatetimeExpressionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDatetimeExpressionBenchmark.scala @@ -165,18 +165,11 @@ object CometDatetimeExpressionBenchmark extends CometBenchmarkBase { consumeIntervals() } } - benchmark.addCase("Comet (codegen dispatch)") { _ => + benchmark.addCase("Comet") { _ => withSQLConf(cometConfigs.toSeq: _*) { consumeIntervals() } } - benchmark.addCase("Comet (native)") { _ => - val configs = - cometConfigs ++ Map(CometConf.getExprAllowIncompatConfigKey("MakeInterval") -> "true") - withSQLConf(configs.toSeq: _*) { - consumeIntervals() - } - } benchmark.run() } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index 50f723d0bfe..56938c939d0 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -25,7 +25,8 @@ import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers import org.apache.arrow.memory.RootAllocator -import org.apache.arrow.vector.{BigIntVector, IntervalMonthDayNanoVector, IntVector, VectorSchemaRoot} +import org.apache.arrow.vector.{BigIntVector, IntVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.StructVector import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} import org.apache.spark.sql.catalyst.expressions.GenericInternalRow import org.apache.spark.sql.comet.util.Utils @@ -61,19 +62,19 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { Utils.fromArrowField(field) shouldBe CalendarIntervalType val root = VectorSchemaRoot.create(new Schema(Seq(field).asJava), allocator) try { - val expected = new CalendarInterval(14, -3, 1234567L) + val expected = new CalendarInterval(14, -3, Long.MaxValue) val writer = ArrowWriter.create(root) writer.write(new GenericInternalRow(Array[Any](expected))) writer.write(new GenericInternalRow(Array[Any](null))) writer.finish() - val arrow = root.getVector(0).asInstanceOf[IntervalMonthDayNanoVector] - IntervalMonthDayNanoVector.getMonths(arrow.getDataBuffer, 0) shouldBe expected.months - IntervalMonthDayNanoVector.getDays(arrow.getDataBuffer, 0) shouldBe expected.days - IntervalMonthDayNanoVector.getNanoseconds(arrow.getDataBuffer, 0) shouldBe - expected.microseconds * 1000L + val arrow = root.getVector(0).asInstanceOf[StructVector] + arrow.getChild("months").asInstanceOf[IntVector].get(0) shouldBe expected.months + arrow.getChild("days").asInstanceOf[IntVector].get(0) shouldBe expected.days + arrow.getChild("microseconds").asInstanceOf[BigIntVector].get(0) shouldBe + expected.microseconds - val comet = new CometPlainVector(arrow, false) + val comet = CometVector.getVector(arrow, null) comet.getInterval(0) shouldBe expected comet.getInterval(1) shouldBe null } finally { From a0cb4ddbe917596e46549300019323452e7269cf Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 7 Aug 2026 23:14:18 +0800 Subject: [PATCH 02/20] scalafix --- .../org/apache/comet/vector/CometStructVector.java | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/spark/src/main/java/org/apache/comet/vector/CometStructVector.java b/spark/src/main/java/org/apache/comet/vector/CometStructVector.java index 25514f07f33..259793b831a 100644 --- a/spark/src/main/java/org/apache/comet/vector/CometStructVector.java +++ b/spark/src/main/java/org/apache/comet/vector/CometStructVector.java @@ -27,7 +27,6 @@ import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.util.TransferPair; import org.apache.spark.sql.vectorized.ColumnVector; -import org.apache.spark.unsafe.types.CalendarInterval; /** * A {@link CometDecodedVector} for Spark struct columns, wrapping an Arrow {@link StructVector} and @@ -59,15 +58,6 @@ public ColumnVector getChild(int i) { return children.get(i); } - @Override - public CalendarInterval getInterval(int rowId) { - if (isNullAt(rowId)) return null; - return new CalendarInterval( - children.get(0).getInt(rowId), - children.get(1).getInt(rowId), - children.get(2).getLong(rowId)); - } - @Override public CometVector slice(int offset, int length) { TransferPair tp = this.valueVector.getTransferPair(this.valueVector.getAllocator()); From 9b764d83eb28931ddd8a39f5b699baa526727364 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 8 Aug 2026 01:50:26 +0800 Subject: [PATCH 03/20] fix: use Arrow null check for struct vectors --- .../codegen/CometBatchKernelCodegenInput.scala | 7 +------ .../apache/comet/CometCodegenSourceSuite.scala | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 6 deletions(-) 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 69586898da4..964a7b1f644 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala @@ -110,12 +110,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim { if (!spec.nullable) { s" case $ord: return false;" } else { - // CometPlainVector exposes `isNullAt`; Arrow-typed fields expose `isNull`. Same semantics. - val method = spec.vectorClass match { - case cls if wrapsInCometPlainVector(cls) => "isNullAt" - case cls if cls == classOf[StructVector] => "isNullAt" - case _ => "isNull" - } + val method = nullCheckMethod(spec) s" case $ord: return this.col$ord.$method(this.rowIdx);" } } diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala index adc4440915b..f24a6abbecb 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala @@ -114,6 +114,24 @@ class CometCodegenSourceSuite extends AnyFunSuite { s"expected nullable isNullAt to delegate to the Arrow vector; got:\n$src") } + test("nullable struct delegates isNullAt to Arrow isNull") { + val structType = StructType(Seq(StructField("i", IntegerType, nullable = true))) + val structSpec = StructColumnSpec( + nullable = true, + fields = Seq( + StructFieldSpec( + "i", + IntegerType, + nullable = true, + ScalarColumnSpec( + CometBatchKernelCodegen.vectorClassBySimpleName("IntVector"), + nullable = true)))) + val src = gen(BoundReference(0, structType, nullable = true), structSpec) + assert( + src.contains("case 0: return this.col0.isNull(this.rowIdx);"), + s"expected nullable StructVector to use Arrow isNull; got:\n$src") + } + test("VarCharVector getUTF8String uses zero-copy fromAddress") { val expr = Length(BoundReference(0, StringType, nullable = true)) val src = gen(expr, nullableString) From ee5ea96a40ee430c11772b78fbb4bbcca13fcadf Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 13 Aug 2026 14:01:31 +0800 Subject: [PATCH 04/20] andy's review --- .../src/datetime_funcs/make_interval.rs | 94 ++++++++++++++----- .../apache/spark/sql/comet/util/Utils.scala | 4 +- .../expressions/datetime/make_interval.sql | 14 +++ .../datetime/make_interval_ansi.sql | 5 + .../datetime/make_interval_dispatch.sql | 47 ---------- .../datetime/make_interval_dispatch_ansi.sql | 48 ---------- .../spark/sql/comet/util/UtilsSuite.scala | 9 +- 7 files changed, 102 insertions(+), 119 deletions(-) delete mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch.sql delete mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch_ansi.sql diff --git a/native/spark-expr/src/datetime_funcs/make_interval.rs b/native/spark-expr/src/datetime_funcs/make_interval.rs index bb721442b6f..fc1a9d5f063 100644 --- a/native/spark-expr/src/datetime_funcs/make_interval.rs +++ b/native/spark-expr/src/datetime_funcs/make_interval.rs @@ -19,7 +19,7 @@ use crate::arithmetic_overflow_error; use arrow::array::{Array, ArrayRef, Decimal128Array, Int32Array, Int64Array, StructArray}; use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, Field, Fields}; -use datafusion::common::Result; +use datafusion::common::{DataFusionError, Result}; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; @@ -76,14 +76,29 @@ fn make_interval( hours: i32, minutes: i32, seconds_micros: i128, -) -> Option<(i32, i32, i64)> { - let months = years.checked_mul(12)?.checked_add(months)?; - let days = weeks.checked_mul(7)?.checked_add(days)?; +) -> std::result::Result<(i32, i32, i64), &'static str> { + let months = years + .checked_mul(12) + .and_then(|years| years.checked_add(months)) + .ok_or("integer")?; + let days = weeks + .checked_mul(7) + .and_then(|weeks| weeks.checked_add(days)) + .ok_or("integer")?; let micros = i64::try_from(seconds_micros) - .ok()? - .checked_add(i64::from(hours).checked_mul(MICROS_PER_HOUR)?)? - .checked_add(i64::from(minutes).checked_mul(MICROS_PER_MINUTE)?)?; - Some((months, days, micros)) + .ok() + .and_then(|seconds| { + i64::from(hours) + .checked_mul(MICROS_PER_HOUR) + .and_then(|hours| seconds.checked_add(hours)) + }) + .and_then(|micros| { + i64::from(minutes) + .checked_mul(MICROS_PER_MINUTE) + .and_then(|minutes| micros.checked_add(minutes)) + }) + .ok_or("long")?; + Ok((months, days, micros)) } impl ScalarUDFImpl for SparkMakeInterval { @@ -106,16 +121,48 @@ impl ScalarUDFImpl for SparkMakeInterval { .into_iter() .map(|arg| arg.into_array(number_rows)) .collect::>>()?; - let years = arrays[0].as_any().downcast_ref::().unwrap(); - let months = arrays[1].as_any().downcast_ref::().unwrap(); - let weeks = arrays[2].as_any().downcast_ref::().unwrap(); - let days = arrays[3].as_any().downcast_ref::().unwrap(); - let hours = arrays[4].as_any().downcast_ref::().unwrap(); - let minutes = arrays[5].as_any().downcast_ref::().unwrap(); + let years = arrays[0] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution("make_interval: expected Int32 years".into()) + })?; + let months = arrays[1] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution("make_interval: expected Int32 months".into()) + })?; + let weeks = arrays[2] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution("make_interval: expected Int32 weeks".into()) + })?; + let days = arrays[3] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution("make_interval: expected Int32 days".into()) + })?; + let hours = arrays[4] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution("make_interval: expected Int32 hours".into()) + })?; + let minutes = arrays[5] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution("make_interval: expected Int32 minutes".into()) + })?; let seconds = arrays[6] .as_any() .downcast_ref::() - .unwrap(); + .ok_or_else(|| { + DataFusionError::Execution("make_interval: expected Decimal128 seconds".into()) + })?; let mut result_months = Vec::with_capacity(years.len()); let mut result_days = Vec::with_capacity(years.len()); @@ -140,16 +187,16 @@ impl ScalarUDFImpl for SparkMakeInterval { minutes.value(i), seconds.value(i), ) { - Some((months, days, micros)) => { + Ok((months, days, micros)) => { result_months.push(months); result_days.push(days); result_micros.push(micros); valid.push(true); } - None if self.fail_on_error => { - return Err(arithmetic_overflow_error("interval").into()); + Err(from_type) if self.fail_on_error => { + return Err(arithmetic_overflow_error(from_type).into()); } - None => { + Err(_) => { result_months.push(0); result_days.push(0); result_micros.push(0); @@ -182,9 +229,12 @@ mod tests { fn preserves_spark_microsecond_range_and_overflow() { assert_eq!( make_interval(1, 2, 3, 4, 2_562_048, 0, 123_456_789_012_123_456), - Some((14, 25, 132_680_161_812_123_456)) + Ok((14, 25, 132_680_161_812_123_456)) + ); + assert_eq!(make_interval(i32::MAX, 0, 0, 0, 0, 0, 0), Err("integer")); + assert_eq!( + make_interval(0, 0, 0, 0, i32::MAX, i32::MAX, i128::MAX), + Err("long") ); - assert!(make_interval(i32::MAX, 0, 0, 0, 0, 0, 0).is_none()); - assert!(make_interval(0, 0, 0, 0, i32::MAX, i32::MAX, i128::MAX).is_none()); } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 9ab0ac2f6d1..02f24fe1f3c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -169,7 +169,9 @@ object Utils extends CometTypeShim with Logging { // Spark stores DayTimeIntervalType as microseconds in an int64, matching Arrow // Duration(Microsecond) rather than the lossy Interval(DayTime) {days, millis} layout. case _: DayTimeIntervalType => new ArrowType.Duration(TimeUnit.MICROSECOND) - case CalendarIntervalType => ArrowType.Struct.INSTANCE + case CalendarIntervalType => + throw new UnsupportedOperationException( + "CalendarIntervalType requires toArrowField to preserve struct fields and metadata") case _ => throw new UnsupportedOperationException( s"Unsupported data type: [${dt.getClass.getName}] ${dt.catalogString}") diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql index 9226098e7a7..bf0afc93ecb 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql @@ -15,6 +15,8 @@ -- specific language governing permissions and limitations -- under the License. +-- Config: spark.comet.shuffle.mode=native + statement CREATE TABLE test_make_interval( years int, @@ -45,6 +47,7 @@ INSERT INTO test_make_interval VALUES (NULL, 1, 2, 3, 4, 5, 6.000000), (2, NULL, 2, 3, 4, 5, 6.000000), (3, 1, 2, 3, 4, 5, NULL), + (-178956970, -8, -306783378, -2, -2147483648, -2147483648, -999999999999.999999), (-2147483648, 0, 0, 0, 0, 0, 0.000000) query @@ -76,6 +79,17 @@ SELECT make_interval(years), FROM test_make_interval WHERE years = 100 +query +-- Build ARRAY from native make_interval below a native shuffle, then consume its +-- interval child with native GetArrayItem above the shuffle. This pins the CalendarInterval +-- metadata across both native expression and shuffle boundaries. +SELECT intervals, intervals[0] +FROM ( + SELECT years, array(make_interval(years, months, weeks, days, hours, mins, secs)) AS intervals + FROM test_make_interval + DISTRIBUTE BY years +) + query SELECT make_interval(0, 1, 0, 1, 0, 0, 100.000001) diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql index f3e2f8ca32a..3add6f52759 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql @@ -54,6 +54,11 @@ FROM test_make_interval_ansi WHERE id BETWEEN 0 AND 9 ORDER BY id +-- Spark <= 4.1 expects the literal "integer overflow": +-- https://github.com/apache/spark/blob/v4.1.3/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala#L704-L713 +-- Spark 4.2 expects the literal "overflow": +-- https://github.com/apache/spark/blob/v4.2.0/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala#L703-L712 +-- Match the common text across versions. query expect_error(overflow. If necessary set) SELECT make_interval(years) FROM test_make_interval_ansi diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch.sql deleted file mode 100644 index 57747625f29..00000000000 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch.sql +++ /dev/null @@ -1,47 +0,0 @@ --- 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. - -statement -CREATE TABLE test_make_interval_dispatch( - years int, - months int, - weeks int, - days int, - hours int, - mins int, - secs decimal(18, 6)) USING parquet - -statement -INSERT INTO test_make_interval_dispatch VALUES - (1, 2, 3, 4, 5, 6, 7.123456), - (0, 1, 0, 1, 0, 0, 100.000001), - (-1, -2, -1, -1, -1, -1, -1.500000), - (NULL, 1, 2, 3, 4, 5, 6.000000), - (2, NULL, 2, 3, 4, 5, 6.000000), - (3, 1, 2, 3, 4, 5, NULL), - (0, 0, 0, 0, 2562048, 0, 0.000000) - -query -SELECT make_interval(years, months, weeks, days, hours, mins, secs) -FROM test_make_interval_dispatch -WHERE hours != 2562048 -ORDER BY years - -query -SELECT make_interval(0, 0, 0, 0, hours) -FROM test_make_interval_dispatch -WHERE hours = 2562048 diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch_ansi.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch_ansi.sql deleted file mode 100644 index 1d9d885cad2..00000000000 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch_ansi.sql +++ /dev/null @@ -1,48 +0,0 @@ --- 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. - --- Config: spark.sql.ansi.enabled=true - -statement -CREATE TABLE test_make_interval_dispatch_ansi(years int, weeks int, hours int) USING parquet - -statement -INSERT INTO test_make_interval_dispatch_ansi VALUES - (1, 0, 5), - (2147483647, 0, 0), - (0, 2147483647, 0), - (0, 0, 2562048) - -query -SELECT make_interval(years, 0, weeks, 0, hours) -FROM test_make_interval_dispatch_ansi -WHERE years = 1 - -query expect_error(overflow. If necessary set) -SELECT make_interval(years) -FROM test_make_interval_dispatch_ansi -WHERE years = 2147483647 - -query expect_error(overflow. If necessary set) -SELECT make_interval(0, 0, weeks) -FROM test_make_interval_dispatch_ansi -WHERE weeks = 2147483647 - -query -SELECT make_interval(0, 0, 0, 0, hours) -FROM test_make_interval_dispatch_ansi -WHERE hours = 2562048 diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index c3b00a2814c..a1288897bfe 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -21,11 +21,18 @@ package org.apache.spark.sql.comet.util import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.vectorized.ConstantColumnVector -import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType, TimestampType} +import org.apache.spark.sql.types.{CalendarIntervalType, IntegerType, StringType, StructField, StructType, TimestampType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} class UtilsSuite extends CometTestBase { + test("CalendarIntervalType requires field-level Arrow conversion") { + val error = intercept[UnsupportedOperationException] { + Utils.toArrowType(CalendarIntervalType, "UTC") + } + assert(error.getMessage.contains("requires toArrowField")) + } + test("serializeBatches preserves row count for a zero-column batch") { val numRows = 5 val batch = new ColumnarBatch(Array.empty[ColumnVector], numRows) From c0f2f9b692c9069c65d1de5b1cf600bf2402cd9e Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 14 Aug 2026 13:01:25 +0800 Subject: [PATCH 05/20] Update make_interval_ansi.sql --- .../expressions/datetime/make_interval_ansi.sql | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql index 3add6f52759..90fdb0ab046 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql @@ -54,11 +54,12 @@ FROM test_make_interval_ansi WHERE id BETWEEN 0 AND 9 ORDER BY id --- Spark <= 4.1 expects the literal "integer overflow": +-- Spark <= 4.1 preserves the JDK literal "integer overflow": -- https://github.com/apache/spark/blob/v4.1.3/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala#L704-L713 --- Spark 4.2 expects the literal "overflow": --- https://github.com/apache/spark/blob/v4.2.0/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala#L703-L712 --- Match the common text across versions. +-- SPARK-55714 canonicalizes simple " overflow" messages to "overflow" in Spark 4.2: +-- https://github.com/apache/spark/blob/v4.2.0/sql/api/src/main/scala/org/apache/spark/sql/errors/ExecutionErrors.scala#L118-L144 +-- The native unit test pins the "integer"/"long" overflow type; this assertion matches the +-- common rendered suffix across supported Spark versions. query expect_error(overflow. If necessary set) SELECT make_interval(years) FROM test_make_interval_ansi From eb4e1cec6c7b7cfe8eac0daa41f1d1444a1bd846 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Mon, 24 Aug 2026 16:08:49 +0800 Subject: [PATCH 06/20] fix scalalint --- .../spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index efc15a36445..505e0b568c9 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -25,7 +25,7 @@ import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers import org.apache.arrow.memory.{AllocationListener, RootAllocator} -import org.apache.arrow.vector.{BaseValueVector, BigIntVector, BitVector, DecimalVector, IntervalMonthDayNanoVector, IntVector, VectorSchemaRoot} +import org.apache.arrow.vector.{BaseValueVector, BigIntVector, BitVector, DecimalVector, IntVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.StructVector import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} import org.apache.spark.sql.catalyst.expressions.GenericInternalRow From 8b61b8a1fefef06f0ba0ab2ebbae088d3737c0fe Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 28 Aug 2026 00:14:38 +0800 Subject: [PATCH 07/20] docs: explain metadata-marker requirement in isCalendarIntervalStructField Co-Authored-By: Claude Fable 5 --- .../scala/org/apache/spark/sql/comet/util/Utils.scala | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 02f24fe1f3c..4540337e06a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -237,6 +237,15 @@ object Utils extends CometTypeShim with Logging { } } + /** + * Returns true only for the Spark-tagged `CalendarIntervalType` struct produced by + * [[toArrowField]]. Matching the field shape (`months: Int32`, `days: Int32`, `microseconds: + * Int64`, all non-nullable) is not sufficient: a user-defined `struct` has the same shape, so this also requires the + * `SPARK::calendarInterval::struct` metadata marker on the `months` child. Only fields carrying + * the marker round-trip back to `CalendarIntervalType`; unmarked structs stay plain + * `StructType`s. + */ def isCalendarIntervalStructField(field: Field): Boolean = { val children = field.getChildren def child(index: Int, name: String, bits: Int): Boolean = { From fef4551ad2de77c7a2bbae8e1640a1fe9bbedf64 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 4 Sep 2026 22:48:52 +0800 Subject: [PATCH 08/20] fix: keep datafusion-spark for benchmarks --- native/Cargo.lock | 1 + native/spark-expr/Cargo.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/native/Cargo.lock b/native/Cargo.lock index fa630bc0c81..17945a0e21a 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2075,6 +2075,7 @@ dependencies = [ "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", + "datafusion-spark", "futures", "jni 0.22.4", "num", diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 12c2d820c9f..b210b8ccdee 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -50,6 +50,7 @@ uuid = "1.23.3" [dev-dependencies] arrow = {workspace = true} criterion = { version = "0.7", features = ["async", "async_tokio", "async_std"] } +datafusion-spark = { workspace = true } rand = { workspace = true} tokio = { version = "1", features = ["rt-multi-thread"] } datafusion = { workspace = true, features = ["sql"] } From b55bde0b552fccc5d45b0bd890974a0550d2734c Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 13 Sep 2026 13:46:07 +0800 Subject: [PATCH 09/20] fix: restore datafusion-spark as a spark-expr library dependency This branch moved `datafusion-spark` to `[dev-dependencies]` once the native `make_interval` stopped using `SparkMakeInterval`, keeping it only for the timestamp benches. Since then `main` added library uses of the crate in `string_funcs/concat_ws.rs` (#5725) and `agg_funcs/collect.rs`, so the merged crate failed to compile with `E0433: cannot find module or crate datafusion_spark`, breaking every native build, rust-test, the Spark SQL and Iceberg builds, the benchmark check and the Delta build gate. Move it back under `[dependencies]`, where `main` keeps it; benches still see it, so the dev-dependency entry is dropped. `Cargo.lock` is unchanged because it already recorded the dependency. Co-Authored-By: Claude Opus 5 --- native/spark-expr/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index b37720bc836..f00f1a7c495 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -30,6 +30,7 @@ edition = { workspace = true } arrow = { workspace = true } chrono = { workspace = true } datafusion = { workspace = true } +datafusion-spark = { workspace = true } num = { workspace = true } regex = { workspace = true } # preserve_order: needed for get_json_object to match Spark's JSON key ordering @@ -50,7 +51,6 @@ uuid = "1.23.3" [dev-dependencies] arrow = {workspace = true} criterion = { version = "0.7", features = ["async", "async_tokio", "async_std"] } -datafusion-spark = { workspace = true } rand = { workspace = true} tokio = { version = "1", features = ["rt-multi-thread"] } datafusion = { workspace = true, features = ["sql"] } From 99087b8e3c866fff0ba4db95f56bfc2ea63802f2 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 15 Sep 2026 22:59:44 +0800 Subject: [PATCH 10/20] fix: update native make_interval routing expectations and lint --- .../sql-tests/expressions/math/routing_intervals_disabled.sql | 2 +- .../sql-tests/expressions/math/routing_intervals_enabled.sql | 2 +- .../spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_disabled.sql b/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_disabled.sql index e9e0e2728f7..98e745b5e81 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_disabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_disabled.sql @@ -30,5 +30,5 @@ SELECT abs(i) FROM routing_intervals query expect_fallback(abs: spark.comet.exec.scalaUDF.codegen.enabled=false) SELECT abs(make_dt_interval(i)), abs(make_ym_interval(i)) FROM routing_intervals -query expect_fallback(make_interval: spark.comet.exec.scalaUDF.codegen.enabled=false) +query expect_native(make_interval) SELECT make_interval(i, 0, 0, 0, 0, 0, 0) FROM routing_intervals diff --git a/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_enabled.sql b/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_enabled.sql index 6f8dad1b404..d509cbc80b5 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_enabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_enabled.sql @@ -30,5 +30,5 @@ SELECT abs(i) FROM routing_intervals query expect_dispatch(abs) SELECT abs(make_dt_interval(i)), abs(make_ym_interval(i)) FROM routing_intervals -query expect_dispatch(make_interval) +query expect_native(make_interval) SELECT make_interval(i, 0, 0, 0, 0, 0, 0) FROM routing_intervals diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index 30517755420..939bec0e9b4 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -29,7 +29,7 @@ import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers import org.apache.arrow.memory.{AllocationListener, RootAllocator} -import org.apache.arrow.vector.{BaseFixedWidthVector, BaseValueVector, BigIntVector, BitVector, DecimalVector, IntervalMonthDayNanoVector, IntVector, VarCharVector, VectorSchemaRoot} +import org.apache.arrow.vector.{BaseFixedWidthVector, BaseValueVector, BigIntVector, BitVector, DecimalVector, IntVector, VarCharVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.StructVector import org.apache.arrow.vector.dictionary.{Dictionary => ArrowDictionary} import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider From 6742caab46184599fc19f22361f41e675998740d Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 29 Sep 2026 02:07:38 +0800 Subject: [PATCH 11/20] fix: hash CalendarInterval values like Spark in the native hash kernels The tagged months/days/microseconds struct that now carries a Spark CalendarInterval reached the generic struct arm of the native hash and xxhash64 kernels, which hash its three children in field order. Spark hashes an interval as hashInt(months, hashLong(microseconds, seed)), so the results were silently wrong (calendar_interval.sql failed in the Spark 4.1 expressions job). Recognize an interval by its SPARK::calendarInterval::struct marker rather than its shape, and hash it with Spark's formula ahead of the generic struct arm. Nested arrays, structs and maps reach it through the recursive hash. The Interval(MonthDayNano) arm is removed: nothing produces that type for CalendarIntervalType any more, so an unexpected one fails instead of being hashed as an interval. --- .../src/datetime_funcs/make_interval.rs | 45 +++++++ native/spark-expr/src/datetime_funcs/mod.rs | 2 +- native/spark-expr/src/hash_funcs/murmur3.rs | 94 ++++++++++++++ native/spark-expr/src/hash_funcs/utils.rs | 117 +++++++++++++----- native/spark-expr/src/hash_funcs/xxhash64.rs | 71 +++++++++++ native/spark-expr/src/lib.rs | 10 +- .../datetime/calendar_interval.sql | 20 +++ 7 files changed, 323 insertions(+), 36 deletions(-) diff --git a/native/spark-expr/src/datetime_funcs/make_interval.rs b/native/spark-expr/src/datetime_funcs/make_interval.rs index fc1a9d5f063..2885f74aa1a 100644 --- a/native/spark-expr/src/datetime_funcs/make_interval.rs +++ b/native/spark-expr/src/datetime_funcs/make_interval.rs @@ -42,6 +42,29 @@ pub fn calendar_interval_type() -> DataType { ])) } +/// Whether `fields` are the children of a Spark `CalendarIntervalType` value, i.e. of the tagged +/// struct built by [`calendar_interval_type`]. +/// +/// The `SPARK::calendarInterval::struct` marker on `months` is what identifies an interval: a +/// user's own `struct` has the same shape and must keep +/// ordinary struct semantics. Child nullability is deliberately not compared: the struct's own +/// validity decides whether an interval is null, and a nullability widening that keeps field +/// metadata must not turn an interval into a plain struct, which for the hash kernels would mean +/// silently hashing it as one. The JVM counterpart is `Utils.isCalendarIntervalStructField`. +pub fn is_calendar_interval_fields(fields: &Fields) -> bool { + let child = |index: usize, name: &str, data_type: &DataType| { + fields[index].name() == name && fields[index].data_type() == data_type + }; + fields.len() == 3 + && child(0, "months", &DataType::Int32) + && child(1, "days", &DataType::Int32) + && child(2, "microseconds", &DataType::Int64) + && fields[0] + .metadata() + .get(CALENDAR_INTERVAL_STRUCT_KEY) + .is_some_and(|value| value == "true") +} + #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkMakeInterval { signature: Signature, @@ -237,4 +260,26 @@ mod tests { Err("long") ); } + + #[test] + fn recognizes_tagged_interval_fields_by_marker() { + let DataType::Struct(fields) = calendar_interval_type() else { + unreachable!() + }; + assert!(is_calendar_interval_fields(&fields)); + + // Children widened to nullable with the metadata kept are still an interval. + let widened: Fields = fields + .iter() + .map(|f| Arc::new(f.as_ref().clone().with_nullable(true))) + .collect(); + assert!(is_calendar_interval_fields(&widened)); + + // The same shape without the marker is an ordinary user struct. + let unmarked: Fields = fields + .iter() + .map(|f| Arc::new(f.as_ref().clone().with_metadata(HashMap::new()))) + .collect(); + assert!(!is_calendar_interval_fields(&unmarked)); + } } diff --git a/native/spark-expr/src/datetime_funcs/mod.rs b/native/spark-expr/src/datetime_funcs/mod.rs index 5501dd0db6c..ac58d9b3631 100644 --- a/native/spark-expr/src/datetime_funcs/mod.rs +++ b/native/spark-expr/src/datetime_funcs/mod.rs @@ -42,7 +42,7 @@ pub use extract_date_part::SparkMinute; pub use extract_date_part::SparkSecond; pub use hours::SparkHoursTransform; pub use make_date::SparkMakeDate; -pub use make_interval::{calendar_interval_type, SparkMakeInterval}; +pub use make_interval::{calendar_interval_type, is_calendar_interval_fields, SparkMakeInterval}; pub use make_time::SparkMakeTime; pub use next_day::SparkNextDay; pub use seconds_to_timestamp::SparkSecondsToTimestamp; diff --git a/native/spark-expr/src/hash_funcs/murmur3.rs b/native/spark-expr/src/hash_funcs/murmur3.rs index 29a4820ca52..2c75344663d 100644 --- a/native/spark-expr/src/hash_funcs/murmur3.rs +++ b/native/spark-expr/src/hash_funcs/murmur3.rs @@ -391,6 +391,100 @@ mod tests { ); } + /// Spark hashes a `CalendarInterval` as `hashInt(months, hashLong(microseconds, seed))`, not + /// as a struct of its fields, so the tagged interval struct must not reach the generic struct + /// arm. That holds bare, nested in an array or struct, and with the children widened to + /// nullable. + #[test] + fn test_calendar_interval_uses_spark_interval_hash() { + use crate::hash_funcs::utils::test_utils::calendar_interval_array; + use crate::murmur3::{spark_compatible_murmur3_hash, spark_murmur3_hash}; + use arrow::array::{Array, ListArray, StructArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field, Fields}; + use datafusion::common::ScalarValue; + use datafusion::physical_plan::ColumnarValue; + + // `make_interval(1)`, a value that differs only in days, one with every field set, and NULL. + let intervals = [ + Some((12, 0, 0)), + Some((12, 5, 0)), + Some((14, -3, 7_008_009)), + None, + ]; + let interval_hash = |months: i32, micros: i64| { + spark_compatible_murmur3_hash( + months.to_le_bytes(), + spark_compatible_murmur3_hash(micros.to_le_bytes(), 42), + ) + }; + let expected: Vec = intervals + .iter() + .map(|v| v.map_or(42, |(months, _, micros)| interval_hash(months, micros))) + .collect(); + // Spark 4.1.3 returns -351543533 for `hash(make_interval(y))` with y = 1. + assert_eq!(expected[0] as i32, -351543533); + + for nullable_children in [false, true] { + let array = calendar_interval_array(&intervals, nullable_children); + assert_eq!( + hash_of(Arc::clone(&array), intervals.len()), + expected, + "nullable_children = {nullable_children}" + ); + + // The expression entry point, as `hash(c)` reaches it. + let result = spark_murmur3_hash(&[ + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Int32(Some(42))), + ]) + .unwrap(); + let ColumnarValue::Array(result) = result else { + panic!("expected an array result") + }; + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.value(0), -351543533); + + // `array(c)` and `struct(c)` hash a non-null interval exactly as the bare value, and + // skip a null one, as Spark's generated code does. + let list: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", array.data_type().clone(), true)), + OffsetBuffer::from_lengths(vec![1; intervals.len()]), + Arc::clone(&array), + None, + )); + assert_eq!(hash_of(list, intervals.len()), expected); + let wrapper_fields: Fields = + vec![Arc::new(Field::new("c", array.data_type().clone(), true))].into(); + let wrapped: ArrayRef = Arc::new(StructArray::new( + wrapper_fields, + vec![Arc::clone(&array)], + None, + )); + assert_eq!(hash_of(wrapped, intervals.len()), expected); + } + + // A user's struct with the same shape but without the marker keeps struct semantics: + // every field is hashed in order, days included. + let tagged = calendar_interval_array(&intervals[..3], false); + let tagged = tagged.as_any().downcast_ref::().unwrap(); + let plain_fields: Fields = vec![ + Arc::new(Field::new("months", DataType::Int32, false)), + Arc::new(Field::new("days", DataType::Int32, false)), + Arc::new(Field::new("microseconds", DataType::Int64, false)), + ] + .into(); + let plain: ArrayRef = Arc::new(StructArray::new( + plain_fields, + tagged.columns().to_vec(), + None, + )); + let mut field_by_field = vec![42u32; 3]; + create_murmur3_hashes(tagged.columns(), &mut field_by_field).unwrap(); + assert_eq!(hash_of(plain, 3), field_by_field); + assert_ne!(field_by_field, expected[..3]); + } + /// One `struct` element: `None` is a null struct, and the fields are /// independently nullable. type StructElem = Option<(Option, Option<&'static str>)>; diff --git a/native/spark-expr/src/hash_funcs/utils.rs b/native/spark-expr/src/hash_funcs/utils.rs index 7411b8dbfe5..da9c71dfdd0 100644 --- a/native/spark-expr/src/hash_funcs/utils.rs +++ b/native/spark-expr/src/hash_funcs/utils.rs @@ -157,47 +157,68 @@ macro_rules! hash_array_primitive_float { }; } +/// Hashes a Spark `CalendarIntervalType` column, which Comet carries as the tagged +/// `struct` built by +/// `calendar_interval_type`. `create_hashes_internal!` must route such a column here before its +/// generic struct arm: Spark hashes an interval with its own formula, not as a struct of fields. +/// +/// Only the struct's own validity is consulted. A null interval leaves the seed untouched, and a +/// valid one always has its children populated by every producer (the native `make_interval` +/// kernel, the JVM Arrow writer and the codegen dispatcher), whether or not the children are +/// declared nullable. #[macro_export] -macro_rules! hash_array_interval_month_day_nano { +macro_rules! hash_array_calendar_interval { ($column: ident, $hashes: ident, $hash_method: ident) => { let array = $column .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap_or_else(|| { panic!( "Failed to downcast column to {}. Actual data type: {:?}.", - stringify!(IntervalMonthDayNanoArray), + stringify!(StructArray), $column.data_type() ) }); + let months = array + .column(0) + .as_any() + .downcast_ref::() + .unwrap_or_else(|| { + panic!( + "Failed to downcast interval months to Int32Array. Actual data type: {:?}.", + array.column(0).data_type() + ) + }); + let microseconds = array + .column(2) + .as_any() + .downcast_ref::() + .unwrap_or_else(|| { + panic!( + "Failed to downcast interval microseconds to Int64Array. Actual data type: {:?}.", + array.column(2).data_type() + ) + }); - // The `nanoseconds / 1_000` below is exact: Spark's `CalendarInterval` is - // microsecond-based, and both JVM-to-Arrow producers - // (`ArrowWriters.CalendarIntervalWriter` and the codegen dispatch kernel) convert - // with `Math.multiplyExact(microseconds, 1000L)`, so the nanoseconds field is - // always an exact multiple of 1000 and out-of-range intervals throw at - // conversion time instead of reaching this hasher. + // Match Spark 4.2 generated code, which hashes microseconds and then months and omits + // the days field: + // https://github.com/apache/spark/blob/v4.2.0/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/hash.scala#L428-L431 + // Nested arrays, structs and maps reach this arm through the recursive hash, which is + // also how Spark's generated code hashes a nested interval. SPARK-58236 includes days + // starting in Spark 4.3; the version switch for that is tracked in + // https://github.com/apache/datafusion-comet/issues/5498. if array.null_count() == 0 { // Fast path: no nulls, use direct indexing for i in 0..$hashes.len() { - let value = array.value(i); - // Match Spark 4.2 generated code, which omits the days field: - // https://github.com/apache/spark/blob/v4.2.0/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/hash.scala#L428-L431 - // SPARK-58236 includes days starting in Spark 4.3; the version - // switch for that is tracked in - // https://github.com/apache/datafusion-comet/issues/5498. - $hashes[i] = - $hash_method((value.nanoseconds / 1_000).to_le_bytes(), $hashes[i]); - $hashes[i] = $hash_method(value.months.to_le_bytes(), $hashes[i]); + $hashes[i] = $hash_method(microseconds.value(i).to_le_bytes(), $hashes[i]); + $hashes[i] = $hash_method(months.value(i).to_le_bytes(), $hashes[i]); } } else { // Slow path: check nulls for i in 0..$hashes.len() { if !array.is_null(i) { - let value = array.value(i); - $hashes[i] = - $hash_method((value.nanoseconds / 1_000).to_le_bytes(), $hashes[i]); - $hashes[i] = $hash_method(value.months.to_le_bytes(), $hashes[i]); + $hashes[i] = $hash_method(microseconds.value(i).to_le_bytes(), $hashes[i]); + $hashes[i] = $hash_method(months.value(i).to_le_bytes(), $hashes[i]); } } } @@ -836,7 +857,7 @@ macro_rules! hash_list_array { #[macro_export] macro_rules! create_hashes_internal { ($arrays: ident, $hashes_buffer: ident, $hash_method: ident, $create_dictionary_hash_method: ident, $recursive_hash_method: ident) => { - use arrow::datatypes::{DataType, IntervalUnit, TimeUnit}; + use arrow::datatypes::{DataType, TimeUnit}; use arrow::array::{types::*, *}; use datafusion_comet_common::children_with_parent_nulls; @@ -980,13 +1001,6 @@ macro_rules! create_hashes_internal { $hash_method ); } - DataType::Interval(IntervalUnit::MonthDayNano) => { - $crate::hash_array_interval_month_day_nano!( - col, - $hashes_buffer, - $hash_method - ); - } DataType::Utf8 => { $crate::hash_array!(StringArray, col, $hashes_buffer, $hash_method); } @@ -1098,6 +1112,12 @@ macro_rules! create_hashes_internal { $crate::hash_list_with_primitive_elements!(fixed_size: list_array, values, list_size, field, $hashes_buffer, $hash_method, $recursive_hash_method); } + // A Spark `CalendarIntervalType` value is a tagged struct, recognised by its field + // marker rather than its shape. Test it before the generic struct arm, which would + // hash its three children in field order. + DataType::Struct(fields) if $crate::is_calendar_interval_fields(fields) => { + $crate::hash_array_calendar_interval!(col, $hashes_buffer, $hash_method); + } DataType::Struct(_) => { let struct_array = col.as_any().downcast_ref::().unwrap(); // Hash each field of the struct - Spark hashes all fields recursively, and a @@ -1336,4 +1356,41 @@ pub(crate) mod test_utils { ); }; } + + /// Builds a Spark `CalendarIntervalType` column in Comet's tagged-struct layout, one + /// `(months, days, microseconds)` per row. `nullable_children` widens the three children to + /// nullable, keeping the marker. + #[cfg(test)] + pub(crate) fn calendar_interval_array( + values: &[Option<(i32, i32, i64)>], + nullable_children: bool, + ) -> arrow::array::ArrayRef { + use crate::calendar_interval_type; + use arrow::array::{Int32Array, Int64Array, StructArray}; + use arrow::buffer::NullBuffer; + use arrow::datatypes::{DataType, Fields}; + use std::sync::Arc; + + let DataType::Struct(fields) = calendar_interval_type() else { + unreachable!("calendar_interval_type is a struct") + }; + let fields: Fields = if nullable_children { + fields + .iter() + .map(|f| Arc::new(f.as_ref().clone().with_nullable(true))) + .collect() + } else { + fields + }; + let value = |v: &Option<(i32, i32, i64)>| v.unwrap_or_default(); + let months = Int32Array::from_iter_values(values.iter().map(|v| value(v).0)); + let days = Int32Array::from_iter_values(values.iter().map(|v| value(v).1)); + let micros = Int64Array::from_iter_values(values.iter().map(|v| value(v).2)); + let nulls = NullBuffer::from_iter(values.iter().map(Option::is_some)); + Arc::new(StructArray::new( + fields, + vec![Arc::new(months), Arc::new(days), Arc::new(micros)], + Some(nulls), + )) + } } diff --git a/native/spark-expr/src/hash_funcs/xxhash64.rs b/native/spark-expr/src/hash_funcs/xxhash64.rs index bdda7c85f1d..37cee5fd844 100644 --- a/native/spark-expr/src/hash_funcs/xxhash64.rs +++ b/native/spark-expr/src/hash_funcs/xxhash64.rs @@ -351,6 +351,77 @@ mod tests { assert_eq!(from_hidden, from_null); } + /// Companion to the murmur3 case. `spark_xxhash64` keeps every struct on the Comet kernel, + /// so the tagged interval struct must reach its interval arm there too, bare, nested, and + /// with the children widened to nullable. + #[test] + fn test_calendar_interval_uses_spark_interval_hash() { + use super::{spark_compatible_xxhash64, spark_xxhash64}; + use crate::hash_funcs::utils::test_utils::calendar_interval_array; + use arrow::array::{Array, ListArray, StructArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Fields}; + use datafusion::common::ScalarValue; + use datafusion::physical_plan::ColumnarValue; + + let intervals = [ + Some((12, 0, 0)), + Some((12, 5, 0)), + Some((14, -3, 7_008_009)), + None, + ]; + let expected: Vec = intervals + .iter() + .map(|v| { + v.map_or(42, |(months, _, micros): (i32, i32, i64)| { + spark_compatible_xxhash64( + months.to_le_bytes(), + spark_compatible_xxhash64(micros.to_le_bytes(), 42), + ) + }) + }) + .collect(); + // Spark 4.1.3 returns 604378839101286624 for `xxhash64(make_interval(y))` with y = 1. + assert_eq!(expected[0] as i64, 604378839101286624); + + for nullable_children in [false, true] { + let array = calendar_interval_array(&intervals, nullable_children); + let hash = |column: ArrayRef| { + let mut hashes = vec![42u64; intervals.len()]; + create_xxhash64_hashes(&[column], &mut hashes).unwrap(); + hashes + }; + assert_eq!(hash(Arc::clone(&array)), expected); + + let result = spark_xxhash64(&[ + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Int64(Some(42))), + ]) + .unwrap(); + let ColumnarValue::Array(result) = result else { + panic!("expected an array result") + }; + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.value(0), 604378839101286624); + + let list: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", array.data_type().clone(), true)), + OffsetBuffer::from_lengths(vec![1; intervals.len()]), + Arc::clone(&array), + None, + )); + assert_eq!(hash(list), expected); + let wrapper_fields: Fields = + vec![Arc::new(Field::new("c", array.data_type().clone(), true))].into(); + let wrapped: ArrayRef = Arc::new(StructArray::new( + wrapper_fields, + vec![Arc::clone(&array)], + None, + )); + assert_eq!(hash(wrapped), expected); + } + } + #[test] fn test_i8() { test_xxhash64_hash::( diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 738ae08835e..536e966d380 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -77,11 +77,11 @@ pub use comet_scalar_funcs::{ }; pub use csv_funcs::*; pub use datetime_funcs::{ - calendar_interval_type, spark_day_name, spark_month_name, spark_seconds_of_time, spark_to_time, - SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkDayOfWeek, SparkHour, - SparkHoursTransform, SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkMinute, - SparkNextDay, SparkSecond, SparkSecondsToTimestamp, SparkUnixTimestamp, SparkWeekDay, - TimestampTruncExpr, + calendar_interval_type, is_calendar_interval_fields, spark_day_name, spark_month_name, + spark_seconds_of_time, spark_to_time, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, + SparkDayOfWeek, SparkHour, SparkHoursTransform, SparkMakeDate, SparkMakeInterval, + SparkMakeTime, SparkMinute, SparkNextDay, SparkSecond, SparkSecondsToTimestamp, + SparkUnixTimestamp, SparkWeekDay, TimestampTruncExpr, }; pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult}; pub use hash_funcs::*; diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql index 2f335f35127..96cb5d312d7 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql @@ -91,3 +91,23 @@ FROM ( AS t(n, c) DISTRIBUTE BY n ) + +statement +CREATE TABLE test_interval_hash(y int) USING parquet + +statement +INSERT INTO test_interval_hash VALUES (1), (-2), (NULL) + +-- Native make_interval output hashed natively over a Parquet scan. The interval is a tagged Arrow +-- struct, and it must hash as Spark's CalendarInterval, not as a struct of its three fields. For +-- y = 1, Spark returns -351543533 and 604378839101286624. +query +SELECT + y, + hash(make_interval(y)), + xxhash64(make_interval(y)), + hash(make_interval(y, y, y, y, y, y)), + xxhash64(make_interval(y, y, y, y, y, y)), + hash(array(make_interval(y, y))), + xxhash64(named_struct('i', make_interval(0, 0, 0, y))) +FROM test_interval_hash From 7cb1f76304b50781eea9cade3f41f4c41cf255c9 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 29 Sep 2026 02:07:38 +0800 Subject: [PATCH 12/20] fix: keep CalendarInterval layout through collect_list normalization coerce_collect_child_nullability widened every nested field of a collect_list / collect_set argument to nullable, including the children of the tagged CalendarInterval struct. The JVM then no longer recognized the aggregate result as ARRAY, so a codegen consumer such as transform failed with "getInterval not implemented", and the final stage of a global aggregate, which reads its state back in the Spark-declared layout, failed validating its own output batch. Keep an interval in its tagged layout during the normalization, and recognize one on the JVM by its marker alone, as the native side does, so a metadata-preserving nullability widening elsewhere cannot turn it into a plain struct. --- native/core/src/execution/planner.rs | 55 +++++++++++++++++-- .../apache/spark/sql/comet/util/Utils.scala | 17 ++++-- .../datetime/calendar_interval.sql | 20 +++++++ .../spark/sql/comet/util/UtilsSuite.scala | 28 +++++++++- 4 files changed, 109 insertions(+), 11 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 10cc4ae37d7..fdac8d969f4 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -89,10 +89,10 @@ use datafusion::{ prelude::SessionContext, }; use datafusion_comet_spark_expr::{ - create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle, - BloomFilterAgg, BloomFilterMightContain, CometCollectList, CometCollectSet, CsvWriteOptions, - EvalMode, SparkArraysZipFunc, SparkBloomFilterVersion, SparkListAgg, SparkPercentile, - SumInteger, ToCsv, + calendar_interval_type, create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, + is_calendar_interval_fields, BinaryOutputStyle, BloomFilterAgg, BloomFilterMightContain, + CometCollectList, CometCollectSet, CsvWriteOptions, EvalMode, SparkArraysZipFunc, + SparkBloomFilterVersion, SparkListAgg, SparkPercentile, SumInteger, ToCsv, }; use iceberg::expr::Bind; @@ -178,6 +178,13 @@ struct JoinParameters { /// Return a copy of `data_type` with every nested field marked nullable. Map key fields are left /// non-nullable to preserve Arrow's map invariant. Primitive types are returned unchanged. +/// +/// A Spark `CalendarIntervalType` value, carried as the tagged struct of +/// [`calendar_interval_type`], is a scalar with a fixed layout rather than a struct, so it is kept +/// in exactly that layout wherever it appears. Widening its children would make a global +/// `collect_list` declare a result type that differs from the `ARRAY` state its final +/// stage reads back in the Spark-declared layout, and the final aggregate would then fail +/// validating its own output batch. fn make_all_fields_nullable(data_type: &DataType) -> DataType { fn nullable_field(field: &Field, nullable: bool) -> FieldRef { Arc::new( @@ -190,6 +197,7 @@ fn make_all_fields_nullable(data_type: &DataType) -> DataType { ) } match data_type { + DataType::Struct(fields) if is_calendar_interval_fields(fields) => calendar_interval_type(), DataType::Struct(fields) => { DataType::Struct(fields.iter().map(|f| nullable_field(f, true)).collect()) } @@ -7409,6 +7417,45 @@ mod tests { ); } + /// A `CalendarIntervalType` value keeps its tagged layout through the coercion, bare or + /// nested, while ordinary fields around it are still widened. A global `collect_list` reads + /// its state back in the Spark-declared `ARRAY` layout in the final stage, so a + /// widened declaration there would fail the aggregate's output-batch validation. + #[test] + fn test_collect_agg_coercion_keeps_calendar_interval_layout() { + let interval = datafusion_comet_spark_expr::calendar_interval_type(); + let raw = Arc::new(Column::new("s", 0)) as Arc; + let coerced_type = |data_type: DataType| { + let plan_schema = collect_agg_schema(data_type); + PhysicalPlanner::coerce_collect_child_nullability(Arc::clone(&raw), &plan_schema) + .unwrap() + .data_type(plan_schema.as_ref()) + .unwrap() + }; + + assert_eq!(coerced_type(interval.clone()), interval); + + let nested = |nullable: bool| { + DataType::Struct(Fields::from(vec![ + Field::new("i", interval.clone(), true), + Field::new("n", DataType::Int32, nullable), + ])) + }; + assert_eq!(coerced_type(nested(false)), nested(true)); + + // An interval whose children were already widened is normalized back to the layout. + let DataType::Struct(fields) = &interval else { + unreachable!() + }; + let widened = DataType::Struct( + fields + .iter() + .map(|f| Arc::new(f.as_ref().clone().with_nullable(true))) + .collect(), + ); + assert_eq!(coerced_type(widened), interval); + } + /// Primitive arguments carry no field-level nullability, so no cast is inserted. #[test] fn test_collect_agg_no_coercion_for_primitive_child() { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index fce01af1293..984af3ddf1f 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -249,17 +249,22 @@ object Utils extends CometTypeShim with Logging { /** * Returns true only for the Spark-tagged `CalendarIntervalType` struct produced by * [[toArrowField]]. Matching the field shape (`months: Int32`, `days: Int32`, `microseconds: - * Int64`, all non-nullable) is not sufficient: a user-defined `struct` has the same shape, so this also requires the - * `SPARK::calendarInterval::struct` metadata marker on the `months` child. Only fields carrying - * the marker round-trip back to `CalendarIntervalType`; unmarked structs stay plain - * `StructType`s. + * Int64`) is not sufficient: a user-defined `struct` + * has the same shape, so this also requires the `SPARK::calendarInterval::struct` metadata + * marker on the `months` child. Only fields carrying the marker round-trip back to + * `CalendarIntervalType`; unmarked structs stay plain `StructType`s. + * + * Child nullability is deliberately not compared. [[toArrowField]] declares the children + * non-nullable, but a nullability widening that keeps field metadata must not turn an interval + * into a plain struct, whose consumers then fail to read it as an interval. The struct's own + * validity decides whether an interval is null, and `ColumnVector.getInterval` reads the three + * children under it. The native counterpart is `is_calendar_interval_fields`. */ def isCalendarIntervalStructField(field: Field): Boolean = { val children = field.getChildren def child(index: Int, name: String, bits: Int): Boolean = { val f = children.get(index) - f.getName == name && f.getType == new ArrowType.Int(bits, true) && !f.isNullable + f.getName == name && f.getType == new ArrowType.Int(bits, true) } field.getType == ArrowType.Struct.INSTANCE && children.size == 3 && diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql index 96cb5d312d7..6c46017ff09 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql @@ -111,3 +111,23 @@ SELECT hash(array(make_interval(y, y))), xxhash64(named_struct('i', make_interval(0, 0, 0, y))) FROM test_interval_hash + +-- Native collect_list over intervals, consumed by the JVM codegen dispatcher (transform). The +-- result must be read back as ARRAY, and the final stage of the global aggregate must +-- accept the interval state it reads back from the shuffle. Both failed while the aggregate's +-- argument normalization widened the interval's children to nullable. +query +SELECT transform(collect_list(make_interval(y, y, y, y, y, y)), x -> x) +FROM test_interval_hash +WHERE y = 1 + +query +SELECT y, transform(collect_list(make_interval(y, 1)), x -> x) +FROM test_interval_hash +GROUP BY y + +-- The aggregate result also hashes natively as ARRAY. +query +SELECT hash(collect_list(make_interval(y, y))), xxhash64(collect_list(make_interval(y, y))) +FROM test_interval_hash +WHERE y = 1 diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index 0f74fa19d8a..1443f058f65 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -19,10 +19,13 @@ package org.apache.spark.sql.comet.util +import scala.jdk.CollectionConverters._ + import org.apache.arrow.c.CDataDictionaryProvider +import org.apache.arrow.vector.types.pojo.{Field, FieldType} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.vectorized.ConstantColumnVector -import org.apache.spark.sql.types.{CalendarIntervalType, IntegerType, StringType, StructField, StructType, TimestampType} +import org.apache.spark.sql.types.{ArrayType, CalendarIntervalType, IntegerType, StringType, StructField, StructType, TimestampType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.comet.CometArrowAllocator @@ -37,6 +40,29 @@ class UtilsSuite extends CometTestBase { assert(error.getMessage.contains("requires toArrowField")) } + test("CalendarIntervalType is recognized by its marker, whatever the child nullability") { + val field = Utils.toArrowField("i", CalendarIntervalType, nullable = true, "UTC") + def mapChildren(field: Field)(f: Field => Field): Field = + new Field(field.getName, field.getFieldType, field.getChildren.asScala.map(f).asJava) + + // Children widened to nullable with the metadata kept are still an interval. + val widened = mapChildren(field) { child => + val fieldType = new FieldType(true, child.getType, child.getDictionary, child.getMetadata) + new Field(child.getName, fieldType, child.getChildren) + } + assert(Utils.fromArrowField(widened) === CalendarIntervalType) + val list = Utils.toArrowField("l", ArrayType(CalendarIntervalType), nullable = true, "UTC") + val widenedList = mapChildren(list)(_ => widened) + assert(Utils.fromArrowField(widenedList) === ArrayType(CalendarIntervalType)) + + // The same shape without the marker is an ordinary struct. + val unmarked = mapChildren(field) { child => + val fieldType = new FieldType(child.isNullable, child.getType, child.getDictionary) + new Field(child.getName, fieldType, child.getChildren) + } + assert(Utils.fromArrowField(unmarked).isInstanceOf[StructType]) + } + test("serializeBatches preserves row count for a zero-column batch") { val numRows = 5 val batch = new ColumnarBatch(Array.empty[ColumnVector], numRows) From 05fedc4ed9829f1fdc8d6091af0b3cbfcc32727d Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 29 Sep 2026 02:07:38 +0800 Subject: [PATCH 13/20] fix: preserve make_interval NULL short-circuiting for eager arguments Spark's MakeInterval evaluates its arguments left to right and returns NULL at the first NULL one, while the native scalar function evaluates every argument up front. With ANSI on, make_interval(y, CAST(s AS INT)) over a row (NULL, 'bad') failed natively where Spark returns NULL. Keep make_interval native only when every argument after the first nullable one is a column, a literal, or a lossless up-cast of one, none of which can throw or carry state. Route other shapes through the JVM codegen dispatcher, as CometSequence and CometArrayJoin already do. --- docs/source/user-guide/latest/expressions.md | 2 +- .../org/apache/comet/serde/datetime.scala | 41 +++++++++++++++++-- .../datetime/make_interval_ansi.sql | 30 ++++++++++++++ 3 files changed, 69 insertions(+), 4 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index ed75b00cdee..d4b70313f0f 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -278,7 +278,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | `localtimestamp` | ✅ | — | | | `make_date` | ✅ | Native | | | `make_dt_interval` | ✅ | Codegen dispatch | | -| `make_interval` | ✅ | Native | | +| `make_interval` | ✅ | Hybrid | Runs natively when each argument after the first nullable one is a column, a literal, or a lossless up-cast of one; other shapes use the JVM codegen dispatcher to keep Spark's NULL short-circuit ([details](compatibility/expressions/datetime.md)) | | `make_time` | ✅ | — | Spark 4.1+; requires `spark.sql.timeType.enabled=true`, which Spark leaves off by default. Runs natively; remaining TIME type work is tracked by [#4288](https://github.com/apache/datafusion-comet/issues/4288) | | `make_timestamp` | ✅ | Hybrid | | | `make_timestamp_ltz` | ✅ | — | 2-arg TIME form falls back | 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 61fde98c8a1..c05596fcd40 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, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, TimestampAdd, TimestampDiff, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year} +import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, BoundReference, Cast, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, TimestampAdd, TimestampDiff, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{CalendarIntervalType, DataType, DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String @@ -884,8 +884,43 @@ object CometMakeYMInterval extends CometCodegenDispatch[MakeYMInterval] object CometMakeDTInterval extends CometCodegenDispatch[MakeDTInterval] -object CometMakeInterval extends CometExpressionSerde[MakeInterval] { - override def getSupportLevel(expr: MakeInterval): SupportLevel = Compatible() +object CometMakeInterval extends CometExpressionSerde[MakeInterval] with CodegenDispatchFallback { + + private val eagerArgumentReason = + "An argument that follows a nullable argument and is not a column reference, a literal, or " + + "a lossless up-cast of one, such as an ANSI `CAST` that can fail: Spark stops at the " + + "first NULL argument without evaluating the rest, while the native kernel evaluates every " + + "argument first" + + override def getUnsupportedReasons(): Seq[String] = Seq(eagerArgumentReason) + + override def getSupportLevel(expr: MakeInterval): SupportLevel = + if (argumentsCanBeEvaluatedEagerly(expr)) Compatible() + else Unsupported(Some(eagerArgumentReason)) + + /** + * Spark's `MakeInterval` is a `SeptenaryExpression`: it evaluates its arguments left to right + * and returns NULL at the first NULL one, so an argument that follows a nullable argument only + * runs on the rows where every earlier argument is non-null. The native scalar function + * evaluates every argument over the whole batch before its kernel checks for nulls. Running + * such an argument on the extra rows is unobservable only when it cannot throw, carry state or + * have a side effect. That holds for a column reference, a literal, and a lossless up-cast of + * either, which covers the implicit casts Spark adds to widen integral and decimal arguments. + * Anything else routes through the JVM codegen dispatcher, which runs Spark's own nested + * evaluation. `foldable` is not a usable test: an unfolded foldable argument can still throw. + * Arguments before the first nullable one run on every row in Spark too, so they are + * unrestricted. + */ + private def argumentsCanBeEvaluatedEagerly(expr: MakeInterval): Boolean = { + def unobservable(arg: Expression): Boolean = arg match { + case _: Literal | _: Attribute | _: BoundReference => true + case cast: Cast => + Cast.canUpCast(cast.child.dataType, cast.dataType) && unobservable(cast.child) + case _ => false + } + val firstNullable = expr.children.indexWhere(_.nullable) + firstNullable < 0 || expr.children.drop(firstNullable + 1).forall(unobservable) + } override def convert( expr: MakeInterval, diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql index 90fdb0ab046..57b84e9f28b 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql @@ -72,3 +72,33 @@ WHERE id = 11 query SELECT make_interval(0, 0, 0, 0, 2562048) + +statement +CREATE TABLE test_make_interval_short_circuit(id int, y int, sm smallint, s string) USING parquet + +statement +INSERT INTO test_make_interval_short_circuit VALUES + (0, NULL, 1, 'bad'), + (1, 1, 2, '2'), + (2, 3, 3, 'bad') + +-- Spark stops at the first NULL argument without evaluating the rest, so the invalid ANSI cast +-- never runs on the row whose years is NULL and the result is NULL. Native execution would +-- evaluate every argument first and fail, so this shape runs through the JVM codegen dispatcher. +query expect_dispatch(make_interval) +SELECT id, make_interval(y, CAST(s AS INT)) +FROM test_make_interval_short_circuit +WHERE id <= 1 + +-- Where years is not NULL, Spark does evaluate the cast, and so does the dispatcher. +query expect_error(CAST_INVALID_INPUT) +SELECT make_interval(y, CAST(s AS INT)) +FROM test_make_interval_short_circuit +WHERE id = 2 + +-- Columns, literals and lossless up-casts of them (here SMALLINT to INT, and the decimal literal +-- to DECIMAL(18, 6)) cannot throw or carry state, so evaluating them eagerly is unobservable and +-- the expression stays native. +query expect_native(make_interval) +SELECT id, make_interval(y, sm, id, 0, 0, 0, 1.5) +FROM test_make_interval_short_circuit From 10e26d3f1b5296081c1f042518f0a5ad7228a92e Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 2 Oct 2026 23:05:24 +0800 Subject: [PATCH 14/20] fix: drop the obsolete calendar interval range cap on folded literals A calendar interval now crosses into Arrow as a months/days/microseconds struct, and the codegen dispatcher writes microseconds into it unchanged, so the Long.MaxValue / 1000 microsecond cap that kept folded literals off the nanosecond IntervalMonthDayNano encoding no longer guards anything, and a folded map('k', make_interval(0, 0, 0, 0, 2562048)) still fell back to Spark. Remove the cap, MaxArrowIntervalMicros and their Scaladoc, and turn the test that pinned the fallback into one that requires native execution for values past the old limit and near the Long range limit. --- .../org/apache/comet/serde/literals.scala | 40 +++++++------------ .../comet/CometArrayExpressionSuite.scala | 33 +++++++-------- 2 files changed, 32 insertions(+), 41 deletions(-) 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 98a3757ffc8..d5ba0440243 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, CreateArray, CreateMap, Expression, KnownNullable, Literal, MapFromArrays} import org.apache.spark.sql.catalyst.util.{ArrayData, MapData, TypeUtils} import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructType, TimestampNTZType, TimestampType} -import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String} +import org.apache.spark.unsafe.types.UTF8String import com.google.protobuf.ByteString @@ -342,39 +342,32 @@ object CometLiteral extends CometExpressionSerde[Literal] with CometTypeShim wit /** * True when some scalar inside `value` has no faithful Arrow encoding, in which case the - * literal has to stay on Spark. Two of the copies a Literal takes towards the native side are - * not total: - * - a `UTF8String` holding malformed UTF-8 (`CAST(X'FF' AS STRING)`). A rebuilt `CreateMap` - * reaches the JVM codegen dispatcher, whose generated writer copies the bytes verbatim into - * a `VarCharVector`; native string kernels then read them through unchecked Arrow FFI and - * abort the JVM on a `hint::unreachable_unchecked`, or silently mis-count characters. The - * direct paths (`convert`'s `StringType` arm and [[makeListLiteral]]'s) instead go through - * `UTF8String.toString`, which substitutes U+FFFD. Neither is fixable on the wire, because - * `string_val` and `string_values` are proto `string` fields. - * - a `CalendarInterval` beyond about 292 years of elapsed time overflows - * `Math.multiplyExact(microseconds, 1000L)` on the way into `IntervalMonthDayNanoVector`, - * whose nanosecond field cannot represent it (#5279). Spark itself accepts the value. + * literal has to stay on Spark. One of the copies a Literal takes towards the native side is + * not total: a `UTF8String` holding malformed UTF-8 (`CAST(X'FF' AS STRING)`). A rebuilt + * `CreateMap` reaches the JVM codegen dispatcher, whose generated writer copies the bytes + * verbatim into a `VarCharVector`; native string kernels then read them through unchecked Arrow + * FFI and abort the JVM on a `hint::unreachable_unchecked`, or silently mis-count characters. + * The direct paths (`convert`'s `StringType` arm and [[makeListLiteral]]'s) instead go through + * `UTF8String.toString`, which substitutes U+FFFD. Neither is fixable on the wire, because + * `string_val` and `string_values` are proto `string` fields. * - * The same hazards exist for a non-folded `map(...)` / a scanned string column, which never - * reach this serde; this only keeps a Literal from newly admitting them. + * The same hazard exists for a non-folded `map(...)` / a scanned string column, which never + * reach this serde; this only keeps a Literal from newly admitting it. */ private def hasUnwritableValue(value: Any, dataType: DataType): Boolean = mayHoldUnwritableValue(dataType) && unwritableValueIn(value, dataType) /** - * Type-level precondition for [[hasUnwritableValue]]: only a string or a calendar interval - * reaches one of its hazardous arms, so a large all-numeric literal is never walked element by - * element. Applied once per collection rather than once per element. + * Type-level precondition for [[hasUnwritableValue]]: only a string reaches its hazardous arm, + * so a large literal without strings is never walked element by element. Applied once per + * collection rather than once per element. */ private def mayHoldUnwritableValue(dataType: DataType): Boolean = - SupportLevel.containsType(dataType, classOf[StringType], classOf[CalendarIntervalType]) + SupportLevel.containsType(dataType, classOf[StringType]) private def unwritableValueIn(value: Any, dataType: DataType): Boolean = dataType match { case _ if value == null => false case _: StringType => !isValidUtf8(value.asInstanceOf[UTF8String]) - case CalendarIntervalType => - val micros = value.asInstanceOf[CalendarInterval].microseconds - micros > MaxArrowIntervalMicros || micros < -MaxArrowIntervalMicros case ArrayType(et, _) => unwritableElementIn(value.asInstanceOf[ArrayData], et) case MapType(kt, vt, _) => val mapData = value.asInstanceOf[MapData] @@ -393,9 +386,6 @@ object CometLiteral extends CometExpressionSerde[Literal] with CometTypeShim wit (0 until arr.numElements()).exists(i => unwritableValueIn(arr.get(i, elementType), elementType)) - /** Largest `CalendarInterval.microseconds` that survives conversion to Arrow nanoseconds. */ - private final val MaxArrowIntervalMicros: Long = Long.MaxValue / 1000 - /** * True when every map key type reachable inside `dataType` can be rebuilt and consumed * natively. A folded map is rebuilt as a `CreateMap`; once native it may reach a map consumer diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 1c4ad3ce152..78310e2a6df 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1357,24 +1357,25 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } - // https://github.com/apache/datafusion-comet/issues/5544: Arrow's IntervalMonthDayNano holds - // elapsed time in nanoseconds, so the dispatcher's `Math.multiplyExact(microseconds, 1000L)` - // overflows past about 292 years, where Spark accepts the value (#5279). Expansion has to - // decline the value, not just the type: the calendar-interval restriction in - // `mapKeyTypesExpandable` only covers map keys. The `array(map(...))` spelling also covers the - // outer `ArrayType` recursion of the value walk. - test("folded map value with an out-of-range calendar interval falls back (multirow)") { + // https://github.com/apache/datafusion-comet/issues/5279: a calendar interval crosses into Arrow + // as a months/days/microseconds struct, and the dispatcher writes `microseconds` into it as-is, + // so a folded map value keeps Spark's whole microsecond range. Expansion used to decline any + // value past about 292 years, which the nanosecond `IntervalMonthDayNano` encoding could not + // hold. 2562048 hours is the first whole hour past that limit, and 2147483647 hours and minutes + // plus 999999999999.999999 seconds is about 96% of `Long.MaxValue` microseconds. The + // `array(map(...))` spelling also covers the outer `ArrayType` recursion of expansion. + test("folded map value with a calendar interval past the nanosecond range (multirow)") { withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { Seq( - "array(map(1, make_interval(0, 0, 0, 0, 3000000, 0, 0)))" -> "ArrayType", - "map(1, make_interval(0, 0, 0, 0, 3000000, 0, 0))" -> "MapType", - "map(1, make_interval(0, 0, 0, 0, -3000000, 0, 0))" -> "MapType").foreach { - case (value, declaredType) => - checkSparkAnswerAndFallbackReason( - s"SELECT _1 AS id, $value AS v FROM tbl", - s"Unsupported data type $declaredType") - } - // An interval inside the nanosecond range still runs natively. + "map('k', make_interval(0, 0, 0, 0, 2562048))", + "map(1, make_interval(0, 0, 0, 0, -3000000, 0, 0))", + "array(map(1, make_interval(0, 0, 0, 0, 3000000, 0, 0)))", + "map(1, make_interval(0, 0, 0, 0, 2147483647, 2147483647, 999999999999.999999))", + "map(1, make_interval(0, 0, 0, 0, -2147483647, -2147483647, -999999999999.999999))") + .foreach { value => + checkSparkAnswerAndOperator(s"SELECT _1 AS id, $value AS v FROM tbl") + } + // An interval inside the old nanosecond range runs natively as before. checkSparkAnswerAndOperator( "SELECT _1 AS id, array(map(1, make_interval(0, 0, 0, 0, 24, 0, 0))) AS a FROM tbl") } From c017d8f20e57f81cfd00e011f1e714eb37bd9c92 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 2 Oct 2026 23:05:24 +0800 Subject: [PATCH 15/20] test: cover calendar interval negation, addition and subtraction CometUnaryMinus, CometAdd and CometSubtract only admit numeric types and do not mix in CodegenDispatchFallback, so interval arithmetic falls back to Spark and never hands the interval struct to a numeric kernel. Pin that, together with Spark's answers, for make_interval columns directly, over intervals built natively below a native shuffle, and under ANSI, where Spark uses the *Exact variants. --- .../expressions/datetime/make_interval.sql | 19 +++++++++++++++++++ .../datetime/make_interval_ansi.sql | 9 +++++++++ 2 files changed, 28 insertions(+) diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql index bf0afc93ecb..c6878c5a2f0 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql @@ -90,6 +90,25 @@ FROM ( DISTRIBUTE BY years ) +-- Spark negates, adds and subtracts calendar intervals with IntervalUtils. CometUnaryMinus, +-- CometAdd and CometSubtract only admit numeric types, so these projections stay on Spark rather +-- than handing the interval struct to a numeric kernel. +query expect_fallback(Unsupported datatype CalendarIntervalType) +SELECT -make_interval(years, months, weeks, days, hours, mins, secs), + make_interval(years) + make_interval(0, months), + make_interval(years, months, weeks, days) - make_interval(0, 0, 0, 0, hours, mins, secs) +FROM test_make_interval + +-- The same arithmetic over intervals that native make_interval built below a native shuffle, so +-- Spark reads them back from the interval struct. +query expect_fallback(Unsupported datatype CalendarIntervalType) +SELECT -i, i + make_interval(years), i - make_interval(0, months) +FROM ( + SELECT years, months, make_interval(years, months, weeks, days, hours, mins, secs) AS i + FROM test_make_interval + DISTRIBUTE BY years +) + query SELECT make_interval(0, 1, 0, 1, 0, 0, 100.000001) diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql index 57b84e9f28b..ddd1e26b973 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_ansi.sql @@ -54,6 +54,15 @@ FROM test_make_interval_ansi WHERE id BETWEEN 0 AND 9 ORDER BY id +-- Under ANSI, Spark uses IntervalUtils.negateExact, addExact and subtractExact. Comet still +-- declines interval arithmetic, so these stay on Spark too. +query expect_fallback(Unsupported datatype CalendarIntervalType) +SELECT -make_interval(years, months, weeks, days, hours, mins, secs), + make_interval(years, months) + make_interval(0, 0, weeks, days), + make_interval(years, months, weeks, days) - make_interval(0, 0, 0, 0, hours, mins, secs) +FROM test_make_interval_ansi +WHERE id BETWEEN 0 AND 9 + -- Spark <= 4.1 preserves the JDK literal "integer overflow": -- https://github.com/apache/spark/blob/v4.1.3/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala#L704-L713 -- SPARK-55714 canonicalizes simple " overflow" messages to "overflow" in Spark 4.2: From 603e7a4262425cd91b0ffcc95dffb0600bfbec6e Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 4 Oct 2026 23:15:43 +0800 Subject: [PATCH 16/20] fix: pass Decimal(18, 6) seconds in the make_interval bench The native kernel accepts only Decimal128(18, 6) seconds, the type Spark gives the secs argument. The bench still passed a Float64 column, so cargo bench panicked in unwrap on the first benchmark. --- native/spark-expr/benches/make_interval.rs | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/native/spark-expr/benches/make_interval.rs b/native/spark-expr/benches/make_interval.rs index 3590db5fffd..ee94b85fd5d 100644 --- a/native/spark-expr/benches/make_interval.rs +++ b/native/spark-expr/benches/make_interval.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use arrow::array::{ArrayRef, Decimal128Array}; use arrow::datatypes::{DataType, Field}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::common::config::ConfigOptions; @@ -25,14 +26,25 @@ use std::sync::Arc; #[path = "common/mod.rs"] mod common; -use common::{f64_array, i32_array, NULL_RATIOS, ROW_COUNTS}; +use common::{i32_array, is_null, NULL_RATIOS, ROW_COUNTS}; + +/// The kernel accepts only `Decimal(18, 6)` seconds, the type Spark gives the `secs` argument. +fn seconds_array(rows: usize, null_ratio: f64) -> ArrayRef { + let arr = (0..rows) + .map(|i| (!is_null(i, null_ratio)).then_some((i % 60_000_000) as i128)) + .collect::() + .with_precision_and_scale(18, 6) + .unwrap(); + Arc::new(arr) +} fn criterion_benchmark(c: &mut Criterion) { let udf = SparkMakeInterval::new(false); let mut group = c.benchmark_group("make_interval"); for rows in ROW_COUNTS { for (null_ratio, tag) in NULL_RATIOS { - // make_interval(years, months, weeks, days, hours, mins: Int32, secs: Float64) + // make_interval(years, months, weeks, days, hours, mins: Int32, + // secs: Decimal128(18, 6)) let args = vec![ ColumnarValue::Array(i32_array(rows, null_ratio, |i| (i % 10) as i32)), ColumnarValue::Array(i32_array(rows, null_ratio, |i| (i % 12) as i32)), @@ -40,7 +52,7 @@ fn criterion_benchmark(c: &mut Criterion) { ColumnarValue::Array(i32_array(rows, null_ratio, |i| (i % 28) as i32)), ColumnarValue::Array(i32_array(rows, null_ratio, |i| (i % 24) as i32)), ColumnarValue::Array(i32_array(rows, null_ratio, |i| (i % 60) as i32)), - ColumnarValue::Array(f64_array(rows, null_ratio, |i| (i % 60) as f64)), + ColumnarValue::Array(seconds_array(rows, null_ratio)), ]; group.bench_with_input( BenchmarkId::from_parameter(format!("{rows}/{tag}")), From 6cd802abf22d0dd0ed9c5746c66cfaa3ac71e17a Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 4 Oct 2026 23:33:03 +0800 Subject: [PATCH 17/20] fix: dispatch make_interval when an argument has no native path On branch-1.1, make_interval always ran in the JVM codegen dispatcher. With the native kernel, an argument without a native path made convert return None, so the whole projection fell back to Spark. An example is make_interval(0, 0, 0, hash(d), y) with d DECIMAL(38, 10). convert now runs the whole expression in the dispatcher when the native conversion fails. The failed conversion leaves explain tags on the argument trees: the fallback reason of the argument without a native path, and the native names of the arguments that converted first. The dispatcher runs these arguments now, so CometExplainInfo.saveTags restores the tags from before the attempt. When the dispatcher also declines, the tags stay and explain shows why. --- docs/source/user-guide/latest/expressions.md | 2 +- .../apache/comet/ExtendedExplainInfo.scala | 25 ++++++++++++++++ .../org/apache/comet/serde/datetime.scala | 14 ++++++--- .../expressions/datetime/make_interval.sql | 20 +++++++++++++ .../math/routing_intervals_disabled.sql | 8 +++++ .../math/routing_intervals_enabled.sql | 4 +++ .../org/apache/comet/CometCodegenSuite.scala | 29 ++++++++++++++++++- 7 files changed, 96 insertions(+), 6 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index d3c89c3d554..96ed4dc1b0e 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -282,7 +282,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | `localtimestamp` | ✅ | — | | | `make_date` | ✅ | Native | | | `make_dt_interval` | ✅ | Codegen dispatch | | -| `make_interval` | ✅ | Hybrid | Runs natively when each argument after the first nullable one is a column, a literal, or a lossless up-cast of one; other shapes use the JVM codegen dispatcher to keep Spark's NULL short-circuit ([details](compatibility/expressions/datetime.md)) | +| `make_interval` | ✅ | Hybrid | Runs natively when each argument after the first nullable one is a column, a literal, or a lossless up-cast of one. Other shapes use the JVM codegen dispatcher to keep Spark's NULL short-circuit. An argument without a native path also uses the dispatcher ([details](compatibility/expressions/datetime.md)) | | `make_time` | ✅ | — | Spark 4.1+; requires `spark.sql.timeType.enabled=true`, which Spark leaves off by default. Runs natively; remaining TIME type work is tracked by [#4288](https://github.com/apache/datafusion-comet/issues/4288) | | `make_timestamp` | ✅ | Hybrid | | | `make_timestamp_ltz` | ✅ | — | 2-arg TIME form falls back | diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala index 07a7f57ccd8..01279146baa 100644 --- a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala +++ b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala @@ -365,6 +365,31 @@ object CometExplainInfo { collectTagValues(exprs.filterNot(isNeverTagged), tag) } + /** + * Saves the explain tags on each node of `exprs`. The returned function restores them when a + * serde replaces a failed native conversion of these trees. + */ + def saveTags(exprs: Seq[Expression]): () => Unit = { + val saved = exprs.flatMap(_.collect { case e => e }).flatMap { node => + Seq( + new SavedTag(node, FALLBACK_REASONS), + new SavedTag(node, EXTENSION_INFO), + new SavedTag(node, CODEGEN_DISPATCH_EXPRS), + new SavedTag(node, DISPATCHED_SELF), + new SavedTag(node, NATIVE_EXPRS)) + } + () => saved.foreach(_.restore()) + } + + private final class SavedTag[T](node: Expression, tag: TreeNodeTag[T]) { + private val value = node.getTagValue(tag) + + def restore(): Unit = value match { + case Some(v) => node.setTagValue(tag, v) + case None => node.unsetTagValue(tag) + } + } + /** * Nodes that never carry their own coverage or info tags, so those tags can only arrive by the * copying described in [[collectExprTagValues]]. This set must match 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 d228846823a..4a4242fc68a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{CalendarIntervalType, DataType, DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String -import org.apache.comet.CometConf +import org.apache.comet.{CometConf, CometExplainInfo} import org.apache.comet.expressions.{CometCast, CometEvalMode} import org.apache.comet.serde.CometGetDateField.CometGetDateField import org.apache.comet.serde.ExprOuterClass.Expr @@ -942,13 +942,19 @@ object CometMakeInterval extends CometExpressionSerde[MakeInterval] with Codegen expr: MakeInterval, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { + val restoreArgumentTags = CometExplainInfo.saveTags(expr.children) val childExprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) - val optExpr = scalarFunctionExprToProtoWithReturnType( + scalarFunctionExprToProtoWithReturnType( "make_interval", CalendarIntervalType, expr.failOnError, - childExprs: _*) - optExpr + childExprs: _*).orElse { + // An argument has no native path, so the dispatcher runs the whole expression instead. + val dispatched = CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding) + // The arguments now run in the dispatcher, so the tags of the failed conversion are stale. + if (dispatched.isDefined) restoreArgumentTags() + dispatched + } } } diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql index c6878c5a2f0..4943d27c023 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql @@ -129,3 +129,23 @@ SELECT make_interval(0, 0, 0, 0, 2562048) query SELECT make_interval(0, 0, 0, 0, 0, 0, 1234567890123456789) + +statement +CREATE TABLE test_make_interval_hash(d decimal(38, 10), y int) USING parquet + +statement +INSERT INTO test_make_interval_hash VALUES + (1.5, 1), + (-123456789.0123456789, NULL), + (NULL, 2), + (9999999999999999999999999999.9999999999, 3) + +-- hash has no native path above decimal precision 18. make_interval then runs as a whole through +-- the JVM codegen dispatcher, so the projection stays in Comet. +query expect_dispatch(make_interval) +SELECT make_interval(0, 0, 0, hash(d), y) FROM test_make_interval_hash + +-- abs converts natively before hash fails. It then runs in the dispatcher kernel, so the explain +-- output reports it as dispatched only. +query expect_dispatch(make_interval, hash, abs) +SELECT make_interval(0, 0, 0, hash(d), abs(y)) FROM test_make_interval_hash diff --git a/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_disabled.sql b/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_disabled.sql index 98e745b5e81..1b9829c234f 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_disabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_disabled.sql @@ -32,3 +32,11 @@ SELECT abs(make_dt_interval(i)), abs(make_ym_interval(i)) FROM routing_intervals query expect_native(make_interval) SELECT make_interval(i, 0, 0, 0, 0, 0, 0) FROM routing_intervals + +-- hash has no native path above decimal precision 18, and the dispatcher is off. The fallback +-- reasons name both causes. +query expect_fallback(make_interval: spark.comet.exec.scalaUDF.codegen.enabled=false) +SELECT make_interval(0, 0, 0, hash(CAST(i AS DECIMAL(38, 10))), i) FROM routing_intervals + +query expect_fallback(`DecimalType` with precision > 18 is not supported) +SELECT make_interval(0, 0, 0, hash(CAST(i AS DECIMAL(38, 10))), i) FROM routing_intervals diff --git a/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_enabled.sql b/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_enabled.sql index d509cbc80b5..965b8527060 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_enabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_enabled.sql @@ -32,3 +32,7 @@ SELECT abs(make_dt_interval(i)), abs(make_ym_interval(i)) FROM routing_intervals query expect_native(make_interval) SELECT make_interval(i, 0, 0, 0, 0, 0, 0) FROM routing_intervals + +-- hash has no native path above decimal precision 18, so the dispatcher runs make_interval. +query expect_dispatch(make_interval) +SELECT make_interval(0, 0, 0, hash(CAST(i AS DECIMAL(38, 10))), i) FROM routing_intervals diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 12ab72e93ed..46479a178ef 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.{SparkConf, SparkEnv, TaskContext} import org.apache.spark.sql.{CometTestBase, Row} import org.apache.spark.sql.api.java.UDF1 import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, Literal, MapConcat, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, Literal, MapConcat, Murmur3Hash, ScalaUDF} import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.catalyst.util.GenericArrayData @@ -275,6 +275,33 @@ class CometCodegenSuite } } + test("make_interval drops the fallback reason of an argument that the dispatcher runs") { + // `hash` has no native path above decimal precision 18, so the dispatcher runs all of + // make_interval. The reason that the failed conversion put on `hash` no longer applies. + withTable("t") { + sql("CREATE TABLE t (d DECIMAL(38, 10), y INT, s STRING) USING parquet") + sql("INSERT INTO t VALUES (1.5, 1, 'a'), (NULL, NULL, NULL)") + val (_, plan) = checkSparkAnswerAndImpl( + sql("SELECT make_interval(0, 0, 0, hash(d), y) FROM t"), + dispatched = Seq("make_interval", "hash")) + val hashes = collect(plan) { case p: CometProjectExec => p } + .flatMap(_.projectList) + .flatMap(_.collect { case h: Murmur3Hash => h }) + assert(hashes.nonEmpty, s"expected hash in a native projection:\n$plan") + assert(hashes.forall(_.getTagValue(CometExplainInfo.FALLBACK_REASONS).isEmpty)) + + // A projection that falls back for another reason reports that reason alone. + val lengthKey = CometConf.getExprEnabledConfigKey("Length") + withSQLConf(lengthKey -> "false") { + val (_, fallbackPlan) = + checkSparkAnswer(sql("SELECT make_interval(0, 0, 0, hash(d), y), length(s) FROM t")) + val reasons = new ExtendedExplainInfo().getFallbackReasons(fallbackPlan) + assert(reasons.exists(_.contains(lengthKey)), s"reasons: $reasons") + assert(!reasons.exists(_.contains("precision > 18")), s"reasons: $reasons") + } + } + } + test("ScalaUDF over concat(c1, c2) suppresses the null short-circuit") { // Concat is not NullIntolerant. The dispatcher's short-circuit guard inspects every node in // the bound tree and must skip the whole-tree null short-circuit because one child is From 7c043d1d667ea18e89183ca2a17ec75d197971ac Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 4 Oct 2026 23:33:14 +0800 Subject: [PATCH 18/20] fix: keep Python operators with calendar intervals on Spark With spark.comet.exec.pyarrowUDF.enabled=true, the Comet Python runner sends each vector with its own Arrow field. Comet now carries a calendar interval as a tagged struct, but Spark sends Python an Arrow MONTH_DAY_NANO interval, as branch-1.1 did. The mapInArrow and mapInPandas guard now declines a plan whose input or output holds CalendarIntervalType at any depth, as it does for Variant. --- .../rules/EliminateRedundantTransitions.scala | 15 ++++-- .../sql/comet/CometMapInBatchSuite.scala | 54 +++++++++++++------ 2 files changed, 50 insertions(+), 19 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index dcaa22c9575..7aed894045d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -31,10 +31,11 @@ import org.apache.spark.sql.comet.util.Utils.containsVariantType import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec +import org.apache.spark.sql.types.CalendarIntervalType import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.{withFallbackReason, withInfo} -import org.apache.comet.serde.NativeOptIn +import org.apache.comet.serde.{NativeOptIn, SupportLevel} import org.apache.comet.shims.ShimSQLConf // This rule is responsible for eliminating redundant transitions between row-based and @@ -303,13 +304,21 @@ case class EliminateRedundantTransitions(session: SparkSession) matchMapInArrow(plan) .orElse(matchMapInPandas(plan)) .flatMap { info => + val types = (info.output ++ info.child.output).map(_.dataType) // TODO: Remove this guard once Comet Python operators preserve Variant identity // and Spark's Arrow layout for both input and output. // https://github.com/apache/datafusion-comet/issues/5437 - if ((info.output ++ info.child.output).exists(attr => - containsVariantType(attr.dataType))) { + if (types.exists(containsVariantType)) { withFallbackReason(plan, "Comet Python operators do not support type VariantType") None + } else if (types.exists( + SupportLevel.containsType(_, classOf[CalendarIntervalType]))) { + // Comet carries an interval as a tagged struct, but Spark gives Python an Arrow + // MONTH_DAY_NANO interval. + withFallbackReason( + plan, + "Comet Python operators do not support type CalendarIntervalType") + None } else { extractColumnarChild(info.child).map(child => (info, child)) } diff --git a/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala b/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala index 2a3a0f9920e..cf9c31b016f 100644 --- a/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala +++ b/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, ExprId, PythonUDF} import org.apache.spark.sql.execution.{ColumnarToRowExec, LeafExecNode} import org.apache.spark.sql.execution.python.MapInArrowExec -import org.apache.spark.sql.types.{LongType, StructField, StructType, VariantType} +import org.apache.spark.sql.types.{ArrayType, CalendarIntervalType, LongType, StructField, StructType, VariantType} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.{CometConf, ExtendedExplainInfo} @@ -105,27 +105,49 @@ class CometMapInBatchSuite extends CometTestBase { } } + private def assertStaysOnSpark( + input: Seq[Attribute], + output: Seq[Attribute], + reason: String): Unit = { + val udf = stubPythonUDF.copy( + children = input, + dataType = StructType(output.map(attr => StructField(attr.name, attr.dataType)))) + val plan = MapInArrowExec( + udf, + output, + ColumnarToRowExec(StubCometLeaf(input)), + isBarrier = false, + profile = None) + val rewritten = EliminateRedundantTransitions(spark).apply(plan) + assert(rewritten.isInstanceOf[MapInArrowExec]) + assert(!rewritten.exists(_.isInstanceOf[CometMapInBatchExec])) + assert(new ExtendedExplainInfo().getFallbackReasons(rewritten).exists(_.contains(reason))) + } + test("Variant inputs and outputs keep Python operators on Spark") { val plain = Seq(AttributeReference("id", LongType)()) val variant = Seq(AttributeReference("v", VariantType)()) withSQLConf(CometConf.COMET_PYARROW_UDF_ENABLED.key -> "true") { for ((input, output) <- Seq(variant -> plain, plain -> variant)) { - val udf = stubPythonUDF.copy( - children = input, - dataType = StructType(output.map(attr => StructField(attr.name, attr.dataType)))) - val plan = MapInArrowExec( - udf, + assertStaysOnSpark( + input, + output, + "Comet Python operators do not support type VariantType") + } + } + } + + test("CalendarInterval inputs and outputs keep Python operators on Spark") { + // Spark gives Python an Arrow MONTH_DAY_NANO interval, not the struct that Comet carries. + val plain = Seq(AttributeReference("id", LongType)()) + val interval = Seq(AttributeReference("i", CalendarIntervalType)()) + val nested = Seq(AttributeReference("a", ArrayType(CalendarIntervalType))()) + withSQLConf(CometConf.COMET_PYARROW_UDF_ENABLED.key -> "true") { + for ((input, output) <- Seq(interval -> plain, plain -> interval, nested -> plain)) { + assertStaysOnSpark( + input, output, - ColumnarToRowExec(StubCometLeaf(input)), - isBarrier = false, - profile = None) - val rewritten = EliminateRedundantTransitions(spark).apply(plan) - assert(rewritten.isInstanceOf[MapInArrowExec]) - assert(!rewritten.exists(_.isInstanceOf[CometMapInBatchExec])) - assert( - new ExtendedExplainInfo() - .getFallbackReasons(rewritten) - .exists(_.contains("Comet Python operators do not support type VariantType"))) + "Comet Python operators do not support type CalendarIntervalType") } } } From b0abd19eb27955cd90a2c4ed166c1fd4b21c2c4f Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 4 Oct 2026 23:33:14 +0800 Subject: [PATCH 19/20] perf: write calendar interval child vectors directly in the Arrow writer CalendarIntervalStructWriter boxed months, days and microseconds into a GenericInternalRow for each row, and children.indices.foreach built a Range and a closure each time. It now writes the three child vectors directly, as CometBatchKernelCodegenOutput does. A null interval still sets the struct and each child null at that index, as StructWriter does. The existing writer test now checks that state, the value counts, and a second batch after reset. It passes on the old and the new writer. --- .../comet/execution/arrow/ArrowWriters.scala | 36 +++++++++++-------- .../arrow/CometArrowStreamSuite.scala | 24 ++++++++++--- 2 files changed, 41 insertions(+), 19 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala index 99d70ad331d..8921ca26ba9 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala @@ -27,7 +27,6 @@ import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector._ import org.apache.arrow.vector.complex._ import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.GenericInternalRow import org.apache.spark.sql.catalyst.expressions.SpecializedGetters import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.errors.QueryExecutionErrors @@ -96,10 +95,7 @@ private[arrow] object ArrowWriter { case (CalendarIntervalType, vector: IntervalMonthDayNanoVector) => new IntervalMonthDayNanoWriter(vector) case (CalendarIntervalType, vector: StructVector) => - val children = (0 until vector.size()).map { ordinal => - createFieldWriter(vector.getChildByOrdinal(ordinal)) - } - new CalendarIntervalStructWriter(vector, children.toArray) + new CalendarIntervalStructWriter(vector) case (dt, _) => throw QueryExecutionErrors.notSupportTypeError(dt) } @@ -649,20 +645,30 @@ private[arrow] class StructWriter( } } -private[arrow] class CalendarIntervalStructWriter( - valueVector: StructVector, - children: Array[ArrowFieldWriter]) - extends StructWriter(valueVector, children) { +/** + * Writes the three child vectors directly, as the codegen output does, so it boxes no value. A + * null sets the struct and each child null, as [[StructWriter]] does. + */ +private[arrow] class CalendarIntervalStructWriter(val valueVector: StructVector) + extends ArrowFieldWriter { - private val row = new GenericInternalRow(3) + private val months = valueVector.getChildByOrdinal(0).asInstanceOf[IntVector] + private val days = valueVector.getChildByOrdinal(1).asInstanceOf[IntVector] + private val microseconds = valueVector.getChildByOrdinal(2).asInstanceOf[BigIntVector] + + override def setNull(): Unit = { + months.setNull(count) + days.setNull(count) + microseconds.setNull(count) + valueVector.setNull(count) + } override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - valueVector.setIndexDefined(count) val interval = input.getInterval(ordinal) - row.update(0, interval.months) - row.update(1, interval.days) - row.update(2, interval.microseconds) - children.indices.foreach(i => children(i).write(row, i)) + valueVector.setIndexDefined(count) + months.setSafe(count, interval.months) + days.setSafe(count, interval.days) + microseconds.setSafe(count, interval.microseconds) } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index 68278b7a031..f3ba2b4dfe2 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -79,14 +79,30 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { writer.finish() val arrow = root.getVector(0).asInstanceOf[StructVector] - arrow.getChild("months").asInstanceOf[IntVector].get(0) shouldBe expected.months - arrow.getChild("days").asInstanceOf[IntVector].get(0) shouldBe expected.days - arrow.getChild("microseconds").asInstanceOf[BigIntVector].get(0) shouldBe - expected.microseconds + val months = arrow.getChild("months").asInstanceOf[IntVector] + val days = arrow.getChild("days").asInstanceOf[IntVector] + val micros = arrow.getChild("microseconds").asInstanceOf[BigIntVector] + months.get(0) shouldBe expected.months + days.get(0) shouldBe expected.days + micros.get(0) shouldBe expected.microseconds + // A null interval is null in the struct and in each child. + Seq(arrow, months, days, micros).foreach { vector => + vector.getValueCount shouldBe 2 + vector.isNull(1) shouldBe true + } val comet = CometVector.getVector(arrow, null) comet.getInterval(0) shouldBe expected comet.getInterval(1) shouldBe null + + // After a reset, the writer fills the same vectors again, with the null first this time. + writer.reset() + writer.write(new GenericInternalRow(Array[Any](null))) + writer.write(new GenericInternalRow(Array[Any](expected))) + writer.finish() + val reused = CometVector.getVector(arrow, null) + reused.getInterval(0) shouldBe null + reused.getInterval(1) shouldBe expected } finally { root.close() allocator.close() From 4f4b8f1f83b1647a8bcbfba3d553f3ac748d06a5 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 4 Oct 2026 23:33:14 +0800 Subject: [PATCH 20/20] test: cover more consumers of native make_interval output Add make_interval fixture queries for the paths whose layout changed: - a NULL interval literal in CASE and IF branches and inside coalesce, - a dispatched CAST(i AS STRING) above a DISTRIBUTE BY shuffle, - first and last over make_interval(y) with GROUP BY. Each query matches Spark. Spark removes ELSE NULL and NULL coalesce arguments, so the NULL literal reaches Comet as a THEN or IF branch. --- .../expressions/datetime/make_interval.sql | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql index 4943d27c023..9d32e2d2862 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql @@ -149,3 +149,56 @@ SELECT make_interval(0, 0, 0, hash(d), y) FROM test_make_interval_hash -- output reports it as dispatched only. query expect_dispatch(make_interval, hash, abs) SELECT make_interval(0, 0, 0, hash(d), abs(y)) FROM test_make_interval_hash + +-- A dispatched CAST reads the interval struct that native make_interval wrote below a shuffle. +query expect_dispatch(cast) +SELECT years, CAST(i AS STRING) +FROM ( + SELECT years, make_interval(years, months, weeks, days, hours, mins, secs) AS i + FROM test_make_interval + DISTRIBUTE BY years +) + +statement +CREATE TABLE test_make_interval_merge(id int, y int) USING parquet + +statement +INSERT INTO test_make_interval_merge VALUES (0, 1), (1, NULL), (2, -2), (3, 3), (4, NULL) + +-- Native make_interval merged with a NULL interval literal. Spark drops ELSE NULL and NULL +-- coalesce arguments, so the literal reaches Comet as a THEN or IF branch. +query +SELECT id, + CASE WHEN id % 2 = 0 THEN make_interval(y) ELSE NULL END, + CASE WHEN id % 2 = 0 THEN NULL ELSE make_interval(y, 1) END, + if(id > 1, make_interval(0, y), NULL), + if(id > 1, NULL, make_interval(0, 0, y)) +FROM test_make_interval_merge + +query +SELECT id, + coalesce(make_interval(y), CAST(NULL AS INTERVAL)), + coalesce(make_interval(y), make_interval(0, id)), + coalesce(CASE WHEN id % 2 = 0 THEN NULL ELSE make_interval(y) END, make_interval(0, 0, 0, id)) +FROM test_make_interval_merge + +statement +CREATE TABLE test_make_interval_groups(g int, y int) USING parquet + +statement +INSERT INTO test_make_interval_groups VALUES + (1, 1), (1, 1), (2, NULL), (2, NULL), (3, NULL), (3, 3), (NULL, 4) + +-- first and last keep the interval layout through the partial and final aggregates. Each group +-- here holds one value, so the answer does not depend on row order. +query +SELECT g, first(make_interval(y)), last(make_interval(y)) +FROM test_make_interval_groups +WHERE g <> 3 OR g IS NULL +GROUP BY g + +-- Group 3 mixes NULL and 3 years, so only the ignoreNulls form has one answer there. +query +SELECT g, first(make_interval(y), true), last(make_interval(y), true) +FROM test_make_interval_groups +GROUP BY g