diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 1cc928d1d59..fc5799f3166 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -15,10 +15,7 @@ // specific language governing permissions and limitations // under the License. use arrow::{ - array::{ - make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray, - TimestampMicrosecondArray, TimestampMillisecondArray, - }, + array::{make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray}, compute::CastOptions, datatypes::{DataType, FieldRef, Schema, TimeUnit}, record_batch::RecordBatch, @@ -26,8 +23,7 @@ use arrow::{ 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::{ @@ -142,40 +138,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. @@ -214,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. 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), + ( + 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. @@ -271,23 +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 casted = cast_timestamp_micros_to_millis_array(&array, target_tz.clone()); - Ok(ColumnarValue::Array(casted)) - } - ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(opt_val, _)) => { - let casted = cast_timestamp_micros_to_millis_scalar(opt_val, target_tz.clone()); - Ok(ColumnarValue::Scalar(casted)) - } - _ => 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. @@ -329,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()); } @@ -349,159 +316,84 @@ 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_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 - let input_field = Arc::new(Field::new( - "ts", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - )); - 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, - )); - - // 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_000_000), Some(2_000_000), 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), 1000); - assert_eq!(millis_array.value(1), 2000); - assert!(millis_array.is_null(2)); - } - _ => panic!("Expected Array 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") + )); } } #[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)]); - - // Create target field with TimestampMillisecond - let target_field = Arc::new(Field::new( - "ts", - DataType::Timestamp(TimeUnit::Millisecond, None), - true, - )); - - // Create a literal expression that returns a scalar - let scalar = ScalarValue::TimestampMicrosecond(Some(1_500_000), 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(1500), None)); + 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()); } - _ => panic!("Expected Scalar result"), } } 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/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 96346962bc4..ed057d28ce3 100644 --- a/native/spark-expr/src/conversion_funcs/temporal.rs +++ b/native/spark-expr/src/conversion_funcs/temporal.rs @@ -18,8 +18,10 @@ 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; +use arrow::datatypes::{DataType, Date32Type, TimeUnit}; use chrono::NaiveDate; +use datafusion::common::format::DEFAULT_CAST_OPTIONS; use std::str::FromStr; use std::sync::Arc; @@ -39,49 +41,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), + &DEFAULT_CAST_OPTIONS, + )?); + } + 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..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,13 +15,12 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, Date32Array, Int32Array}; +use arrow::compute::cast_with_options; use arrow::datatypes::DataType; -use datafusion::common::{utils::take_function_args, DataFusionError, Result, ScalarValue}; +use datafusion::common::{format::DEFAULT_CAST_OPTIONS, 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, + &DEFAULT_CAST_OPTIONS, + )?)), + 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..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::array::TimestampMicrosecondArray; use arrow::datatypes::{MAX_DECIMAL128_FOR_EACH_PRECISION, MIN_DECIMAL128_FOR_EACH_PRECISION}; use arrow::error::ArrowError; use arrow::{ @@ -81,13 +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(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)) - } _ => { // Not supported Err(ArrowError::CastError(format!( @@ -376,6 +368,7 @@ pub fn unlikely(b: bool) -> bool { #[cfg(test)] mod tests { use super::*; + 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(); 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 684c4d6581f..fdff3f1a820 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._