From a68a5954ec9a165f0e9e580452a85aa3cc0f2aae Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 1 Aug 2026 00:38:15 +0800 Subject: [PATCH 01/13] refactor: use Arrow casts for temporal conversions --- native/core/src/parquet/cast_column.rs | 154 +++--------------- .../src/conversion_funcs/temporal.rs | 79 +++++---- .../src/datetime_funcs/date_from_unix_date.rs | 37 +---- native/spark-expr/src/utils.rs | 30 +++- 4 files changed, 96 insertions(+), 204 deletions(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 1cc928d1d59..c9567af6dd7 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -15,11 +15,8 @@ // specific language governing permissions and limitations // under the License. use arrow::{ - array::{ - make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray, - TimestampMicrosecondArray, TimestampMillisecondArray, - }, - compute::CastOptions, + array::{make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray}, + compute::{cast_with_options, CastOptions}, datatypes::{DataType, FieldRef, Schema, TimeUnit}, record_batch::RecordBatch, }; @@ -142,40 +139,6 @@ fn relabel_array(array: ArrayRef, target_type: &DataType) -> ArrayRef { } } -/// Casts a Timestamp(Microsecond) array to Timestamp(Millisecond) by dividing values by 1000. -/// Preserves the timezone from the target type. -fn cast_timestamp_micros_to_millis_array( - array: &ArrayRef, - target_tz: Option>, -) -> ArrayRef { - let micros_array = array - .as_any() - .downcast_ref::() - .expect("Expected TimestampMicrosecondArray"); - - let millis_values: TimestampMillisecondArray = - arrow::compute::kernels::arity::unary(micros_array, |v| v / 1000); - - // Apply timezone if present - let result = if let Some(tz) = target_tz { - millis_values.with_timezone(tz) - } else { - millis_values - }; - - Arc::new(result) -} - -/// Casts a Timestamp(Microsecond) scalar to Timestamp(Millisecond) by dividing the value by 1000. -/// Preserves the timezone from the target type. -fn cast_timestamp_micros_to_millis_scalar( - opt_val: Option, - target_tz: Option>, -) -> ScalarValue { - let new_val = opt_val.map(|v| v / 1000); - ScalarValue::TimestampMillisecond(new_val, target_tz) -} - #[derive(Debug, Clone, Eq)] pub struct CometCastColumnExpr { /// The physical expression producing the value to cast. @@ -279,11 +242,15 @@ impl PhysicalExpr for CometCastColumnExpr { DataType::Timestamp(TimeUnit::Millisecond, target_tz), ) => match value { ColumnarValue::Array(array) => { - let casted = cast_timestamp_micros_to_millis_array(&array, target_tz.clone()); + // Arrow adjusts values when adding a timezone, but Spark only relabels them. + let source_type = DataType::Timestamp(TimeUnit::Microsecond, target_tz.clone()); + let array = relabel_array(array, &source_type); + let casted = cast_with_options(&array, target_field, &self.cast_options)?; Ok(ColumnarValue::Array(casted)) } - ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(opt_val, _)) => { - let casted = cast_timestamp_micros_to_millis_scalar(opt_val, target_tz.clone()); + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(value, _)) => { + let casted = ScalarValue::TimestampMicrosecond(value, target_tz.clone()) + .cast_to_with_options(target_field, &self.cast_options)?; Ok(ColumnarValue::Scalar(casted)) } _ => Ok(value), @@ -349,78 +316,12 @@ impl PhysicalExpr for CometCastColumnExpr { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, Int32Array, StringArray}; + use arrow::array::{ + Array, Int32Array, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray, + }; use arrow::datatypes::{Field, Fields}; use datafusion::physical_expr::expressions::Column; - #[test] - fn test_cast_timestamp_micros_to_millis_array() { - // Create a TimestampMicrosecond array with some values - let micros_array: TimestampMicrosecondArray = vec![ - Some(1_000_000), // 1 second in micros - Some(2_500_000), // 2.5 seconds in micros - None, // null value - Some(0), // zero - Some(-1_000_000), // negative value (before epoch) - ] - .into(); - let array_ref: ArrayRef = Arc::new(micros_array); - - // Cast without timezone - let result = cast_timestamp_micros_to_millis_array(&array_ref, None); - let millis_array = result - .as_any() - .downcast_ref::() - .expect("Expected TimestampMillisecondArray"); - - assert_eq!(millis_array.len(), 5); - assert_eq!(millis_array.value(0), 1000); // 1_000_000 / 1000 - assert_eq!(millis_array.value(1), 2500); // 2_500_000 / 1000 - assert!(millis_array.is_null(2)); - assert_eq!(millis_array.value(3), 0); - assert_eq!(millis_array.value(4), -1000); // -1_000_000 / 1000 - } - - #[test] - fn test_cast_timestamp_micros_to_millis_array_with_timezone() { - let micros_array: TimestampMicrosecondArray = vec![Some(1_000_000), Some(2_000_000)].into(); - let array_ref: ArrayRef = Arc::new(micros_array); - - let target_tz: Option> = Some(Arc::from("UTC")); - let result = cast_timestamp_micros_to_millis_array(&array_ref, target_tz); - let millis_array = result - .as_any() - .downcast_ref::() - .expect("Expected TimestampMillisecondArray"); - - assert_eq!(millis_array.value(0), 1000); - assert_eq!(millis_array.value(1), 2000); - // Verify timezone is preserved - assert_eq!( - result.data_type(), - &DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("UTC"))) - ); - } - - #[test] - fn test_cast_timestamp_micros_to_millis_scalar() { - // Test with a value - let result = cast_timestamp_micros_to_millis_scalar(Some(1_500_000), None); - assert_eq!(result, ScalarValue::TimestampMillisecond(Some(1500), None)); - - // Test with null - let null_result = cast_timestamp_micros_to_millis_scalar(None, None); - assert_eq!(null_result, ScalarValue::TimestampMillisecond(None, None)); - - // Test with timezone - let target_tz: Option> = Some(Arc::from("UTC")); - let tz_result = cast_timestamp_micros_to_millis_scalar(Some(2_000_000), target_tz.clone()); - assert_eq!( - tz_result, - ScalarValue::TimestampMillisecond(Some(2000), target_tz) - ); - } - #[test] fn test_comet_cast_column_expr_evaluate_micros_to_millis_array() { // Create input schema with TimestampMicrosecond column @@ -431,12 +332,8 @@ mod tests { )); let schema = Schema::new(vec![Arc::clone(&input_field)]); - // Create target field with TimestampMillisecond - let target_field = Arc::new(Field::new( - "ts", - DataType::Timestamp(TimeUnit::Millisecond, None), - true, - )); + let target_type = DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("+07:00"))); + let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); // Create a column expression let col_expr: Arc = Arc::new(Column::new("ts", 0)); @@ -446,7 +343,7 @@ mod tests { // Create a record batch with TimestampMicrosecond data let micros_array: TimestampMicrosecondArray = - vec![Some(1_000_000), Some(2_000_000), None].into(); + vec![Some(1_500_001), Some(-1_500_001), None].into(); let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(micros_array)]).unwrap(); // Evaluate @@ -458,9 +355,10 @@ mod tests { .as_any() .downcast_ref::() .expect("Expected TimestampMillisecondArray"); - assert_eq!(millis_array.value(0), 1000); - assert_eq!(millis_array.value(1), 2000); + assert_eq!(millis_array.value(0), 1500); + assert_eq!(millis_array.value(1), -1500); assert!(millis_array.is_null(2)); + assert_eq!(millis_array.data_type(), &target_type); } _ => panic!("Expected Array result"), } @@ -476,15 +374,11 @@ mod tests { )); let schema = Schema::new(vec![Arc::clone(&input_field)]); - // Create target field with TimestampMillisecond - let target_field = Arc::new(Field::new( - "ts", - DataType::Timestamp(TimeUnit::Millisecond, None), - true, - )); + let target_type = DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("+07:00"))); + let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); // Create a literal expression that returns a scalar - let scalar = ScalarValue::TimestampMicrosecond(Some(1_500_000), None); + let scalar = ScalarValue::TimestampMicrosecond(Some(-1_500_001), None); let literal_expr: Arc = Arc::new(datafusion::physical_expr::expressions::Literal::new(scalar)); @@ -499,7 +393,11 @@ mod tests { match result { ColumnarValue::Scalar(s) => { - assert_eq!(s, ScalarValue::TimestampMillisecond(Some(1500), None)); + assert_eq!( + s, + ScalarValue::TimestampMillisecond(Some(-1500), Some(Arc::from("+07:00"))) + ); + assert_eq!(s.data_type(), target_type); } _ => panic!("Expected Scalar result"), } diff --git a/native/spark-expr/src/conversion_funcs/temporal.rs b/native/spark-expr/src/conversion_funcs/temporal.rs index 96346962bc4..54024bbf991 100644 --- a/native/spark-expr/src/conversion_funcs/temporal.rs +++ b/native/spark-expr/src/conversion_funcs/temporal.rs @@ -18,7 +18,8 @@ use crate::utils::resolve_local_datetime; use crate::{timezone, SparkCastOptions, SparkResult}; use arrow::array::{ArrayRef, AsArray, TimestampMicrosecondBuilder}; -use arrow::datatypes::{DataType, Date32Type}; +use arrow::compute::{cast_with_options, CastOptions}; +use arrow::datatypes::{DataType, Date32Type, TimeUnit}; use chrono::NaiveDate; use std::str::FromStr; use std::sync::Arc; @@ -39,49 +40,45 @@ pub(crate) fn cast_date_to_timestamp( cast_options: &SparkCastOptions, target_tz: &Option>, ) -> SparkResult { + if target_tz.is_none() { + return Ok(cast_with_options( + array_ref, + &DataType::Timestamp(TimeUnit::Microsecond, None), + &CastOptions::default(), + )?); + } + let date_array = array_ref.as_primitive::(); let mut builder = TimestampMicrosecondBuilder::with_capacity(date_array.len()); - - if target_tz.is_none() { - // TIMESTAMP_NTZ: pure day arithmetic, no session-TZ offset. - // Matches Spark: daysToMicros(d, ZoneOffset.UTC) - for date in date_array.iter() { - match date { - Some(d) => builder.append_value((d as i64) * 86_400 * 1_000_000), - None => builder.append_null(), - } - } + // TIMESTAMP: midnight in session TZ → UTC epoch μs + let tz_str = if cast_options.timezone.is_empty() { + "UTC" } else { - // TIMESTAMP: midnight in session TZ → UTC epoch μs - let tz_str = if cast_options.timezone.is_empty() { - "UTC" - } else { - cast_options.timezone.as_str() - }; - // safe to unwrap since we are falling back to UTC above - let tz = timezone::Tz::from_str(tz_str)?; - let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); - for date in date_array.iter() { - match date { - Some(d) => { - // safe to unwrap since chrono's range ( 262,143 yrs) is higher than - // number of years possible with days as i32 (~ 6 mil yrs) - // convert date in session timezone to timestamp in UTC - let naive_date = epoch + chrono::Duration::days(d as i64); - let local_midnight = naive_date.and_hms_opt(0, 0, 0).unwrap(); - // Use resolve_local_datetime to correctly handle DST transitions: - // - Single: normal case, uses the given offset - // - Ambiguous (fall back): uses the earlier/DST occurrence, matching Spark - // - None (spring forward gap at midnight, e.g. America/Sao_Paulo): uses the - // pre-transition offset to compute the correct UTC time, matching Spark's - // LocalDate.atStartOfDay(zoneId) behaviour. - let local_midnight_in_microsec = - resolve_local_datetime(&tz, local_midnight).timestamp_micros(); - builder.append_value(local_midnight_in_microsec); - } - None => { - builder.append_null(); - } + cast_options.timezone.as_str() + }; + // safe to unwrap since we are falling back to UTC above + let tz = timezone::Tz::from_str(tz_str)?; + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + for date in date_array.iter() { + match date { + Some(d) => { + // safe to unwrap since chrono's range ( 262,143 yrs) is higher than + // number of years possible with days as i32 (~ 6 mil yrs) + // convert date in session timezone to timestamp in UTC + let naive_date = epoch + chrono::Duration::days(d as i64); + let local_midnight = naive_date.and_hms_opt(0, 0, 0).unwrap(); + // Use resolve_local_datetime to correctly handle DST transitions: + // - Single: normal case, uses the given offset + // - Ambiguous (fall back): uses the earlier/DST occurrence, matching Spark + // - None (spring forward gap at midnight, e.g. America/Sao_Paulo): uses the + // pre-transition offset to compute the correct UTC time, matching Spark's + // LocalDate.atStartOfDay(zoneId) behaviour. + let local_midnight_in_microsec = + resolve_local_datetime(&tz, local_midnight).timestamp_micros(); + builder.append_value(local_midnight_in_microsec); + } + None => { + builder.append_null(); } } } diff --git a/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs b/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs index 0e624e6472a..7977886ec6c 100644 --- a/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs +++ b/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs @@ -15,13 +15,12 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, Date32Array, Int32Array}; +use arrow::compute::{cast_with_options, CastOptions}; use arrow::datatypes::DataType; -use datafusion::common::{utils::take_function_args, DataFusionError, Result, ScalarValue}; +use datafusion::common::{utils::take_function_args, Result}; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; -use std::sync::Arc; /// Spark-compatible date_from_unix_date function. /// Converts an integer representing days since Unix epoch (1970-01-01) to a Date32 value. @@ -62,32 +61,14 @@ impl ScalarUDFImpl for SparkDateFromUnixDate { fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [unix_date] = take_function_args(self.name(), args.args)?; match unix_date { - ColumnarValue::Array(arr) => { - let int_array = arr.as_any().downcast_ref::().ok_or_else(|| { - DataFusionError::Execution( - "date_from_unix_date expects Int32Array input".to_string(), - ) - })?; - - // Date32 and Int32 both represent days since epoch, so we can directly - // reinterpret the values. The only operation needed is creating a Date32Array - // from the same underlying i32 values. - let date_array = - Date32Array::new(int_array.values().clone(), int_array.nulls().cloned()); - - Ok(ColumnarValue::Array(Arc::new(date_array))) + ColumnarValue::Array(arr) => Ok(ColumnarValue::Array(cast_with_options( + arr.as_ref(), + &DataType::Date32, + &CastOptions::default(), + )?)), + ColumnarValue::Scalar(scalar) => { + Ok(ColumnarValue::Scalar(scalar.cast_to(&DataType::Date32)?)) } - ColumnarValue::Scalar(scalar) => match scalar { - ScalarValue::Int32(Some(days)) => { - Ok(ColumnarValue::Scalar(ScalarValue::Date32(Some(days)))) - } - ScalarValue::Int32(None) | ScalarValue::Null => { - Ok(ColumnarValue::Scalar(ScalarValue::Date32(None))) - } - _ => Err(DataFusionError::Execution( - "date_from_unix_date expects Int32 scalar input".to_string(), - )), - }, } } diff --git a/native/spark-expr/src/utils.rs b/native/spark-expr/src/utils.rs index 7a785c72259..e7d788ba62a 100644 --- a/native/spark-expr/src/utils.rs +++ b/native/spark-expr/src/utils.rs @@ -29,7 +29,7 @@ use std::sync::Arc; use crate::timezone::Tz; use arrow::array::types::TimestampMillisecondType; -use arrow::array::TimestampMicrosecondArray; +use arrow::compute::{cast_with_options, CastOptions}; use arrow::datatypes::{MAX_DECIMAL128_FOR_EACH_PRECISION, MIN_DECIMAL128_FOR_EACH_PRECISION}; use arrow::error::ArrowError; use arrow::{ @@ -81,12 +81,8 @@ pub fn array_with_timezone( // so the result has the exact annotation the caller expects. timestamp_ntz_to_timestamp(array, timezone.as_str(), Some(target_tz.as_ref())) } - Some(DataType::Timestamp(TimeUnit::Microsecond, None)) => { - // Convert from Timestamp(Millisecond, None) to Timestamp(Microsecond, None) - let millis_array = as_primitive_array::(&array); - let micros_array: TimestampMicrosecondArray = - arrow::compute::kernels::arity::unary(millis_array, |v| v * 1000); - Ok(Arc::new(micros_array)) + Some(to_type @ DataType::Timestamp(TimeUnit::Microsecond, None)) => { + cast_with_options(array.as_ref(), to_type, &CastOptions::default()) } _ => { // Not supported @@ -376,6 +372,7 @@ pub fn unlikely(b: bool) -> bool { #[cfg(test)] mod tests { use super::*; + use arrow::array::{TimestampMicrosecondArray, TimestampMillisecondArray}; fn array_containing(local_datetime: &str) -> ArrayRef { let dt = NaiveDateTime::parse_from_str(local_datetime, "%Y-%m-%d %H:%M:%S").unwrap(); @@ -390,6 +387,25 @@ mod tests { .timestamp_micros() } + #[test] + fn test_array_with_timezone_millis_to_micros() { + let input: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![ + Some(1234), + Some(-1234), + None, + ])); + let target = DataType::Timestamp(TimeUnit::Microsecond, None); + + let output = array_with_timezone(input, "UTC".to_string(), Some(&target)).unwrap(); + let output = as_primitive_array::(&output); + + assert_eq!( + output.iter().collect::>(), + vec![Some(1_234_000), Some(-1_234_000), None] + ); + assert_eq!(output.timezone(), None); + } + #[test] fn test_build_bool_state() { let mut builder = BooleanBufferBuilder::new(0); From 263a61e8dc7b8a97f9f7216f0ea59437547b5012 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 2 Aug 2026 10:22:03 +0800 Subject: [PATCH 02/13] fix: preserve Spark temporal cast semantics --- native/core/src/parquet/cast_column.rs | 107 +++++++++--------- .../src/datetime_funcs/date_from_unix_date.rs | 6 +- native/spark-expr/src/utils.rs | 8 +- 3 files changed, 62 insertions(+), 59 deletions(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index c9567af6dd7..7a83c2350b9 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -15,8 +15,11 @@ // specific language governing permissions and limitations // under the License. use arrow::{ - array::{make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray}, - compute::{cast_with_options, CastOptions}, + array::{ + make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray, + TimestampMicrosecondArray, TimestampMillisecondArray, + }, + compute::CastOptions, datatypes::{DataType, FieldRef, Schema, TimeUnit}, record_batch::RecordBatch, }; @@ -242,16 +245,23 @@ impl PhysicalExpr for CometCastColumnExpr { DataType::Timestamp(TimeUnit::Millisecond, target_tz), ) => match value { ColumnarValue::Array(array) => { - // Arrow adjusts values when adding a timezone, but Spark only relabels them. - let source_type = DataType::Timestamp(TimeUnit::Microsecond, target_tz.clone()); - let array = relabel_array(array, &source_type); - let casted = cast_with_options(&array, target_field, &self.cast_options)?; - Ok(ColumnarValue::Array(casted)) + let micros = array + .as_any() + .downcast_ref::() + .expect("Expected TimestampMicrosecondArray"); + // Spark floors when downscaling negative timestamps; Arrow truncates. + let millis: TimestampMillisecondArray = + arrow::compute::kernels::arity::unary(micros, |v| v.div_euclid(1_000)); + // Applying the target timezone as metadata avoids shifting the values. + Ok(ColumnarValue::Array(Arc::new( + millis.with_timezone_opt(target_tz.clone()), + ))) } ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(value, _)) => { - let casted = ScalarValue::TimestampMicrosecond(value, target_tz.clone()) - .cast_to_with_options(target_field, &self.cast_options)?; - Ok(ColumnarValue::Scalar(casted)) + Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond( + value.map(|v| v.div_euclid(1_000)), + target_tz.clone(), + ))) } _ => Ok(value), }, @@ -316,51 +326,43 @@ impl PhysicalExpr for CometCastColumnExpr { #[cfg(test)] mod tests { use super::*; - use arrow::array::{ - Array, Int32Array, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray, - }; + use arrow::array::{Array, Int32Array, StringArray}; use arrow::datatypes::{Field, Fields}; use datafusion::physical_expr::expressions::Column; #[test] fn test_comet_cast_column_expr_evaluate_micros_to_millis_array() { - // Create input schema with TimestampMicrosecond column - let input_field = Arc::new(Field::new( - "ts", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - )); - let schema = Schema::new(vec![Arc::clone(&input_field)]); - - let target_type = DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("+07:00"))); - let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); - - // Create a column expression - let col_expr: Arc = Arc::new(Column::new("ts", 0)); - - // Create the CometCastColumnExpr - let cast_expr = CometCastColumnExpr::new(col_expr, input_field, target_field, None); - - // Create a record batch with TimestampMicrosecond data - let micros_array: TimestampMicrosecondArray = - vec![Some(1_500_001), Some(-1_500_001), None].into(); - let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(micros_array)]).unwrap(); - - // Evaluate - let result = cast_expr.evaluate(&batch).unwrap(); - - match result { - ColumnarValue::Array(arr) => { - let millis_array = arr - .as_any() - .downcast_ref::() - .expect("Expected TimestampMillisecondArray"); - assert_eq!(millis_array.value(0), 1500); - assert_eq!(millis_array.value(1), -1500); - assert!(millis_array.is_null(2)); - assert_eq!(millis_array.data_type(), &target_type); + for (source_tz, target_tz) in [ + (None, None), + (None, Some(Arc::from("+07:00"))), + (Some(Arc::from("UTC")), Some(Arc::from("America/New_York"))), + ] { + let input_type = DataType::Timestamp(TimeUnit::Microsecond, source_tz.clone()); + let input_field = Arc::new(Field::new("ts", input_type, true)); + let schema = Schema::new(vec![Arc::clone(&input_field)]); + let target_type = DataType::Timestamp(TimeUnit::Millisecond, target_tz); + let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); + let col_expr: Arc = Arc::new(Column::new("ts", 0)); + let cast_expr = CometCastColumnExpr::new(col_expr, input_field, target_field, None); + let micros_array = + TimestampMicrosecondArray::from(vec![Some(1_500_001), Some(-1_500_001), None]) + .with_timezone_opt(source_tz); + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(micros_array)]).unwrap(); + + match cast_expr.evaluate(&batch).unwrap() { + ColumnarValue::Array(arr) => { + let millis_array = arr + .as_any() + .downcast_ref::() + .expect("Expected TimestampMillisecondArray"); + assert_eq!(millis_array.value(0), 1500); + assert_eq!(millis_array.value(1), -1501); + assert!(millis_array.is_null(2)); + assert_eq!(millis_array.data_type(), &target_type); + } + _ => panic!("Expected Array result"), } - _ => panic!("Expected Array result"), } } @@ -374,7 +376,7 @@ mod tests { )); let schema = Schema::new(vec![Arc::clone(&input_field)]); - let target_type = DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("+07:00"))); + let target_type = DataType::Timestamp(TimeUnit::Millisecond, None); let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); // Create a literal expression that returns a scalar @@ -393,10 +395,7 @@ mod tests { match result { ColumnarValue::Scalar(s) => { - assert_eq!( - s, - ScalarValue::TimestampMillisecond(Some(-1500), Some(Arc::from("+07:00"))) - ); + assert_eq!(s, ScalarValue::TimestampMillisecond(Some(-1501), None)); assert_eq!(s.data_type(), target_type); } _ => panic!("Expected Scalar result"), diff --git a/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs b/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs index 7977886ec6c..c3c53fefcac 100644 --- a/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs +++ b/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs @@ -15,9 +15,9 @@ // specific language governing permissions and limitations // under the License. -use arrow::compute::{cast_with_options, CastOptions}; +use arrow::compute::cast_with_options; use arrow::datatypes::DataType; -use datafusion::common::{utils::take_function_args, Result}; +use datafusion::common::{format::DEFAULT_CAST_OPTIONS, utils::take_function_args, Result}; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; @@ -64,7 +64,7 @@ impl ScalarUDFImpl for SparkDateFromUnixDate { ColumnarValue::Array(arr) => Ok(ColumnarValue::Array(cast_with_options( arr.as_ref(), &DataType::Date32, - &CastOptions::default(), + &DEFAULT_CAST_OPTIONS, )?)), ColumnarValue::Scalar(scalar) => { Ok(ColumnarValue::Scalar(scalar.cast_to(&DataType::Date32)?)) diff --git a/native/spark-expr/src/utils.rs b/native/spark-expr/src/utils.rs index e7d788ba62a..a623c22cf5b 100644 --- a/native/spark-expr/src/utils.rs +++ b/native/spark-expr/src/utils.rs @@ -29,7 +29,7 @@ use std::sync::Arc; use crate::timezone::Tz; use arrow::array::types::TimestampMillisecondType; -use arrow::compute::{cast_with_options, CastOptions}; +use arrow::compute::cast_with_options; use arrow::datatypes::{MAX_DECIMAL128_FOR_EACH_PRECISION, MIN_DECIMAL128_FOR_EACH_PRECISION}; use arrow::error::ArrowError; use arrow::{ @@ -37,6 +37,7 @@ use arrow::{ temporal_conversions::as_datetime, }; use chrono::{DateTime, LocalResult, NaiveDateTime, Offset, TimeZone}; +use datafusion::common::format::DEFAULT_CAST_OPTIONS; /// Preprocesses input arrays to add timezone information from Spark to Arrow array datatype or /// to apply timezone offset. @@ -82,7 +83,7 @@ pub fn array_with_timezone( timestamp_ntz_to_timestamp(array, timezone.as_str(), Some(target_tz.as_ref())) } Some(to_type @ DataType::Timestamp(TimeUnit::Microsecond, None)) => { - cast_with_options(array.as_ref(), to_type, &CastOptions::default()) + cast_with_options(array.as_ref(), to_type, &DEFAULT_CAST_OPTIONS) } _ => { // Not supported @@ -404,6 +405,9 @@ mod tests { vec![Some(1_234_000), Some(-1_234_000), None] ); assert_eq!(output.timezone(), None); + + let overflow: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![i64::MAX])); + assert!(array_with_timezone(overflow, "UTC".to_string(), Some(&target)).is_err()); } #[test] From bd1d243f1980419eff75abd3afdfb682c5d2d29d Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 2 Aug 2026 10:36:51 +0800 Subject: [PATCH 03/13] test: match Spark micros-to-millis cases --- native/core/src/parquet/cast_column.rs | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 7a83c2350b9..5c12c569023 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -344,9 +344,14 @@ mod tests { let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); let col_expr: Arc = Arc::new(Column::new("ts", 0)); let cast_expr = CometCastColumnExpr::new(col_expr, input_field, target_field, None); - let micros_array = - TimestampMicrosecondArray::from(vec![Some(1_500_001), Some(-1_500_001), None]) - .with_timezone_opt(source_tz); + // Includes every case from Spark's DateTimeUtilsSuite.microsToMillis test. + let micros_array = TimestampMicrosecondArray::from(vec![ + Some(-9_223_372_036_844_776_001), + Some(-157_700_927_876_544), + Some(1_500_001), + None, + ]) + .with_timezone_opt(source_tz); let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(micros_array)]).unwrap(); @@ -356,9 +361,10 @@ mod tests { .as_any() .downcast_ref::() .expect("Expected TimestampMillisecondArray"); - assert_eq!(millis_array.value(0), 1500); - assert_eq!(millis_array.value(1), -1501); - assert!(millis_array.is_null(2)); + assert_eq!(millis_array.value(0), -9_223_372_036_844_777); + assert_eq!(millis_array.value(1), -157_700_927_877); + assert_eq!(millis_array.value(2), 1_500); + assert!(millis_array.is_null(3)); assert_eq!(millis_array.data_type(), &target_type); } _ => panic!("Expected Array result"), @@ -380,7 +386,7 @@ mod tests { let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); // Create a literal expression that returns a scalar - let scalar = ScalarValue::TimestampMicrosecond(Some(-1_500_001), None); + let scalar = ScalarValue::TimestampMicrosecond(Some(-157_700_927_876_544), None); let literal_expr: Arc = Arc::new(datafusion::physical_expr::expressions::Literal::new(scalar)); @@ -395,7 +401,10 @@ mod tests { match result { ColumnarValue::Scalar(s) => { - assert_eq!(s, ScalarValue::TimestampMillisecond(Some(-1501), None)); + assert_eq!( + s, + ScalarValue::TimestampMillisecond(Some(-157_700_927_877), None) + ); assert_eq!(s.data_type(), target_type); } _ => panic!("Expected Scalar result"), From fbb2a7cb821e730bf240d9e24ec7c98a5a33e271 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 2 Aug 2026 10:43:50 +0800 Subject: [PATCH 04/13] test: link Spark micros-to-millis cases --- native/core/src/parquet/cast_column.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 5c12c569023..74631d7d841 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -344,7 +344,8 @@ mod tests { let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); let col_expr: Arc = Arc::new(Column::new("ts", 0)); let cast_expr = CometCastColumnExpr::new(col_expr, input_field, target_field, None); - // Includes every case from Spark's DateTimeUtilsSuite.microsToMillis test. + // Matches the Spark v4.2.0 test cases: + // https://github.com/apache/spark/blob/32f7299601108917fb01920a54e084595b7b3bf8/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala#L969-L972 let micros_array = TimestampMicrosecondArray::from(vec![ Some(-9_223_372_036_844_776_001), Some(-157_700_927_876_544), From 47211ff59e721f23d9564d13cb16b7af365cbc72 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 2 Aug 2026 10:53:38 +0800 Subject: [PATCH 05/13] test: use Spark tag in source link --- native/core/src/parquet/cast_column.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 74631d7d841..56016f1d54c 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -344,8 +344,7 @@ mod tests { let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); let col_expr: Arc = Arc::new(Column::new("ts", 0)); let cast_expr = CometCastColumnExpr::new(col_expr, input_field, target_field, None); - // Matches the Spark v4.2.0 test cases: - // https://github.com/apache/spark/blob/32f7299601108917fb01920a54e084595b7b3bf8/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala#L969-L972 + // https://github.com/apache/spark/blob/v4.2.0/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala#L969-L972 let micros_array = TimestampMicrosecondArray::from(vec![ Some(-9_223_372_036_844_776_001), Some(-157_700_927_876_544), From 6afef50ed0cb1556c83156e7e0730e05e5d55a71 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 2 Aug 2026 11:27:36 +0800 Subject: [PATCH 06/13] Use imported arity kernel for Spark timestamp downscaling --- native/core/src/parquet/cast_column.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 56016f1d54c..6b000d0dc98 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -19,7 +19,7 @@ use arrow::{ make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray, TimestampMicrosecondArray, TimestampMillisecondArray, }, - compute::CastOptions, + compute::{kernels::arity, CastOptions}, datatypes::{DataType, FieldRef, Schema, TimeUnit}, record_batch::RecordBatch, }; @@ -250,8 +250,9 @@ impl PhysicalExpr for CometCastColumnExpr { .downcast_ref::() .expect("Expected TimestampMicrosecondArray"); // Spark floors when downscaling negative timestamps; Arrow truncates. + // [SparkDateTimeUtils.scala](https://github.com/apache/spark/blob/v4.2.0/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkDateTimeUtils.scala#L92-L101) let millis: TimestampMillisecondArray = - arrow::compute::kernels::arity::unary(micros, |v| v.div_euclid(1_000)); + arity::unary(micros, |v| v.div_euclid(1_000)); // Applying the target timezone as metadata avoids shifting the values. Ok(ColumnarValue::Array(Arc::new( millis.with_timezone_opt(target_tz.clone()), From 88685aae55331762e7869389a612368cb75b53b8 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Mon, 3 Aug 2026 02:57:46 +0800 Subject: [PATCH 07/13] fix: enforce Spark temporal conversion semantics --- native/core/src/parquet/cast_column.rs | 171 ++++++------------ native/core/src/parquet/schema_adapter.rs | 8 +- .../src/conversion_funcs/temporal.rs | 5 +- native/spark-expr/src/utils.rs | 6 + 4 files changed, 65 insertions(+), 125 deletions(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 6b000d0dc98..af3a2cc6aad 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -15,19 +15,15 @@ // specific language governing permissions and limitations // under the License. use arrow::{ - array::{ - make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray, - TimestampMicrosecondArray, TimestampMillisecondArray, - }, - compute::{kernels::arity, CastOptions}, + array::{make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray}, + compute::CastOptions, datatypes::{DataType, FieldRef, Schema, TimeUnit}, record_batch::RecordBatch, }; use crate::parquet::parquet_support::{spark_parquet_convert, SparkParquetOptions}; use datafusion::common::format::DEFAULT_CAST_OPTIONS; -use datafusion::common::Result as DataFusionResult; -use datafusion::common::ScalarValue; +use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::logical_expr::ColumnarValue; use datafusion::physical_expr::PhysicalExpr; use std::{ @@ -180,20 +176,41 @@ impl Hash for CometCastColumnExpr { } impl CometCastColumnExpr { - /// Create a new [`CometCastColumnExpr`]. - pub fn new( + /// Try to create a new [`CometCastColumnExpr`]. + pub fn try_new( expr: Arc, physical_field: FieldRef, target_field: FieldRef, cast_options: Option>, - ) -> Self { - Self { + ) -> DataFusionResult { + let physical_type = physical_field.data_type(); + let target_type = target_field.data_type(); + // `target_field` is the Spark logical field, while `physical_field` comes from the + // Parquet or Iceberg file. Comet represents Spark's TimestampType and TimestampNTZType + // as Arrow microseconds, and Spark maps both TIMESTAMP_MICROS and TIMESTAMP_MILLIS files + // to those logical types. A millisecond target is therefore invalid at this read-adapter + // boundary: + // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala#L318-L324 + if matches!( + (physical_type, target_type), + ( + DataType::Timestamp(TimeUnit::Microsecond, _), + DataType::Timestamp(TimeUnit::Millisecond, _) + ) + ) { + return Err(DataFusionError::Plan(format!( + "Cannot adapt Spark timestamp field '{}' from {physical_type} to {target_type}: Spark read schemas represent logical timestamps in microseconds", + physical_field.name() + ))); + } + + Ok(Self { expr, input_physical_field: physical_field, target_field, cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), parquet_options: None, - } + }) } /// Set Spark parquet options to enable complex nested type conversions. @@ -237,35 +254,7 @@ impl PhysicalExpr for CometCastColumnExpr { let input_physical_field = self.input_physical_field.data_type(); let target_field = self.target_field.data_type(); - // Handle specific type conversions with custom casts match (input_physical_field, target_field) { - // Timestamp(Microsecond) -> Timestamp(Millisecond) - ( - DataType::Timestamp(TimeUnit::Microsecond, _), - DataType::Timestamp(TimeUnit::Millisecond, target_tz), - ) => match value { - ColumnarValue::Array(array) => { - let micros = array - .as_any() - .downcast_ref::() - .expect("Expected TimestampMicrosecondArray"); - // Spark floors when downscaling negative timestamps; Arrow truncates. - // [SparkDateTimeUtils.scala](https://github.com/apache/spark/blob/v4.2.0/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkDateTimeUtils.scala#L92-L101) - let millis: TimestampMillisecondArray = - arity::unary(micros, |v| v.div_euclid(1_000)); - // Applying the target timezone as metadata avoids shifting the values. - Ok(ColumnarValue::Array(Arc::new( - millis.with_timezone_opt(target_tz.clone()), - ))) - } - ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(value, _)) => { - Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond( - value.map(|v| v.div_euclid(1_000)), - target_tz.clone(), - ))) - } - _ => Ok(value), - }, // Nested types that differ only in field names (e.g., List element named // "item" vs "element", or Map entries named "key_value" vs "entries"). // Re-label the array so the DataType metadata matches the logical schema. @@ -307,12 +296,12 @@ impl PhysicalExpr for CometCastColumnExpr { ) -> DataFusionResult> { assert_eq!(children.len(), 1); let child = children.pop().expect("CastColumnExpr child"); - let mut new_expr = Self::new( + let mut new_expr = Self::try_new( child, Arc::clone(&self.input_physical_field), Arc::clone(&self.target_field), Some(self.cast_options.clone()), - ); + )?; if let Some(opts) = &self.parquet_options { new_expr = new_expr.with_parquet_options(opts.clone()); } @@ -332,83 +321,27 @@ mod tests { use datafusion::physical_expr::expressions::Column; #[test] - fn test_comet_cast_column_expr_evaluate_micros_to_millis_array() { - for (source_tz, target_tz) in [ - (None, None), - (None, Some(Arc::from("+07:00"))), - (Some(Arc::from("UTC")), Some(Arc::from("America/New_York"))), - ] { - let input_type = DataType::Timestamp(TimeUnit::Microsecond, source_tz.clone()); - let input_field = Arc::new(Field::new("ts", input_type, true)); - let schema = Schema::new(vec![Arc::clone(&input_field)]); - let target_type = DataType::Timestamp(TimeUnit::Millisecond, target_tz); - let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); - let col_expr: Arc = Arc::new(Column::new("ts", 0)); - let cast_expr = CometCastColumnExpr::new(col_expr, input_field, target_field, None); - // https://github.com/apache/spark/blob/v4.2.0/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala#L969-L972 - let micros_array = TimestampMicrosecondArray::from(vec![ - Some(-9_223_372_036_844_776_001), - Some(-157_700_927_876_544), - Some(1_500_001), - None, - ]) - .with_timezone_opt(source_tz); - let batch = - RecordBatch::try_new(Arc::new(schema), vec![Arc::new(micros_array)]).unwrap(); - - match cast_expr.evaluate(&batch).unwrap() { - ColumnarValue::Array(arr) => { - let millis_array = arr - .as_any() - .downcast_ref::() - .expect("Expected TimestampMillisecondArray"); - assert_eq!(millis_array.value(0), -9_223_372_036_844_777); - assert_eq!(millis_array.value(1), -157_700_927_877); - assert_eq!(millis_array.value(2), 1_500); - assert!(millis_array.is_null(3)); - assert_eq!(millis_array.data_type(), &target_type); - } - _ => panic!("Expected Array result"), - } - } - } - - #[test] - fn test_comet_cast_column_expr_evaluate_micros_to_millis_scalar() { - // Create input schema with TimestampMicrosecond column - let input_field = Arc::new(Field::new( - "ts", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - )); - let schema = Schema::new(vec![Arc::clone(&input_field)]); - - let target_type = DataType::Timestamp(TimeUnit::Millisecond, None); - let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); - - // Create a literal expression that returns a scalar - let scalar = ScalarValue::TimestampMicrosecond(Some(-157_700_927_876_544), None); - let literal_expr: Arc = - Arc::new(datafusion::physical_expr::expressions::Literal::new(scalar)); - - // Create the CometCastColumnExpr - let cast_expr = CometCastColumnExpr::new(literal_expr, input_field, target_field, None); - - // Create an empty batch (scalar doesn't need data) - let batch = RecordBatch::new_empty(Arc::new(schema)); - - // Evaluate - let result = cast_expr.evaluate(&batch).unwrap(); - - match result { - ColumnarValue::Scalar(s) => { - assert_eq!( - s, - ScalarValue::TimestampMillisecond(Some(-157_700_927_877), None) - ); - assert_eq!(s.data_type(), target_type); - } - _ => panic!("Expected Scalar result"), + fn test_rejects_millisecond_logical_timestamp() { + for timezone in [None, Some(Arc::from("UTC"))] { + let input_field = Arc::new(Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, timezone.clone()), + true, + )); + let target_field = Arc::new(Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, timezone), + true, + )); + let expr: Arc = Arc::new(Column::new("ts", 0)); + + let err = CometCastColumnExpr::try_new(expr, input_field, target_field, None) + .expect_err("millisecond logical timestamp must be rejected during planning"); + assert!(matches!( + err, + DataFusionError::Plan(message) + if message.contains("Spark read schemas represent logical timestamps in microseconds") + )); } } diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index c6586b4681e..ccd13994b54 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -603,12 +603,12 @@ impl SparkPhysicalExprAdapter { } let cast_expr: Arc = Arc::new( - CometCastColumnExpr::new( + CometCastColumnExpr::try_new( remapped, Arc::clone(physical_field), Arc::clone(logical_field), None, - ) + )? .with_parquet_options(self.parquet_options.clone()), ); return Ok(Transformed::yes(cast_expr)); @@ -892,12 +892,12 @@ impl SparkPhysicalExprAdapter { | (DataType::Timestamp(_, _), DataType::Int64) ) { let comet_cast: Arc = Arc::new( - CometCastColumnExpr::new( + CometCastColumnExpr::try_new( child, input_field, Arc::clone(cast.target_field()), None, - ) + )? .with_parquet_options(self.parquet_options.clone()), ); return Ok(Transformed::yes(comet_cast)); diff --git a/native/spark-expr/src/conversion_funcs/temporal.rs b/native/spark-expr/src/conversion_funcs/temporal.rs index 54024bbf991..ed057d28ce3 100644 --- a/native/spark-expr/src/conversion_funcs/temporal.rs +++ b/native/spark-expr/src/conversion_funcs/temporal.rs @@ -18,9 +18,10 @@ use crate::utils::resolve_local_datetime; use crate::{timezone, SparkCastOptions, SparkResult}; use arrow::array::{ArrayRef, AsArray, TimestampMicrosecondBuilder}; -use arrow::compute::{cast_with_options, CastOptions}; +use arrow::compute::cast_with_options; use arrow::datatypes::{DataType, Date32Type, TimeUnit}; use chrono::NaiveDate; +use datafusion::common::format::DEFAULT_CAST_OPTIONS; use std::str::FromStr; use std::sync::Arc; @@ -44,7 +45,7 @@ pub(crate) fn cast_date_to_timestamp( return Ok(cast_with_options( array_ref, &DataType::Timestamp(TimeUnit::Microsecond, None), - &CastOptions::default(), + &DEFAULT_CAST_OPTIONS, )?); } diff --git a/native/spark-expr/src/utils.rs b/native/spark-expr/src/utils.rs index a623c22cf5b..f4d411b1721 100644 --- a/native/spark-expr/src/utils.rs +++ b/native/spark-expr/src/utils.rs @@ -83,6 +83,12 @@ pub fn array_with_timezone( timestamp_ntz_to_timestamp(array, timezone.as_str(), Some(target_tz.as_ref())) } Some(to_type @ DataType::Timestamp(TimeUnit::Microsecond, None)) => { + // This defensive conversion intentionally errors in every CAST eval mode: + // Spark's vectorized Parquet reader calls `millisToMicros` for both direct + // and dictionary values, independent of CAST evaluation. + // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java#L817-L833 + // `millisToMicros` uses `Math.multiplyExact`: + // https://github.com/apache/spark/blob/v4.2.0/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkDateTimeUtils.scala#L103-L108 cast_with_options(array.as_ref(), to_type, &DEFAULT_CAST_OPTIONS) } _ => { From 35920f6bf5e5f0dae3b327a44744eaf6c2685360 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 9 Aug 2026 01:49:01 +0800 Subject: [PATCH 08/13] andy's 3rd review --- native/core/src/parquet/cast_column.rs | 58 ++++++++++++++++++- native/core/src/parquet/parquet_support.rs | 20 ++++++- native/spark-expr/src/utils.rs | 35 +---------- .../comet/parquet/ParquetReadSuite.scala | 50 ++++++++++++++++ 4 files changed, 124 insertions(+), 39 deletions(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index af3a2cc6aad..fc5799f3166 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -188,8 +188,8 @@ impl CometCastColumnExpr { // `target_field` is the Spark logical field, while `physical_field` comes from the // Parquet or Iceberg file. Comet represents Spark's TimestampType and TimestampNTZType // as Arrow microseconds, and Spark maps both TIMESTAMP_MICROS and TIMESTAMP_MILLIS files - // to those logical types. A millisecond target is therefore invalid at this read-adapter - // boundary: + // to those logical types. For a top-level timestamp column, a millisecond target is + // therefore invalid at this read-adapter boundary: // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala#L318-L324 if matches!( (physical_type, target_type), @@ -316,9 +316,12 @@ impl PhysicalExpr for CometCastColumnExpr { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, Int32Array, StringArray}; + use arrow::array::{ + Array, Int32Array, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray, + }; use arrow::datatypes::{Field, Fields}; use datafusion::physical_expr::expressions::Column; + use datafusion_comet_spark_expr::EvalMode; #[test] fn test_rejects_millisecond_logical_timestamp() { @@ -345,6 +348,55 @@ mod tests { } } + #[test] + fn test_parquet_millis_to_micros_uses_checked_multiply() { + // Spark's Parquet reader calls the checked `millisToMicros` conversion for both + // direct and dictionary values: + // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java#L817-L833 + for eval_mode in [EvalMode::Legacy, EvalMode::Try, EvalMode::Ansi] { + for (source_tz, target_tz) in [ + (None, None), + (Some(Arc::from("UTC")), Some(Arc::from("UTC"))), + ] { + let input_field = Arc::new(Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, source_tz.clone()), + true, + )); + let schema = Arc::new(Schema::new(vec![Arc::clone(&input_field)])); + let target_type = DataType::Timestamp(TimeUnit::Microsecond, target_tz.clone()); + let target_field = Arc::new(Field::new("ts", target_type.clone(), true)); + let expr: Arc = Arc::new(Column::new("ts", 0)); + let cast_expr = CometCastColumnExpr::try_new(expr, input_field, target_field, None) + .unwrap() + .with_parquet_options(SparkParquetOptions::new(eval_mode, "UTC", false)); + + let input = TimestampMillisecondArray::from(vec![Some(1_234), Some(-1_234), None]) + .with_timezone_opt(source_tz.clone()); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(input)]).unwrap(); + let ColumnarValue::Array(output) = cast_expr.evaluate(&batch).unwrap() else { + panic!("Expected array result"); + }; + let output = output + .as_any() + .downcast_ref::() + .expect("Expected TimestampMicrosecondArray"); + assert_eq!( + output.iter().collect::>(), + vec![Some(1_234_000), Some(-1_234_000), None] + ); + assert_eq!(output.data_type(), &target_type); + + let overflow = + TimestampMillisecondArray::from(vec![i64::MAX]).with_timezone_opt(source_tz); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(overflow)]).unwrap(); + assert!(cast_expr.evaluate(&batch).is_err()); + } + } + } + #[test] fn test_relabel_list_field_name() { // Physical: List(Field("item", Int32)) diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 2ee1230ed87..54089d62d5a 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -22,8 +22,9 @@ use arrow::compute::can_cast_types; use arrow::datatypes::{FieldRef, Fields}; use arrow::{ array::{ - cast::AsArray, new_null_array, types::Int32Type, types::TimestampMicrosecondType, Array, - ArrayRef, DictionaryArray, StructArray, + cast::AsArray, new_null_array, types::Int32Type, types::TimestampMicrosecondType, + types::TimestampMillisecondType, Array, ArrayRef, ArrowNativeTypeOp, DictionaryArray, + StructArray, }, compute::{cast_with_options, take, CastOptions}, datatypes::{DataType, TimeUnit}, @@ -220,6 +221,21 @@ fn parquet_convert_array( list_arr.nulls().cloned(), ))) } + ( + Timestamp(TimeUnit::Millisecond, _), + Timestamp(TimeUnit::Microsecond, target_tz), + ) => { + // Spark's Parquet reader calls the checked `millisToMicros` conversion for both + // direct and dictionary values, independent of CAST evaluation mode: + // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java#L817-L833 + // `millisToMicros` uses `Math.multiplyExact`: + // https://github.com/apache/spark/blob/v4.2.0/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkDateTimeUtils.scala#L103-L108 + let micros = array + .as_primitive::() + .try_unary::<_, TimestampMicrosecondType, _>(|value| value.mul_checked(1_000))? + .with_timezone_opt(target_tz.clone()); + Ok(Arc::new(micros)) + } (Timestamp(TimeUnit::Microsecond, None), Timestamp(TimeUnit::Microsecond, Some(tz))) => { Ok(Arc::new( array diff --git a/native/spark-expr/src/utils.rs b/native/spark-expr/src/utils.rs index f4d411b1721..2b1eafdb33e 100644 --- a/native/spark-expr/src/utils.rs +++ b/native/spark-expr/src/utils.rs @@ -29,7 +29,6 @@ use std::sync::Arc; use crate::timezone::Tz; use arrow::array::types::TimestampMillisecondType; -use arrow::compute::cast_with_options; use arrow::datatypes::{MAX_DECIMAL128_FOR_EACH_PRECISION, MIN_DECIMAL128_FOR_EACH_PRECISION}; use arrow::error::ArrowError; use arrow::{ @@ -37,7 +36,6 @@ use arrow::{ temporal_conversions::as_datetime, }; use chrono::{DateTime, LocalResult, NaiveDateTime, Offset, TimeZone}; -use datafusion::common::format::DEFAULT_CAST_OPTIONS; /// Preprocesses input arrays to add timezone information from Spark to Arrow array datatype or /// to apply timezone offset. @@ -82,15 +80,6 @@ pub fn array_with_timezone( // so the result has the exact annotation the caller expects. timestamp_ntz_to_timestamp(array, timezone.as_str(), Some(target_tz.as_ref())) } - Some(to_type @ DataType::Timestamp(TimeUnit::Microsecond, None)) => { - // This defensive conversion intentionally errors in every CAST eval mode: - // Spark's vectorized Parquet reader calls `millisToMicros` for both direct - // and dictionary values, independent of CAST evaluation. - // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java#L817-L833 - // `millisToMicros` uses `Math.multiplyExact`: - // https://github.com/apache/spark/blob/v4.2.0/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkDateTimeUtils.scala#L103-L108 - cast_with_options(array.as_ref(), to_type, &DEFAULT_CAST_OPTIONS) - } _ => { // Not supported Err(ArrowError::CastError(format!( @@ -379,7 +368,7 @@ pub fn unlikely(b: bool) -> bool { #[cfg(test)] mod tests { use super::*; - use arrow::array::{TimestampMicrosecondArray, TimestampMillisecondArray}; + use arrow::array::TimestampMicrosecondArray; fn array_containing(local_datetime: &str) -> ArrayRef { let dt = NaiveDateTime::parse_from_str(local_datetime, "%Y-%m-%d %H:%M:%S").unwrap(); @@ -394,28 +383,6 @@ mod tests { .timestamp_micros() } - #[test] - fn test_array_with_timezone_millis_to_micros() { - let input: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![ - Some(1234), - Some(-1234), - None, - ])); - let target = DataType::Timestamp(TimeUnit::Microsecond, None); - - let output = array_with_timezone(input, "UTC".to_string(), Some(&target)).unwrap(); - let output = as_primitive_array::(&output); - - assert_eq!( - output.iter().collect::>(), - vec![Some(1_234_000), Some(-1_234_000), None] - ); - assert_eq!(output.timezone(), None); - - let overflow: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![i64::MAX])); - assert!(array_with_timezone(overflow, "UTC".to_string(), Some(&target)).is_err()); - } - #[test] fn test_build_bool_state() { let mut builder = BooleanBufferBuilder::new(0); diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index a5fc92a5ebb..158757df810 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -202,6 +202,56 @@ abstract class ParquetReadSuite extends CometTestBase { } } + test("TIMESTAMP_MILLIS overflow fails in native scan") { + // Spark routes both TimestampType and TimestampNTZType through LongAsMicrosUpdater: + // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java#L140-L164 + // The updater calls checked millisToMicros for direct and dictionary values: + // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java#L800-L833 + // Matches Spark's positive and negative overflow cases: + // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/test/resources/sql-tests/inputs/timestamp.sql#L74-L83 + def isOverflow(error: Throwable): Boolean = + Iterator + .iterate(error)(_.getCause) + .takeWhile(_ != null) + .exists(cause => Option(cause.getMessage).exists(_.toLowerCase.contains("overflow"))) + + Seq(false, true).foreach { dictionaryEnabled => + Seq(92233720368547758L, -92233720368547758L).foreach { millis => + withTempDir { dir => + val path = new Path(dir.toURI.toString, "part-r-0.parquet") + val schema = MessageTypeParser.parseMessageType(""" + |message root { + | optional int64 ts(TIMESTAMP_MILLIS); + | optional int64 ts_ntz(TIMESTAMP(MILLIS,false)); + |} + |""".stripMargin) + val writer = createParquetWriter(schema, path, dictionaryEnabled) + val record = new SimpleGroup(schema) + record.add(0, millis) + record.add(1, millis) + writer.write(record) + writer.close() + + Seq(false, true).foreach { ansiEnabled => + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) { + readParquetFile(path.toString) { df => + Seq("ts", "ts_ntz").foreach { column => + val selected = df.select(column) + assert(collect(selected.queryExecution.executedPlan) { + case _: CometNativeScanExec => true + }.nonEmpty) + + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(selected) + assert(Seq(sparkError, cometError).forall(_.exists(isOverflow))) + } + } + } + } + } + } + } + } + test("timestamp as int96") { import testImplicits._ From ae0152f5466d5eed4801940fedf8886c476bc9c4 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 28 Aug 2026 09:45:07 +0800 Subject: [PATCH 09/13] address review: drop temporal.rs refactor, improve adapter error message - Revert cast_date_to_timestamp to main's version to avoid conflicting with #5457, which rewrites the same function as a safety fix - Document the Int32 input guarantee and zero-copy reinterpret path in date_from_unix_date - Make the microsecond-physical/millisecond-target rejection message actionable (report link + scan fallback workaround) Co-Authored-By: Claude Fable 5 --- native/core/src/parquet/cast_column.rs | 2 +- .../src/conversion_funcs/temporal.rs | 80 ++++++++++--------- .../src/datetime_funcs/date_from_unix_date.rs | 3 + 3 files changed, 45 insertions(+), 40 deletions(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index fc5799f3166..231f88c5074 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -199,7 +199,7 @@ impl CometCastColumnExpr { ) ) { return Err(DataFusionError::Plan(format!( - "Cannot adapt Spark timestamp field '{}' from {physical_type} to {target_type}: Spark read schemas represent logical timestamps in microseconds", + "Cannot adapt Spark timestamp field '{}' from {physical_type} to {target_type}: Spark read schemas represent logical timestamps in microseconds. This indicates a bug in Comet's schema handling; please report it at https://github.com/apache/datafusion-comet/issues. As a workaround, set spark.comet.scan.enabled=false", physical_field.name() ))); } diff --git a/native/spark-expr/src/conversion_funcs/temporal.rs b/native/spark-expr/src/conversion_funcs/temporal.rs index ed057d28ce3..96346962bc4 100644 --- a/native/spark-expr/src/conversion_funcs/temporal.rs +++ b/native/spark-expr/src/conversion_funcs/temporal.rs @@ -18,10 +18,8 @@ use crate::utils::resolve_local_datetime; use crate::{timezone, SparkCastOptions, SparkResult}; use arrow::array::{ArrayRef, AsArray, TimestampMicrosecondBuilder}; -use arrow::compute::cast_with_options; -use arrow::datatypes::{DataType, Date32Type, TimeUnit}; +use arrow::datatypes::{DataType, Date32Type}; use chrono::NaiveDate; -use datafusion::common::format::DEFAULT_CAST_OPTIONS; use std::str::FromStr; use std::sync::Arc; @@ -41,45 +39,49 @@ pub(crate) fn cast_date_to_timestamp( cast_options: &SparkCastOptions, target_tz: &Option>, ) -> SparkResult { - if target_tz.is_none() { - return Ok(cast_with_options( - array_ref, - &DataType::Timestamp(TimeUnit::Microsecond, None), - &DEFAULT_CAST_OPTIONS, - )?); - } - let date_array = array_ref.as_primitive::(); let mut builder = TimestampMicrosecondBuilder::with_capacity(date_array.len()); - // TIMESTAMP: midnight in session TZ → UTC epoch μs - let tz_str = if cast_options.timezone.is_empty() { - "UTC" - } else { - cast_options.timezone.as_str() - }; - // safe to unwrap since we are falling back to UTC above - let tz = timezone::Tz::from_str(tz_str)?; - let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); - for date in date_array.iter() { - match date { - Some(d) => { - // safe to unwrap since chrono's range ( 262,143 yrs) is higher than - // number of years possible with days as i32 (~ 6 mil yrs) - // convert date in session timezone to timestamp in UTC - let naive_date = epoch + chrono::Duration::days(d as i64); - let local_midnight = naive_date.and_hms_opt(0, 0, 0).unwrap(); - // Use resolve_local_datetime to correctly handle DST transitions: - // - Single: normal case, uses the given offset - // - Ambiguous (fall back): uses the earlier/DST occurrence, matching Spark - // - None (spring forward gap at midnight, e.g. America/Sao_Paulo): uses the - // pre-transition offset to compute the correct UTC time, matching Spark's - // LocalDate.atStartOfDay(zoneId) behaviour. - let local_midnight_in_microsec = - resolve_local_datetime(&tz, local_midnight).timestamp_micros(); - builder.append_value(local_midnight_in_microsec); + + if target_tz.is_none() { + // TIMESTAMP_NTZ: pure day arithmetic, no session-TZ offset. + // Matches Spark: daysToMicros(d, ZoneOffset.UTC) + for date in date_array.iter() { + match date { + Some(d) => builder.append_value((d as i64) * 86_400 * 1_000_000), + None => builder.append_null(), } - None => { - builder.append_null(); + } + } else { + // TIMESTAMP: midnight in session TZ → UTC epoch μs + let tz_str = if cast_options.timezone.is_empty() { + "UTC" + } else { + cast_options.timezone.as_str() + }; + // safe to unwrap since we are falling back to UTC above + let tz = timezone::Tz::from_str(tz_str)?; + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + for date in date_array.iter() { + match date { + Some(d) => { + // safe to unwrap since chrono's range ( 262,143 yrs) is higher than + // number of years possible with days as i32 (~ 6 mil yrs) + // convert date in session timezone to timestamp in UTC + let naive_date = epoch + chrono::Duration::days(d as i64); + let local_midnight = naive_date.and_hms_opt(0, 0, 0).unwrap(); + // Use resolve_local_datetime to correctly handle DST transitions: + // - Single: normal case, uses the given offset + // - Ambiguous (fall back): uses the earlier/DST occurrence, matching Spark + // - None (spring forward gap at midnight, e.g. America/Sao_Paulo): uses the + // pre-transition offset to compute the correct UTC time, matching Spark's + // LocalDate.atStartOfDay(zoneId) behaviour. + let local_midnight_in_microsec = + resolve_local_datetime(&tz, local_midnight).timestamp_micros(); + builder.append_value(local_midnight_in_microsec); + } + None => { + builder.append_null(); + } } } } diff --git a/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs b/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs index c3c53fefcac..2d08fd2b768 100644 --- a/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs +++ b/native/spark-expr/src/datetime_funcs/date_from_unix_date.rs @@ -60,6 +60,9 @@ impl ScalarUDFImpl for SparkDateFromUnixDate { fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [unix_date] = take_function_args(self.name(), args.args)?; + // The input is guaranteed to be Int32 by `Signature::exact` and the Comet serde + // (Spark's DateFromUnixDate only accepts IntegerType), so the casts below take + // Arrow's zero-copy `Int32 -> Date32` reinterpret path. match unix_date { ColumnarValue::Array(arr) => Ok(ColumnarValue::Array(cast_with_options( arr.as_ref(), From afa3673da3d571b59557ffc5a1ee5e976d745f41 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 28 Aug 2026 09:48:29 +0800 Subject: [PATCH 10/13] chore: remove unused imports in parquet_support Co-Authored-By: Claude Fable 5 --- native/core/src/parquet/parquet_support.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 09ffc1aaf40..9842d6220f5 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -22,9 +22,8 @@ use arrow::compute::can_cast_types; use arrow::datatypes::{FieldRef, Fields}; use arrow::{ array::{ - cast::AsArray, new_null_array, types::Int32Type, types::TimestampMicrosecondType, - types::TimestampMillisecondType, Array, ArrayRef, ArrowNativeTypeOp, DictionaryArray, - StructArray, + cast::AsArray, new_null_array, types::TimestampMicrosecondType, + types::TimestampMillisecondType, Array, ArrayRef, ArrowNativeTypeOp, StructArray, }, compute::{cast_with_options, CastOptions}, datatypes::{DataType, TimeUnit}, From c833b72a66749fbf48d52820c4b16ad0de33cc9a Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 29 Aug 2026 01:04:52 +0800 Subject: [PATCH 11/13] test: exercise dictionary-encoded pages in TIMESTAMP_MILLIS overflow regression A single row falls back to PLAIN even with dictionary encoding enabled, so the dictionary leg never covered a dictionary read. Write 16 repeated rows and assert hasDictionaryEncodedPages matches the writer setting. Co-Authored-By: Claude Fable 5 --- .../comet/parquet/ParquetReadSuite.scala | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index e5048eb0a6d..fa431cfec1f 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -285,12 +285,33 @@ abstract class ParquetReadSuite extends CometTestBase { |} |""".stripMargin) val writer = createParquetWriter(schema, path, dictionaryEnabled) - val record = new SimpleGroup(schema) - record.add(0, millis) - record.add(1, millis) - writer.write(record) + // A single row falls back to PLAIN even with dictionary encoding enabled, because + // a one-entry dictionary page is not smaller than the raw values. Write enough + // repeated rows for the writer to keep dictionary-encoded data pages. + (0 until 16).foreach { _ => + val record = new SimpleGroup(schema) + record.add(0, millis) + record.add(1, millis) + writer.write(record) + } writer.close() + val footerReader = org.apache.parquet.hadoop.ParquetFileReader.open( + org.apache.parquet.hadoop.util.HadoopInputFile + .fromPath(path, spark.sessionState.newHadoopConf())) + try { + footerReader.getFooter.getBlocks.forEach { block => + block.getColumns.forEach { column => + assert( + column.getEncodingStats.hasDictionaryEncodedPages == dictionaryEnabled, + s"expected hasDictionaryEncodedPages=$dictionaryEnabled " + + s"for column ${column.getPath}") + } + } + } finally { + footerReader.close() + } + Seq(false, true).foreach { ansiEnabled => withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) { readParquetFile(path.toString) { df => From d46519298fb2366fdfa00906a185186c1d374cf4 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 29 Aug 2026 16:59:56 +0800 Subject: [PATCH 12/13] fix: preserve timestamp pruning by rewriting predicates to the millisecond domain The predicate over a TIMESTAMP_MILLIS file column was wrapped in CometCastColumnExpr, which DataFusion's pruning analyzer cannot see through, so row groups Spark prunes from millisecond statistics were read and hit the checked millis->micros conversion. Rewrite predicate comparisons into the millisecond domain (exact integer rescaling of the literal, mirroring Spark's ParquetFilters) and unwrap IS NULL checks, so pruning works and predicate evaluation never converts file values. The scan output conversion stays checked. Co-Authored-By: Claude Fable 5 --- native/core/src/parquet/cast_column.rs | 15 ++ native/core/src/parquet/schema_adapter.rs | 250 +++++++++++++++++- .../comet/parquet/ParquetReadSuite.scala | 42 +++ 3 files changed, 305 insertions(+), 2 deletions(-) diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 231f88c5074..8501420089a 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -214,6 +214,21 @@ impl CometCastColumnExpr { } /// Set Spark parquet options to enable complex nested type conversions. + /// The wrapped input expression. + pub fn expr(&self) -> &Arc { + &self.expr + } + + /// The Parquet/Iceberg file field this expression reads from. + pub fn input_physical_field(&self) -> &FieldRef { + &self.input_physical_field + } + + /// The Spark logical field this expression converts to. + pub fn target_field(&self) -> &FieldRef { + &self.target_field + } + pub fn with_parquet_options(mut self, options: SparkParquetOptions) -> Self { self.parquet_options = Some(options); self diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index ccd13994b54..c956184c400 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -18,11 +18,14 @@ use crate::parquet::cast_column::CometCastColumnExpr; use crate::parquet::parquet_support::{spark_parquet_convert, SparkParquetOptions}; use arrow::array::new_empty_array; -use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef, TimeUnit}; use arrow::record_batch::RecordBatch; use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion::common::{DataFusionError, Result as DataFusionResult}; -use datafusion::physical_expr::expressions::Column; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{ + BinaryExpr, Column, IsNotNullExpr, IsNullExpr, Literal, +}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ColumnarValue; use datafusion::scalar::ScalarValue; @@ -496,6 +499,15 @@ impl PhysicalExprAdapter for SparkPhysicalExprAdapter { } }; + // Rewrite predicate comparisons over a millisecond timestamp file column into + // the millisecond domain so DataFusion can prune row groups from statistics and + // the row filter never runs the checked millis->micros conversion on values the + // filter would discard. Values that survive the filter still go through the + // checked conversion when the scan output is adapted. + let expr = expr + .transform(|e| self.rewrite_millis_timestamp_comparison(e)) + .data()?; + // For case-insensitive mode: remap column names from logical back to // original physical names. The default adapter was given a remapped // physical schema (with logical names) so it could find columns. But @@ -522,7 +534,114 @@ impl PhysicalExprAdapter for SparkPhysicalExprAdapter { } } +/// The wrapped expression and the file column's timezone of a millis->micros cast. +type MillisCastParts = (Arc, Option>); + +/// If `expr` is a checked millis->micros [`CometCastColumnExpr`], return its wrapped +/// expression and the file column's timezone. +fn as_millis_to_micros_cast(expr: &Arc) -> Option { + let cast = expr.downcast_ref::()?; + match ( + cast.input_physical_field().data_type(), + cast.target_field().data_type(), + ) { + ( + DataType::Timestamp(TimeUnit::Millisecond, file_tz), + DataType::Timestamp(TimeUnit::Microsecond, _), + ) => Some((Arc::clone(cast.expr()), file_tz.clone())), + _ => None, + } +} + +/// If `expr` is a non-null microsecond timestamp literal, return its value. +fn as_micros_literal(expr: &Arc) -> Option { + match expr.downcast_ref::()?.value() { + ScalarValue::TimestampMicrosecond(Some(micros), _) => Some(*micros), + _ => None, + } +} + impl SparkPhysicalExprAdapter { + /// Rewrite predicate expressions over a `TIMESTAMP_MILLIS` file column read as + /// microseconds into the millisecond domain, mirroring Spark's `ParquetFilters`, + /// which pushes timestamp predicates down in the file's physical unit: + /// https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilters.scala#L192-L196 + /// + /// The wrapped `CometCastColumnExpr` is opaque to DataFusion's pruning-predicate + /// analyzer, so leaving the conversion inside the predicate defeats row-group + /// statistics pruning — and the checked conversion then raises overflow errors for + /// values Spark never reads, because Spark prunes them from the millisecond + /// statistics. Comparing raw millisecond values against a rescaled literal is exact + /// (`m * 1000 OP lit` over the integers), never overflows, and DataFusion can prune + /// with it. Values that survive the filter still go through the checked conversion + /// when the scan output is adapted, so genuine reads of overflowing values keep + /// failing like Spark's `millisToMicros`. + fn rewrite_millis_timestamp_comparison( + &self, + expr: Arc, + ) -> DataFusionResult>> { + // `ts IS NULL` / `ts IS NOT NULL`: the checked conversion preserves null-ness + // (it errors rather than producing nulls), so test the raw column directly. + if let Some(is_null) = expr.downcast_ref::() { + if let Some((inner, _)) = as_millis_to_micros_cast(is_null.arg()) { + return Ok(Transformed::yes(Arc::new(IsNullExpr::new(inner)))); + } + } + if let Some(is_not_null) = expr.downcast_ref::() { + if let Some((inner, _)) = as_millis_to_micros_cast(is_not_null.arg()) { + return Ok(Transformed::yes(Arc::new(IsNotNullExpr::new(inner)))); + } + } + + let Some(binary) = expr.downcast_ref::() else { + return Ok(Transformed::no(expr)); + }; + let matched = if let (Some(cast), Some(micros)) = ( + as_millis_to_micros_cast(binary.left()), + as_micros_literal(binary.right()), + ) { + Some((cast, *binary.op(), micros)) + } else if let (Some(micros), Some(cast)) = ( + as_micros_literal(binary.left()), + as_millis_to_micros_cast(binary.right()), + ) { + binary.op().swap().map(|op| (cast, op, micros)) + } else { + None + }; + let Some(((inner, file_tz), op, micros)) = matched else { + return Ok(Transformed::no(expr)); + }; + + // For a file value `m` (milliseconds) the logical value is exactly `m * 1000` + // microseconds, so `m * 1000 OP L` rewrites to an exact comparison on `m`. + let floor = micros.div_euclid(1_000); + let ceil = floor + i64::from(micros.rem_euclid(1_000) != 0); + let exact = micros % 1_000 == 0; + let (op, millis) = match op { + Operator::Lt => (Operator::Lt, ceil), + Operator::LtEq => (Operator::LtEq, floor), + Operator::Gt => (Operator::Gt, floor), + Operator::GtEq => (Operator::GtEq, ceil), + Operator::Eq if exact => (Operator::Eq, floor), + Operator::NotEq if exact => (Operator::NotEq, floor), + // A sub-millisecond literal can never equal `m * 1000`. `col < i64::MIN` + // (resp. `col >= i64::MIN`) is false (resp. true) for every non-null value + // and NULL for nulls, matching `=` / `!=` null semantics while remaining a + // plain, prunable comparison. + Operator::Eq => (Operator::Lt, i64::MIN), + Operator::NotEq => (Operator::GtEq, i64::MIN), + _ => return Ok(Transformed::no(expr)), + }; + let literal = Arc::new(Literal::new(ScalarValue::TimestampMillisecond( + Some(millis), + file_tz, + ))); + Ok(Transformed::yes(Arc::new(BinaryExpr::new( + inner, op, literal, + )))) + } + /// Wrap ALL Column expressions that have type mismatches with CometCastColumnExpr. /// This is the fallback path when the default adapter fails (e.g., for complex /// nested type casts like List or Map). Uses `spark_parquet_convert` @@ -1775,4 +1894,131 @@ mod test { "Expected duplicate field error, got: {err_msg}" ); } + + use arrow::datatypes::TimeUnit; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{ + BinaryExpr, Column, IsNotNullExpr, IsNullExpr, Literal, + }; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::scalar::ScalarValue; + use datafusion_physical_expr_adapter::PhysicalExprAdapter; + + fn millis_file_adapter() -> Arc { + let tz: Arc = Arc::from("UTC"); + let logical = Arc::new(Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::clone(&tz))), + true, + )])); + let physical = Arc::new(Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, Some(tz)), + true, + )])); + SparkPhysicalExprAdapterFactory::new( + SparkParquetOptions::new(EvalMode::Legacy, "UTC", false), + None, + ) + .create(logical, physical) + .unwrap() + } + + fn micros_lit(value: i64) -> Arc { + Arc::new(Literal::new(ScalarValue::TimestampMicrosecond( + Some(value), + Some(Arc::from("UTC")), + ))) + } + + fn assert_millis_comparison( + rewritten: &Arc, + expected_op: Operator, + expected_millis: i64, + ) { + let binary = rewritten + .downcast_ref::() + .expect("expected BinaryExpr"); + assert!( + binary.left().downcast_ref::().is_some(), + "expected raw column on the left, got {binary}" + ); + assert_eq!(*binary.op(), expected_op); + let literal = binary + .right() + .downcast_ref::() + .expect("expected Literal"); + assert_eq!( + literal.value(), + &ScalarValue::TimestampMillisecond(Some(expected_millis), Some(Arc::from("UTC"))) + ); + } + + #[test] + fn test_millis_timestamp_predicate_rewrites_to_millis_domain() { + let adapter = millis_file_adapter(); + let col: Arc = Arc::new(Column::new("ts", 0)); + + // (predicate op, literal micros, expected op, expected millis) + let cases = vec![ + // Exact multiples of 1000 rescale directly. + (Operator::Lt, 2_000, Operator::Lt, 2), + (Operator::LtEq, 2_000, Operator::LtEq, 2), + (Operator::Gt, 2_000, Operator::Gt, 2), + (Operator::GtEq, 2_000, Operator::GtEq, 2), + (Operator::Eq, 2_000, Operator::Eq, 2), + (Operator::NotEq, 2_000, Operator::NotEq, 2), + // m * 1000 < 1500 iff m < 2, and m * 1000 <= 1500 iff m <= 1. + (Operator::Lt, 1_500, Operator::Lt, 2), + (Operator::LtEq, 1_500, Operator::LtEq, 1), + (Operator::Gt, 1_500, Operator::Gt, 1), + (Operator::GtEq, 1_500, Operator::GtEq, 2), + // Negative literals round toward the correct side (floor/ceil, not + // truncation): m * 1000 < -1500 iff m < -1. + (Operator::Lt, -1_500, Operator::Lt, -1), + (Operator::LtEq, -1_500, Operator::LtEq, -2), + (Operator::Gt, -1_500, Operator::Gt, -2), + (Operator::GtEq, -1_500, Operator::GtEq, -1), + // Sub-millisecond equality can never hold for any file value; the + // rewrite keeps NULL semantics with an always-false / always-true + // comparison. + (Operator::Eq, 1_500, Operator::Lt, i64::MIN), + (Operator::NotEq, 1_500, Operator::GtEq, i64::MIN), + ]; + for (op, micros, expected_op, expected_millis) in cases { + let pred: Arc = + Arc::new(BinaryExpr::new(Arc::clone(&col), op, micros_lit(micros))); + let rewritten = adapter.rewrite(pred).unwrap(); + assert_millis_comparison(&rewritten, expected_op, expected_millis); + + // The literal-on-left form swaps the operator and rescales the same way. + let pred: Arc = Arc::new(BinaryExpr::new( + micros_lit(micros), + op.swap().unwrap(), + Arc::clone(&col), + )); + let rewritten = adapter.rewrite(pred).unwrap(); + assert_millis_comparison(&rewritten, expected_op, expected_millis); + } + } + + #[test] + fn test_millis_timestamp_null_checks_rewrite_to_raw_column() { + let adapter = millis_file_adapter(); + let col: Arc = Arc::new(Column::new("ts", 0)); + + let pred: Arc = Arc::new(IsNullExpr::new(Arc::clone(&col))); + let rewritten = adapter.rewrite(pred).unwrap(); + let is_null = rewritten + .downcast_ref::() + .expect("expected IsNullExpr"); + assert!(is_null.arg().downcast_ref::().is_some()); + + let pred: Arc = Arc::new(IsNotNullExpr::new(col)); + let rewritten = adapter.rewrite(pred).unwrap(); + let is_not_null = rewritten + .downcast_ref::() + .expect("expected IsNotNullExpr"); + assert!(is_not_null.arg().downcast_ref::().is_some()); + } } diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index fa431cfec1f..06dd9198dba 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -332,6 +332,48 @@ abstract class ParquetReadSuite extends CometTestBase { } } + test("TIMESTAMP_MILLIS overflow rows skipped by filter pruning do not fail") { + // Spark prunes the row group from the TIMESTAMP_MILLIS statistics before the + // vectorized reader ever calls the checked millisToMicros conversion, so the + // query returns no rows instead of failing. The native scan must preserve that + // behavior: values Spark never reads must not raise overflow errors. + Seq(false, true).foreach { dictionaryEnabled => + withTempDir { dir => + val path = new Path(dir.toURI.toString, "part-r-0.parquet") + val schema = MessageTypeParser.parseMessageType(""" + |message root { + | optional int64 ts(TIMESTAMP_MILLIS); + |} + |""".stripMargin) + val writer = createParquetWriter(schema, path, dictionaryEnabled) + (0 until 16).foreach { _ => + val record = new SimpleGroup(schema) + // Milliseconds that overflow Long when converted to microseconds + record.add(0, 9223372036854776L) + writer.write(record) + } + writer.close() + + Seq(false, true).foreach { ansiEnabled => + Seq(false, true).foreach { rowFilterPushdown => + withSQLConf( + SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString, + CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key -> + rowFilterPushdown.toString) { + readParquetFile(path.toString) { df => + val filtered = df.where("ts < timestamp'1970-01-01 00:00:00'") + assert(collect(filtered.queryExecution.executedPlan) { + case _: CometNativeScanExec => true + }.nonEmpty) + checkSparkAnswer(filtered) + } + } + } + } + } + } + } + test("timestamp as int96") { import testImplicits._ From 40b92b7d0cedbfb4307072539abc835e2a94a168 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 30 Aug 2026 02:43:01 +0800 Subject: [PATCH 13/13] fix: cover IN, null-safe equality, and nested predicates in the millis-domain rewrite Extend the millisecond-domain predicate rewrite to InListExpr and IsDistinctFrom/IsNotDistinctFrom, both of which DataFusion's pruning predicate can analyze, so row-group pruning protects flat TIMESTAMP_MILLIS columns for those forms too. Nested-field predicates can be neither pruned (PruningPredicate has no nested-field support) nor evaluated as Parquet row filters (struct columns are classified non-pushable), so no rewrite can keep the checked conversion from failing queries Spark answers via nested statistics pruning. Scans whose data filters reference nested fields fall back to the safe cast (overflow -> NULL, the pre-existing behavior), and the checked conversion is scoped to top-level columns for the same reason. Co-Authored-By: Claude Fable 5 --- native/core/src/parquet/parquet_exec.rs | 19 +- native/core/src/parquet/parquet_support.rs | 110 ++++++- native/core/src/parquet/schema_adapter.rs | 269 +++++++++++++++++- .../comet/parquet/ParquetReadSuite.scala | 21 +- 4 files changed, 402 insertions(+), 17 deletions(-) diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 1308ce97fca..7eb9d26412f 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -33,11 +33,18 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion::prelude::SessionContext; use datafusion::scalar::ScalarValue; -use datafusion_comet_spark_expr::EvalMode; +use datafusion_comet_spark_expr::{EvalMode, GetStructField}; use datafusion_datasource::TableSchema; use std::collections::HashMap; use std::sync::Arc; +/// True when the expression reads a struct field (directly or in any child), i.e. the +/// predicate references a nested column. +fn references_nested_field(expr: &Arc) -> bool { + expr.downcast_ref::().is_some() + || expr.children().iter().any(|c| references_nested_field(c)) +} + /// Initializes a DataSourceExec plan with a ParquetSource for Comet's native Parquet scan. /// /// `required_schema`: Schema to be projected by the scan. @@ -93,6 +100,16 @@ pub(crate) fn init_datasource_exec( ); spark_parquet_options.use_field_id = use_field_id; spark_parquet_options.ignore_missing_field_id = ignore_missing_field_id; + // Spark only avoids the TIMESTAMP_MILLIS overflow error for filtered-out values through + // row-group statistics pruning. DataFusion can neither prune nested-field predicates + // (`PruningPredicate` has no nested-field support) nor evaluate them as Parquet row + // filters (struct columns are classified non-pushable), so a scan carrying such a + // predicate would decode row groups Spark prunes and fail on the checked conversion of + // any overflowing top-level column. Fall back to the safe cast (overflow -> NULL) for + // these scans; the filter above the scan then discards the rows like Spark's pruning. + spark_parquet_options.checked_timestamp_overflow = !data_filters + .as_ref() + .is_some_and(|filters| filters.iter().any(references_nested_field)); // Determine the schema and projection to use for ParquetSource. // When data_schema is provided, use it as the base schema so DataFusion knows the full diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 9842d6220f5..49422d81a65 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -99,6 +99,16 @@ pub struct SparkParquetOptions { /// (Spark 3.x, SPARK-36182). Mirrors Comet's per-Spark-version constant /// in ShimCometConf. pub allow_timestamp_ltz_to_ntz: bool, + /// When true (the default), a top-level TIMESTAMP_MILLIS column that overflows during + /// the millis->micros upscale raises an error, matching Spark's checked + /// `millisToMicros`. Scans whose data filters reference nested struct fields set this + /// to false and fall back to the safe cast (overflow -> NULL): Spark only avoids the + /// overflow error for filtered-out values through row-group statistics pruning, and + /// DataFusion can neither prune (`PruningPredicate` has no nested-field support) nor + /// row-filter (struct columns are classified non-pushable in + /// `can_expr_be_pushed_down_with_schemas`) such predicates, so a checked conversion + /// would fail queries Spark answers with zero rows. + pub checked_timestamp_overflow: bool, } impl SparkParquetOptions { @@ -115,6 +125,7 @@ impl SparkParquetOptions { ignore_missing_field_id: false, allow_type_promotion: false, allow_timestamp_ltz_to_ntz: false, + checked_timestamp_overflow: true, } } @@ -131,6 +142,7 @@ impl SparkParquetOptions { ignore_missing_field_id: false, allow_type_promotion: false, allow_timestamp_ltz_to_ntz: false, + checked_timestamp_overflow: true, } } } @@ -167,6 +179,15 @@ fn parquet_convert_array( array: ArrayRef, to_type: &DataType, parquet_options: &SparkParquetOptions, +) -> DataFusionResult { + parquet_convert_array_impl(array, to_type, parquet_options, true) +} + +fn parquet_convert_array_impl( + array: ArrayRef, + to_type: &DataType, + parquet_options: &SparkParquetOptions, + top_level: bool, ) -> DataFusionResult { use DataType::*; let from_type = array.data_type(); @@ -182,10 +203,11 @@ fn parquet_convert_array( )?), (List(_), List(to_inner_type)) => { let list_arr: &ListArray = array.as_list(); - let cast_field = parquet_convert_array( + let cast_field = parquet_convert_array_impl( Arc::clone(list_arr.values()), to_inner_type.data_type(), parquet_options, + false, )?; Ok(Arc::new(ListArray::new( @@ -198,12 +220,19 @@ fn parquet_convert_array( ( Timestamp(TimeUnit::Millisecond, _), Timestamp(TimeUnit::Microsecond, target_tz), - ) => { + ) if top_level && parquet_options.checked_timestamp_overflow => { // Spark's Parquet reader calls the checked `millisToMicros` conversion for both // direct and dictionary values, independent of CAST evaluation mode: // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java#L817-L833 // `millisToMicros` uses `Math.multiplyExact`: // https://github.com/apache/spark/blob/v4.2.0/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkDateTimeUtils.scala#L103-L108 + // + // The checked conversion is limited to TOP-LEVEL columns. Spark only avoids the + // error for filtered-out values through row-group statistics pruning, and + // DataFusion's PruningPredicate does not support nested fields yet, so a checked + // conversion on a nested field would fail queries whose predicates Spark prunes + // (e.g. `WHERE s.ts < X` over an all-overflowing file). Nested fields keep the + // pre-existing safe-cast behavior below (overflow -> NULL). let micros = array .as_primitive::() .try_unary::<_, TimestampMicrosecondType, _>(|value| value.mul_checked(1_000))? @@ -315,10 +344,11 @@ fn parquet_convert_struct_to_struct( }; if let Some(from_index) = from_index { - cast_fields.push(parquet_convert_array( + cast_fields.push(parquet_convert_array_impl( Arc::clone(array.column(from_index)), to_field.data_type(), parquet_options, + false, )?); field_overlap = true; } else { @@ -366,15 +396,17 @@ fn parquet_convert_map_to_map( "map is missing value field".to_string(), ))?; - let key_array = parquet_convert_array( + let key_array = parquet_convert_array_impl( Arc::clone(from.keys()), key_field.data_type(), parquet_options, + false, )?; - let value_array = parquet_convert_array( + let value_array = parquet_convert_array_impl( Arc::clone(from.values()), value_field.data_type(), parquet_options, + false, )?; Ok(Arc::new(MapArray::new( @@ -671,4 +703,72 @@ mod tests { } } } + + #[test] + fn test_millis_to_micros_overflow_checked_only_at_top_level() { + use crate::parquet::parquet_support::{parquet_convert_array, SparkParquetOptions}; + use arrow::array::{Array, ArrayRef, StructArray, TimestampMillisecondArray}; + use arrow::datatypes::{DataType, Field, Fields, TimeUnit}; + use datafusion_comet_spark_expr::EvalMode; + use std::sync::Arc; + + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let overflow_millis = 9_223_372_036_854_776_i64; + let millis: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![ + Some(overflow_millis), + None, + ])); + let micros_type = DataType::Timestamp(TimeUnit::Microsecond, None); + + // Top-level: checked, matching Spark's `millisToMicros` (`Math.multiplyExact`). + let err = parquet_convert_array(Arc::clone(&millis), µs_type, &options) + .expect_err("top-level overflow must error"); + assert!( + err.to_string().to_lowercase().contains("overflow"), + "unexpected error: {err}" + ); + + // Scans whose data filters reference nested fields disable the checked + // conversion (DataFusion cannot prune or row-filter those predicates the way + // Spark's statistics pruning protects them), falling back to overflow -> NULL. + let mut unchecked_options = options.clone(); + unchecked_options.checked_timestamp_overflow = false; + let converted = + parquet_convert_array(Arc::clone(&millis), µs_type, &unchecked_options) + .expect("unchecked overflow must not error"); + assert!(converted.is_null(0), "overflow must become NULL"); + assert!(converted.is_null(1)); + + // Nested: DataFusion's PruningPredicate cannot prune nested fields, so a + // checked conversion would fail queries whose predicates Spark satisfies via + // row-group statistics pruning. The nested field keeps the safe-cast behavior: + // overflow becomes NULL. + let child_field = Arc::new(Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + true, + )); + let strukt: ArrayRef = Arc::new(StructArray::new( + Fields::from(vec![Arc::clone(&child_field)]), + vec![millis], + None, + )); + let target = DataType::Struct(Fields::from(vec![Arc::new(Field::new( + "ts", + micros_type.clone(), + true, + ))])); + let converted = parquet_convert_array(strukt, &target, &options) + .expect("nested overflow must not error"); + let converted_child = Arc::clone( + converted + .as_any() + .downcast_ref::() + .unwrap() + .column(0), + ); + assert_eq!(converted_child.data_type(), µs_type); + assert!(converted_child.is_null(0), "overflow must become NULL"); + assert!(converted_child.is_null(1)); + } } diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index c956184c400..44512ca2241 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -24,7 +24,7 @@ use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::logical_expr::Operator; use datafusion::physical_expr::expressions::{ - BinaryExpr, Column, IsNotNullExpr, IsNullExpr, Literal, + in_list, BinaryExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, Literal, }; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ColumnarValue; @@ -553,17 +553,18 @@ fn as_millis_to_micros_cast(expr: &Arc) -> Option) -> Option { +/// If `expr` is a microsecond timestamp literal (null or not), return its value. +fn as_micros_literal(expr: &Arc) -> Option> { match expr.downcast_ref::()?.value() { - ScalarValue::TimestampMicrosecond(Some(micros), _) => Some(*micros), + ScalarValue::TimestampMicrosecond(micros, _) => Some(*micros), _ => None, } } impl SparkPhysicalExprAdapter { /// Rewrite predicate expressions over a `TIMESTAMP_MILLIS` file column read as - /// microseconds into the millisecond domain, mirroring Spark's `ParquetFilters`, + /// microseconds into the millisecond domain — comparisons (including null-safe + /// `<=>`), `IN` lists, and null checks — mirroring Spark's `ParquetFilters`, /// which pushes timestamp predicates down in the file's physical unit: /// https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilters.scala#L192-L196 /// @@ -593,6 +594,16 @@ impl SparkPhysicalExprAdapter { } } + // `ts IN (...)` / `ts NOT IN (...)`: rescale every microsecond literal in the + // list. DataFusion's pruning predicate understands `InListExpr` (up to 20 + // elements), so keeping the raw column visible restores row-group pruning, + // mirroring Spark's `ParquetFilters` handling of `In`. + if let Some(in_list_expr) = expr.downcast_ref::() { + if let Some((inner, file_tz)) = as_millis_to_micros_cast(in_list_expr.expr()) { + return self.rewrite_millis_in_list(&expr, in_list_expr, inner, file_tz); + } + } + let Some(binary) = expr.downcast_ref::() else { return Ok(Transformed::no(expr)); }; @@ -613,6 +624,31 @@ impl SparkPhysicalExprAdapter { return Ok(Transformed::no(expr)); }; + let Some(micros) = micros else { + // NULL literal. The ordinary comparisons return NULL for every row on both + // sides of the rewrite, so rescaling the literal to a NULL millisecond + // literal is exact. The null-safe comparisons test null-ness of the value, + // which the checked conversion preserves, so probe the raw column. + let rewritten: Arc = match op { + Operator::IsNotDistinctFrom => Arc::new(IsNullExpr::new(inner)), + Operator::IsDistinctFrom => Arc::new(IsNotNullExpr::new(inner)), + Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + | Operator::Eq + | Operator::NotEq => Arc::new(BinaryExpr::new( + inner, + op, + Arc::new(Literal::new(ScalarValue::TimestampMillisecond( + None, file_tz, + ))), + )), + _ => return Ok(Transformed::no(expr)), + }; + return Ok(Transformed::yes(rewritten)); + }; + // For a file value `m` (milliseconds) the logical value is exactly `m * 1000` // microseconds, so `m * 1000 OP L` rewrites to an exact comparison on `m`. let floor = micros.div_euclid(1_000); @@ -625,12 +661,27 @@ impl SparkPhysicalExprAdapter { Operator::GtEq => (Operator::GtEq, ceil), Operator::Eq if exact => (Operator::Eq, floor), Operator::NotEq if exact => (Operator::NotEq, floor), + Operator::IsNotDistinctFrom if exact => (Operator::IsNotDistinctFrom, floor), + Operator::IsDistinctFrom if exact => (Operator::IsDistinctFrom, floor), // A sub-millisecond literal can never equal `m * 1000`. `col < i64::MIN` // (resp. `col >= i64::MIN`) is false (resp. true) for every non-null value // and NULL for nulls, matching `=` / `!=` null semantics while remaining a // plain, prunable comparison. Operator::Eq => (Operator::Lt, i64::MIN), Operator::NotEq => (Operator::GtEq, i64::MIN), + // The null-safe comparisons never return NULL, so an impossible literal + // folds to a constant. DataFusion's pruning predicate evaluates boolean + // literals, so `false` still prunes every row group. + Operator::IsNotDistinctFrom => { + return Ok(Transformed::yes(Arc::new(Literal::new( + ScalarValue::Boolean(Some(false)), + )))); + } + Operator::IsDistinctFrom => { + return Ok(Transformed::yes(Arc::new(Literal::new( + ScalarValue::Boolean(Some(true)), + )))); + } _ => return Ok(Transformed::no(expr)), }; let literal = Arc::new(Literal::new(ScalarValue::TimestampMillisecond( @@ -642,6 +693,61 @@ impl SparkPhysicalExprAdapter { )))) } + /// Rewrite `ts IN (...)` / `ts NOT IN (...)` over a millis->micros conversion into + /// the millisecond domain. Divisible literals rescale exactly; sub-millisecond + /// literals can never equal `m * 1000` and drop out of the list, which preserves + /// IN's three-valued logic (a dropped element contributes `false` to the OR / `true` + /// to the AND for every non-null probe, and the null-probe result stays NULL either + /// way). A list emptied this way degenerates to the same always-false / + /// always-true-with-null-semantics sentinels the `=` / `!=` rewrite uses. + fn rewrite_millis_in_list( + &self, + original: &Arc, + in_list_expr: &InListExpr, + inner: Arc, + file_tz: Option>, + ) -> DataFusionResult>> { + let mut millis_list: Vec> = Vec::new(); + for item in in_list_expr.list() { + // Any non-literal or non-timestamp element: leave the expression alone. + let Some(micros) = as_micros_literal(item) else { + return Ok(Transformed::no(Arc::clone(original))); + }; + match micros { + None => millis_list.push(Arc::new(Literal::new( + ScalarValue::TimestampMillisecond(None, file_tz.clone()), + ))), + Some(v) if v % 1_000 == 0 => millis_list.push(Arc::new(Literal::new( + ScalarValue::TimestampMillisecond(Some(v / 1_000), file_tz.clone()), + ))), + Some(_) => {} + } + } + let negated = in_list_expr.negated(); + if millis_list.is_empty() { + let (op, sentinel) = if negated { + (Operator::GtEq, i64::MIN) + } else { + (Operator::Lt, i64::MIN) + }; + return Ok(Transformed::yes(Arc::new(BinaryExpr::new( + inner, + op, + Arc::new(Literal::new(ScalarValue::TimestampMillisecond( + Some(sentinel), + file_tz, + ))), + )))); + } + let rewritten = in_list( + inner, + millis_list, + &negated, + self.physical_file_schema.as_ref(), + )?; + Ok(Transformed::yes(rewritten)) + } + /// Wrap ALL Column expressions that have type mismatches with CometCastColumnExpr. /// This is the fallback path when the default adapter fails (e.g., for complex /// nested type casts like List or Map). Uses `spark_parquet_convert` @@ -1898,7 +2004,7 @@ mod test { use arrow::datatypes::TimeUnit; use datafusion::logical_expr::Operator; use datafusion::physical_expr::expressions::{ - BinaryExpr, Column, IsNotNullExpr, IsNullExpr, Literal, + in_list, BinaryExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, Literal, }; use datafusion::physical_expr::PhysicalExpr; use datafusion::scalar::ScalarValue; @@ -2021,4 +2127,155 @@ mod test { .expect("expected IsNotNullExpr"); assert!(is_not_null.arg().downcast_ref::().is_some()); } + + #[test] + fn test_millis_timestamp_null_safe_equality_rewrites() { + let adapter = millis_file_adapter(); + let col: Arc = Arc::new(Column::new("ts", 0)); + + // Divisible literal: same operator, rescaled. + for (op, expected_op) in [ + (Operator::IsNotDistinctFrom, Operator::IsNotDistinctFrom), + (Operator::IsDistinctFrom, Operator::IsDistinctFrom), + ] { + let pred: Arc = + Arc::new(BinaryExpr::new(Arc::clone(&col), op, micros_lit(2_000))); + let rewritten = adapter.rewrite(pred).unwrap(); + assert_millis_comparison(&rewritten, expected_op, 2); + } + + // Sub-millisecond literal: null-safe comparisons never return NULL, so the + // expression folds to a boolean constant. + for (op, expected) in [ + (Operator::IsNotDistinctFrom, false), + (Operator::IsDistinctFrom, true), + ] { + let pred: Arc = + Arc::new(BinaryExpr::new(Arc::clone(&col), op, micros_lit(1_500))); + let rewritten = adapter.rewrite(pred).unwrap(); + let literal = rewritten + .downcast_ref::() + .expect("expected boolean Literal"); + assert_eq!(literal.value(), &ScalarValue::Boolean(Some(expected))); + } + + // NULL literal: `<=> NULL` tests null-ness of the raw value. + let null_lit: Arc = Arc::new(Literal::new( + ScalarValue::TimestampMicrosecond(None, Some(Arc::from("UTC"))), + )); + let pred: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&col), + Operator::IsNotDistinctFrom, + Arc::clone(&null_lit), + )); + let rewritten = adapter.rewrite(pred).unwrap(); + assert!(rewritten.downcast_ref::().is_some()); + + let pred: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&col), + Operator::IsDistinctFrom, + Arc::clone(&null_lit), + )); + let rewritten = adapter.rewrite(pred).unwrap(); + assert!(rewritten.downcast_ref::().is_some()); + + // NULL literal with an ordinary comparison: NULL result either way, so the + // literal rescales to a NULL millisecond literal. + let pred: Arc = + Arc::new(BinaryExpr::new(Arc::clone(&col), Operator::Lt, null_lit)); + let rewritten = adapter.rewrite(pred).unwrap(); + let binary = rewritten + .downcast_ref::() + .expect("expected BinaryExpr"); + let literal = binary + .right() + .downcast_ref::() + .expect("expected Literal"); + assert_eq!( + literal.value(), + &ScalarValue::TimestampMillisecond(None, Some(Arc::from("UTC"))) + ); + } + + fn millis_values_of(list: &[Arc]) -> Vec> { + list.iter() + .map(|item| { + match item + .downcast_ref::() + .expect("expected Literal list element") + .value() + { + ScalarValue::TimestampMillisecond(v, _) => *v, + other => panic!("expected millisecond literal, got {other:?}"), + } + }) + .collect() + } + + #[test] + fn test_millis_timestamp_in_list_rewrites_to_millis_domain() { + let adapter = millis_file_adapter(); + let physical = Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("UTC"))), + true, + )]); + let col: Arc = Arc::new(Column::new("ts", 0)); + + for negated in [false, true] { + // Divisible literals rescale, a sub-millisecond literal drops out, and a + // NULL literal stays NULL. + let null_lit: Arc = Arc::new(Literal::new( + ScalarValue::TimestampMicrosecond(None, Some(Arc::from("UTC"))), + )); + let logical = Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))), + true, + )]); + let pred = in_list( + Arc::clone(&col), + vec![micros_lit(2_000), micros_lit(1_500), null_lit], + &negated, + &logical, + ) + .unwrap(); + let rewritten = adapter.rewrite(pred).unwrap(); + let rewritten_in_list = rewritten + .downcast_ref::() + .expect("expected InListExpr"); + assert!(rewritten_in_list.expr().downcast_ref::().is_some()); + assert_eq!(rewritten_in_list.negated(), negated); + assert_eq!( + millis_values_of(rewritten_in_list.list()), + vec![Some(2), None] + ); + assert_eq!( + rewritten_in_list.expr().data_type(&physical).unwrap(), + DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("UTC"))) + ); + + // A list of only impossible literals degenerates to the always-false / + // always-true sentinel comparison with IN's null semantics. + let logical = Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))), + true, + )]); + let pred = in_list( + Arc::clone(&col), + vec![micros_lit(1_500), micros_lit(-1)], + &negated, + &logical, + ) + .unwrap(); + let rewritten = adapter.rewrite(pred).unwrap(); + let expected_op = if negated { + Operator::GtEq + } else { + Operator::Lt + }; + assert_millis_comparison(&rewritten, expected_op, i64::MIN); + } + } } diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 06dd9198dba..075da3c520b 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -343,6 +343,9 @@ abstract class ParquetReadSuite extends CometTestBase { val schema = MessageTypeParser.parseMessageType(""" |message root { | optional int64 ts(TIMESTAMP_MILLIS); + | optional group s { + | optional int64 ts(TIMESTAMP_MILLIS); + | } |} |""".stripMargin) val writer = createParquetWriter(schema, path, dictionaryEnabled) @@ -350,10 +353,16 @@ abstract class ParquetReadSuite extends CometTestBase { val record = new SimpleGroup(schema) // Milliseconds that overflow Long when converted to microseconds record.add(0, 9223372036854776L) + record.addGroup(1).add(0, 9223372036854776L) writer.write(record) } writer.close() + val predicates = Seq( + "ts < timestamp'1970-01-01 00:00:00'", + "ts IN (timestamp'1970-01-01 00:00:00', timestamp'1970-01-02 00:00:00')", + "ts <=> timestamp'1970-01-01 00:00:00'", + "s.ts < timestamp'1970-01-01 00:00:00'") Seq(false, true).foreach { ansiEnabled => Seq(false, true).foreach { rowFilterPushdown => withSQLConf( @@ -361,11 +370,13 @@ abstract class ParquetReadSuite extends CometTestBase { CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key -> rowFilterPushdown.toString) { readParquetFile(path.toString) { df => - val filtered = df.where("ts < timestamp'1970-01-01 00:00:00'") - assert(collect(filtered.queryExecution.executedPlan) { - case _: CometNativeScanExec => true - }.nonEmpty) - checkSparkAnswer(filtered) + predicates.foreach { predicate => + val filtered = df.where(predicate) + assert(collect(filtered.queryExecution.executedPlan) { + case _: CometNativeScanExec => true + }.nonEmpty) + checkSparkAnswer(filtered) + } } } }