rebase_main

Author: coderfender

native/spark-expr/src/conversion_funcs/boolean.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::SparkResult;
-use arrow::array::{ArrayRef, AsArray, Decimal128Array};
+use arrow::array::{Array, ArrayRef, AsArray, Decimal128Array, TimestampMicrosecondBuilder};
 use arrow::datatypes::DataType;
 use std::sync::Arc;
 
@@ -28,7 +28,6 @@ pub fn is_df_cast_from_bool_spark_compatible(to_type: &DataType) -> bool {
     )
 }
 
-// only DF incompatible boolean cast
 pub fn cast_boolean_to_decimal(
     array: &ArrayRef,
     precision: u8,
@@ -43,6 +42,25 @@ pub fn cast_boolean_to_decimal(
     Ok(Arc::new(result.with_precision_and_scale(precision, scale)?))
 }
 
+pub(crate) fn cast_boolean_to_timestamp(
+    array_ref: &ArrayRef,
+    target_tz: &Option<Arc<str>>,
+) -> SparkResult<ArrayRef> {
+    let bool_array = array_ref.as_boolean();
+    let mut builder = TimestampMicrosecondBuilder::with_capacity(bool_array.len());
+
+    for i in 0..bool_array.len() {
+        if bool_array.is_null(i) {
+            builder.append_null();
+        } else {
+            let micros = if bool_array.value(i) { 1 } else { 0 };
+            builder.append_value(micros);
+        }
+    }
+
+    Ok(Arc::new(builder.finish().with_timezone_opt(target_tz.clone())) as ArrayRef)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -53,6 +71,7 @@ mod tests {
         Int64Array, Int8Array, StringArray,
     };
     use arrow::datatypes::DataType::Decimal128;
+    use arrow::datatypes::TimestampMicrosecondType;
     use std::sync::Arc;
 
     fn test_input_bool_array() -> ArrayRef {
@@ -193,4 +212,26 @@ mod tests {
         assert_eq!(arr.value(1), expected_arr.value(1));
         assert!(arr.is_null(2));
     }
+
+    #[test]
+    fn test_cast_boolean_to_timestamp() {
+        let timezones: [Option<Arc<str>>; 3] = [
+            Some(Arc::from("UTC")),
+            Some(Arc::from("America/Los_Angeles")),
+            None,
+        ];
+
+        for tz in &timezones {
+            let bool_array: ArrayRef =
+                Arc::new(BooleanArray::from(vec![Some(true), Some(false), None]));
+
+            let result = cast_boolean_to_timestamp(&bool_array, tz).unwrap();
+            let ts_array = result.as_primitive::<TimestampMicrosecondType>();
+
+            assert_eq!(ts_array.value(0), 1); // true -> 1 microsecond
+            assert_eq!(ts_array.value(1), 0); // false -> 0 (epoch)
+            assert!(ts_array.is_null(2));
+            assert_eq!(ts_array.timezone(), tz.as_ref().map(|s| s.as_ref()));
+        }
+    }
 }

native/spark-expr/src/conversion_funcs/cast.rs

@@ -16,14 +16,15 @@
 // under the License.
 
 use crate::conversion_funcs::boolean::{
-    cast_boolean_to_decimal, is_df_cast_from_bool_spark_compatible,
+    cast_boolean_to_decimal, cast_boolean_to_timestamp, is_df_cast_from_bool_spark_compatible,
 };
 use crate::conversion_funcs::numeric::{
-    cast_float32_to_decimal128, cast_float64_to_decimal128, cast_int_to_decimal128,
-    cast_int_to_timestamp, is_df_cast_from_decimal_spark_compatible,
-    is_df_cast_from_float_spark_compatible, is_df_cast_from_int_spark_compatible,
-    spark_cast_decimal_to_boolean, spark_cast_float32_to_utf8, spark_cast_float64_to_utf8,
-    spark_cast_int_to_int, spark_cast_nonintegral_numeric_to_integral,
+    cast_decimal_to_timestamp, cast_float32_to_decimal128, cast_float64_to_decimal128,
+    cast_float_to_timestamp, cast_int_to_decimal128, cast_int_to_timestamp,
+    is_df_cast_from_decimal_spark_compatible, is_df_cast_from_float_spark_compatible,
+    is_df_cast_from_int_spark_compatible, spark_cast_decimal_to_boolean,
+    spark_cast_float32_to_utf8, spark_cast_float64_to_utf8, spark_cast_int_to_int,
+    spark_cast_nonintegral_numeric_to_integral,
 };
 use crate::conversion_funcs::string::{
     cast_string_to_date, cast_string_to_decimal, cast_string_to_float, cast_string_to_int,
@@ -778,7 +779,7 @@ fn cast_binary_formatter(value: &[u8]) -> String {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::{BooleanArray, StringArray};
+    use arrow::array::StringArray;
     use arrow::datatypes::TimestampMicrosecondType;
     use arrow::datatypes::{Field, Fields};
     #[test]

native/spark-expr/src/conversion_funcs/numeric.rs

@@ -24,7 +24,7 @@ use arrow::array::{
     OffsetSizeTrait, PrimitiveArray, TimestampMicrosecondBuilder,
 };
 use arrow::datatypes::{
-    is_validate_decimal_precision, ArrowPrimitiveType, DataType, Decimal128Type, Float32Type,
+    i256, is_validate_decimal_precision, ArrowPrimitiveType, DataType, Decimal128Type, Float32Type,
     Float64Type, Int16Type, Int32Type, Int64Type, Int8Type,
 };
 use num::{cast::AsPrimitive, ToPrimitive, Zero};
@@ -75,6 +75,56 @@ pub(crate) fn is_df_cast_from_decimal_spark_compatible(to_type: &DataType) -> bo
     )
 }
 
+macro_rules! cast_float_to_timestamp_impl {
+    ($array:expr, $builder:expr, $primitive_type:ty, $eval_mode:expr) => {{
+        let arr = $array.as_primitive::<$primitive_type>();
+        for i in 0..arr.len() {
+            if arr.is_null(i) {
+                $builder.append_null();
+            } else {
+                let val = arr.value(i) as f64;
+                // Path 1: NaN/Infinity check - error says TIMESTAMP
+                if val.is_nan() || val.is_infinite() {
+                    if $eval_mode == EvalMode::Ansi {
+                        return Err(SparkError::CastInvalidValue {
+                            value: val.to_string(),
+                            from_type: "DOUBLE".to_string(),
+                            to_type: "TIMESTAMP".to_string(),
+                        });
+                    }
+                    $builder.append_null();
+                } else {
+                    // Path 2: Multiply then check overflow - error says BIGINT
+                    let micros = val * MICROS_PER_SECOND as f64;
+                    if micros.floor() <= i64::MAX as f64 && micros.ceil() >= i64::MIN as f64 {
+                        $builder.append_value(micros as i64);
+                    } else {
+                        if $eval_mode == EvalMode::Ansi {
+                            let value_str = if micros.is_infinite() {
+                                if micros.is_sign_positive() {
+                                    "Infinity".to_string()
+                                } else {
+                                    "-Infinity".to_string()
+                                }
+                            } else if micros.is_nan() {
+                                "NaN".to_string()
+                            } else {
+                                format!("{:e}", micros).to_uppercase() + "D"
+                            };
+                            return Err(SparkError::CastOverFlow {
+                                value: value_str,
+                                from_type: "DOUBLE".to_string(),
+                                to_type: "BIGINT".to_string(),
+                            });
+                        }
+                        $builder.append_null();
+                    }
+                }
+            }
+        }
+    }};
+}
+
 macro_rules! cast_float_to_string {
     ($from:expr, $eval_mode:expr, $type:ty, $output_type:ty, $offset_type:ty) => {{
 
@@ -913,6 +963,56 @@ pub(crate) fn cast_int_to_timestamp(
     Ok(Arc::new(builder.finish().with_timezone_opt(target_tz.clone())) as ArrayRef)
 }
 
+pub(crate) fn cast_decimal_to_timestamp(
+    array_ref: &ArrayRef,
+    target_tz: &Option<Arc<str>>,
+    scale: i8,
+) -> SparkResult<ArrayRef> {
+    let arr = array_ref.as_primitive::<Decimal128Type>();
+    let scale_factor = 10_i128.pow(scale as u32);
+    let mut builder = TimestampMicrosecondBuilder::with_capacity(arr.len());
+
+    for i in 0..arr.len() {
+        if arr.is_null(i) {
+            builder.append_null();
+        } else {
+            let value = arr.value(i);
+            // Note: spark's big decimal 
+            let value_256 = i256::from_i128(value);
+            let micros_256 = value_256 * i256::from_i128(MICROS_PER_SECOND as i128);
+            let ts = micros_256 / i256::from_i128(scale_factor);
+            builder.append_value(ts.as_i128() as i64);
+        }
+    }
+
+    Ok(Arc::new(builder.finish().with_timezone_opt(target_tz.clone())) as ArrayRef)
+}
+
+pub(crate) fn cast_float_to_timestamp(
+    array_ref: &ArrayRef,
+    target_tz: &Option<Arc<str>>,
+    eval_mode: EvalMode,
+) -> SparkResult<ArrayRef> {
+    let mut builder = TimestampMicrosecondBuilder::with_capacity(array_ref.len());
+
+    match array_ref.data_type() {
+        DataType::Float32 => {
+            cast_float_to_timestamp_impl!(array_ref, builder, Float32Type, eval_mode)
+        }
+        DataType::Float64 => {
+            cast_float_to_timestamp_impl!(array_ref, builder, Float64Type, eval_mode)
+        }
+        dt => {
+            return Err(SparkError::Internal(format!(
+                "Unsupported type for cast_float_to_timestamp: {:?}",
+                dt
+            )))
+        }
+    }
+
+    Ok(Arc::new(builder.finish().with_timezone_opt(target_tz.clone())) as ArrayRef)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1100,4 +1200,113 @@ mod tests {
         assert!(casted.is_null(8));
         assert!(casted.is_null(9));
     }
+
+    #[test]
+    fn test_cast_decimal_to_timestamp() {
+        let timezones: [Option<Arc<str>>; 3] = [
+            Some(Arc::from("UTC")),
+            Some(Arc::from("America/Los_Angeles")),
+            None,
+        ];
+
+        for tz in &timezones {
+            // Decimal128 with scale 6
+            let decimal_array: ArrayRef = Arc::new(
+                Decimal128Array::from(vec![
+                    Some(0_i128),
+                    Some(1_000_000_i128),
+                    Some(-1_000_000_i128),
+                    Some(1_500_000_i128),
+                    Some(123_456_789_i128),
+                    None,
+                ])
+                .with_precision_and_scale(18, 6)
+                .unwrap(),
+            );
+
+            let result = cast_decimal_to_timestamp(&decimal_array, tz, 6).unwrap();
+            let ts_array = result.as_primitive::<TimestampMicrosecondType>();
+
+            assert_eq!(ts_array.value(0), 0);
+            assert_eq!(ts_array.value(1), 1_000_000);
+            assert_eq!(ts_array.value(2), -1_000_000);
+            assert_eq!(ts_array.value(3), 1_500_000);
+            assert_eq!(ts_array.value(4), 123_456_789);
+            assert!(ts_array.is_null(5));
+            assert_eq!(ts_array.timezone(), tz.as_ref().map(|s| s.as_ref()));
+
+            // Test with scale 2
+            let decimal_array: ArrayRef = Arc::new(
+                Decimal128Array::from(vec![Some(100_i128), Some(150_i128), Some(-250_i128)])
+                    .with_precision_and_scale(10, 2)
+                    .unwrap(),
+            );
+
+            let result = cast_decimal_to_timestamp(&decimal_array, tz, 2).unwrap();
+            let ts_array = result.as_primitive::<TimestampMicrosecondType>();
+
+            assert_eq!(ts_array.value(0), 1_000_000);
+            assert_eq!(ts_array.value(1), 1_500_000);
+            assert_eq!(ts_array.value(2), -2_500_000);
+        }
+    }
+
+    #[test]
+    fn test_cast_float_to_timestamp() {
+        let timezones: [Option<Arc<str>>; 3] = [
+            Some(Arc::from("UTC")),
+            Some(Arc::from("America/Los_Angeles")),
+            None,
+        ];
+        let eval_modes = [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try];
+
+        for tz in &timezones {
+            for eval_mode in &eval_modes {
+                // Float64 tests
+                let f64_array: ArrayRef = Arc::new(Float64Array::from(vec![
+                    Some(0.0),
+                    Some(1.0),
+                    Some(-1.0),
+                    Some(1.5),
+                    Some(0.000001),
+                    None,
+                ]));
+
+                let result = cast_float_to_timestamp(&f64_array, tz, *eval_mode).unwrap();
+                let ts_array = result.as_primitive::<TimestampMicrosecondType>();
+
+                assert_eq!(ts_array.value(0), 0);
+                assert_eq!(ts_array.value(1), 1_000_000);
+                assert_eq!(ts_array.value(2), -1_000_000);
+                assert_eq!(ts_array.value(3), 1_500_000);
+                assert_eq!(ts_array.value(4), 1);
+                assert!(ts_array.is_null(5));
+                assert_eq!(ts_array.timezone(), tz.as_ref().map(|s| s.as_ref()));
+
+                // Float32 tests
+                let f32_array: ArrayRef = Arc::new(Float32Array::from(vec![
+                    Some(0.0_f32),
+                    Some(1.0_f32),
+                    Some(-1.0_f32),
+                    None,
+                ]));
+
+                let result = cast_float_to_timestamp(&f32_array, tz, *eval_mode).unwrap();
+                let ts_array = result.as_primitive::<TimestampMicrosecondType>();
+
+                assert_eq!(ts_array.value(0), 0);
+                assert_eq!(ts_array.value(1), 1_000_000);
+                assert_eq!(ts_array.value(2), -1_000_000);
+                assert!(ts_array.is_null(3));
+            }
+        }
+
+        // ANSI mode errors on NaN/Infinity
+        let tz = &Some(Arc::from("UTC"));
+        let f64_nan: ArrayRef = Arc::new(Float64Array::from(vec![Some(f64::NAN)]));
+        assert!(cast_float_to_timestamp(&f64_nan, tz, EvalMode::Ansi).is_err());
+
+        let f64_inf: ArrayRef = Arc::new(Float64Array::from(vec![Some(f64::INFINITY)]));
+        assert!(cast_float_to_timestamp(&f64_inf, tz, EvalMode::Ansi).is_err());
+    }
 }

spark/src/test/scala/org/apache/comet/CometCastSuite.scala

@@ -639,6 +639,10 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     castTest(generateDecimalsPrecision10Scale2(), DataTypes.TimestampType)
   }
 
+  test("cast DecimalType(38,10) to TimestampType") {
+    castTest(generateDecimalsPrecision38Scale18(), DataTypes.TimestampType)
+  }
+
   // CAST from StringType
 
   test("cast StringType to BooleanType") {