diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 1cc928d1d5..af3a2cc6aa 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. 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()); } @@ -354,154 +321,27 @@ mod tests { 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 - 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"), - } - } - - #[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)); - } - _ => 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 c6586b4681..ccd13994b5 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 96346962bc..ed057d28ce 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 0e624e6472..c3c53fefca 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 7a785c7225..f4d411b172 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; 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. @@ -81,12 +82,14 @@ 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)) => { + // 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 @@ -376,6 +379,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 +394,28 @@ 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);