diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index b5abfb88935..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 | 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` | ✅ | 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/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index de458eca329..fd79389f3c2 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -85,10 +85,11 @@ use datafusion_comet_operators::{ WindowFnKind, }; use datafusion_comet_spark_expr::{ - create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle, - BloomFilterAgg, BloomFilterMightContain, CheckedBinaryExpr, CometCollectList, CometCollectSet, - CsvWriteOptions, EvalMode, ListPositionsExpr, SparkArraysZipFunc, SparkBloomFilterVersion, - SparkListAgg, SparkPercentile, Subquery, SumInteger, ToCsv, + calendar_interval_type, create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, + is_calendar_interval_fields, BinaryOutputStyle, BloomFilterAgg, BloomFilterMightContain, + CheckedBinaryExpr, CometCollectList, CometCollectSet, CsvWriteOptions, EvalMode, + ListPositionsExpr, SparkArraysZipFunc, SparkBloomFilterVersion, SparkListAgg, SparkPercentile, + Subquery, SumInteger, ToCsv, }; use datafusion_datasource::TableSchema; use iceberg::expr::Bind; @@ -170,6 +171,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( @@ -182,6 +190,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()) } @@ -6887,6 +6896,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/native/core/src/execution/serde.rs b/native/core/src/execution/serde.rs index 09a0093fac6..34f9a976dc4 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 parquet::{arrow::PARQUET_FIELD_ID_META_KEY, variant::VariantType}; use prost::Message; use std::{io::Cursor, sync::Arc}; @@ -103,9 +104,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::Variant => ArrowDataType::Struct(Fields::from(vec![ Field::new("value", ArrowDataType::Binary, false), Field::new("metadata", ArrowDataType::Binary, false), 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}")), diff --git a/native/spark-expr/src/datetime_funcs/make_interval.rs b/native/spark-expr/src/datetime_funcs/make_interval.rs index 9f7b21ca2a7..2885f74aa1a 100644 --- a/native/spark-expr/src/datetime_funcs/make_interval.rs +++ b/native/spark-expr/src/datetime_funcs/make_interval.rs @@ -16,72 +16,270 @@ // under the License. use crate::arithmetic_overflow_error; -use arrow::array::Array; -use arrow::datatypes::DataType; -use datafusion::common::Result; -use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature}; -use datafusion_spark::function::datetime::make_interval::SparkMakeInterval as DataFusionMakeInterval; +use arrow::array::{Array, ArrayRef, Decimal128Array, Int32Array, Int64Array, StructArray}; +use arrow::buffer::NullBuffer; +use arrow::datatypes::{DataType, Field, Fields}; +use datafusion::common::{DataFusionError, Result}; +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), + ])) +} + +/// 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 { - 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, +) -> 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() + .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 { 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::() + .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::() + .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()); + 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), + ) { + Ok((months, days, micros)) => { + result_months.push(months); + result_days.push(days); + result_micros.push(micros); + valid.push(true); + } + Err(from_type) if self.fail_on_error => { + return Err(arithmetic_overflow_error(from_type).into()); + } + Err(_) => { + 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), + 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") + ); + } + + #[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 5330ba1bf14..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::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 f891087b246..171520198dc 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 55beb6cd378..ebe22423179 100644 --- a/native/spark-expr/src/hash_funcs/utils.rs +++ b/native/spark-expr/src/hash_funcs/utils.rs @@ -147,47 +147,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]); } } } @@ -826,7 +847,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; @@ -966,13 +987,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); } @@ -1084,6 +1098,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 @@ -1322,4 +1342,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 24697114620..762d6bf4303 100644 --- a/native/spark-expr/src/hash_funcs/xxhash64.rs +++ b/native/spark-expr/src/hash_funcs/xxhash64.rs @@ -355,6 +355,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 b1447e37d6f..35ce232b002 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -86,10 +86,11 @@ pub use comet_scalar_funcs::{ }; pub use csv_funcs::*; pub use datetime_funcs::{ - 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/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/codegen/CometBatchKernelCodegenInput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala index d8c6ef07d04..7a35e0e68d9 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, @@ -70,6 +70,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 = { @@ -110,11 +111,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 _ => "isNull" - } + val method = nullCheckMethod(spec) s" case $ord: return this.col$ord.$method(this.rowIdx);" } } @@ -147,7 +144,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 { @@ -426,6 +424,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" } @@ -437,9 +436,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;" @@ -472,6 +475,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 2d2c8f1b2d6..26670e0bc73 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -190,7 +190,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 @@ -237,11 +237,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/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/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index bf27ce2c552..4a4242fc68a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -21,12 +21,12 @@ 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, DivideDTInterval, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, MultiplyYMInterval, 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, DivideDTInterval, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, MultiplyYMInterval, 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 -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 @@ -901,35 +901,60 @@ object CometMakeDTInterval extends CometCodegenDispatch[MakeDTInterval] object CometDivideDTInterval extends CometCodegenDispatch[DivideDTInterval] 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).") + 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 getIncompatibleReasons(): Seq[String] = Seq(incompatReason) + override def getUnsupportedReasons(): Seq[String] = Seq(eagerArgumentReason) override def getSupportLevel(expr: MakeInterval): SupportLevel = - Incompatible(Some(incompatReason)) + 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, 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 optExpr = scalarFunctionExprToProtoWithReturnType( + val restoreArgumentTags = CometExplainInfo.saveTags(expr.children) + val childExprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) + 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/main/scala/org/apache/comet/serde/literals.scala b/spark/src/main/scala/org/apache/comet/serde/literals.scala index 0aa03ad6c58..41ea5684d20 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, YearMonthIntervalType} -import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String} +import org.apache.spark.unsafe.types.UTF8String import com.google.protobuf.ByteString @@ -350,39 +350,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] @@ -401,9 +394,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/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala index 1cab882df93..dd133022a65 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 28e39fa9b49..2851789fd97 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 @@ -96,6 +96,8 @@ private[arrow] object ArrowWriter { case (_: DayTimeIntervalType, vector: DurationVector) => new DurationWriter(vector) case (CalendarIntervalType, vector: IntervalMonthDayNanoVector) => new IntervalMonthDayNanoWriter(vector) + case (CalendarIntervalType, vector: StructVector) => + new CalendarIntervalStructWriter(vector) case (dt, _) => throw QueryExecutionErrors.notSupportTypeError(dt) } @@ -1501,6 +1503,33 @@ private[arrow] class StructWriter( } } +/** + * 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 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 = { + val interval = input.getInterval(ordinal) + valueVector.setIndexDefined(count) + months.setSafe(count, interval.months) + days.setSafe(count, interval.days) + microseconds.setSafe(count, interval.microseconds) + } +} + 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 c884344ef15..32b81d43f72 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 @@ -49,6 +49,7 @@ import org.apache.comet.shims.CometTypeShim import org.apache.comet.vector.CometVector object Utils extends CometTypeShim with Logging { + private val calendarIntervalStructKey = "SPARK::calendarInterval::struct" private val VariantExtensionName = "arrow.parquet.variant" def majorMinorPatchVersion(version: String): Option[(Int, Int, Int)] = @@ -83,6 +84,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 => Option(field.getMetadata) .flatMap(metadata => @@ -178,7 +181,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 => new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO) + 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}") @@ -218,12 +223,60 @@ 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) } } + /** + * Returns true only for the Spark-tagged `CalendarIntervalType` struct produced by + * [[toArrowField]]. Matching the field shape (`months: Int32`, `days: Int32`, `microseconds: + * 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) + } + 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/calendar_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql index 2f335f35127..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 @@ -91,3 +91,43 @@ 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 + +-- 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/resources/sql-tests/expressions/datetime/make_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_interval.sql index 952a9085ee9..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 @@ -15,7 +15,7 @@ -- specific language governing permissions and limitations -- under the License. --- Config: spark.comet.expression.MakeInterval.allowIncompatible=true +-- Config: spark.comet.shuffle.mode=native statement CREATE TABLE test_make_interval( @@ -29,12 +29,25 @@ 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), (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 @@ -43,7 +56,58 @@ 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 +-- 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 +) + +-- 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) @@ -51,17 +115,90 @@ 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 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 + +-- 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 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..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 @@ -15,27 +15,99 @@ -- 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 +-- 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: +-- 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(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) + +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 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 ee2d5e81605..00000000000 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_interval_dispatch.sql +++ /dev/null @@ -1,49 +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. - --- With allowIncompatible unset, MakeInterval uses Spark's JVM codegen dispatcher. - -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 ignore(https://github.com/apache/datafusion-comet/issues/5279) -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 f199a8fb66a..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 ignore(https://github.com/apache/datafusion-comet/issues/5279) -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/resources/sql-tests/expressions/math/routing_intervals_disabled.sql b/spark/src/test/resources/sql-tests/expressions/math/routing_intervals_disabled.sql index e9e0e2728f7..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 @@ -30,5 +30,13 @@ 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 + +-- 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 6f8dad1b404..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 @@ -30,5 +30,9 @@ 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 + +-- 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/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 5aabafe4022..2210c528f4e 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1477,24 +1477,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") } diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala index 400ec6fcd2f..ee190e34c2f 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) diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index c9371fbf6da..3a198b797a6 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -28,11 +28,12 @@ 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 import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec, CometProjectExec} +import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryCatalog} import org.apache.spark.sql.connector.catalog.functions.{BoundFunction, ScalarFunction, UnboundFunction} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper @@ -238,18 +239,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) @@ -259,7 +267,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() @@ -267,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("a closed allocateOutput vector releases all of its memory") { // Struct outputs leaked the children that StructVector's writer allocates in its // constructor. The List and Map cases make sure that these outputs do not start to leak. 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 1ff82a6d450..d1667882059 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 @@ -166,18 +166,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 ef73522577a..d4d5e6495ad 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, VectorLoader, VectorSchemaRoot, VectorUnloader} +import org.apache.arrow.vector.{BaseFixedWidthVector, BaseValueVector, BigIntVector, BitVector, DecimalVector, IntVector, VarCharVector, VectorLoader, VectorSchemaRoot, VectorUnloader} import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.dictionary.{Dictionary => ArrowDictionary} import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider @@ -72,21 +72,37 @@ 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, 2) 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] + 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 = new CometPlainVector(arrow, false) + 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() 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 4510f9d0ac1..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.{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 @@ -30,6 +33,36 @@ import org.apache.comet.vector.CometVector 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("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) 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") } } }