Skip to content

Commit 0c4d1d4

Browse files
coderfendercoderfender
committed
int_timestamp_add_benchmarks
1 parent fd34134 commit 0c4d1d4

1 file changed

Lines changed: 86 additions & 96 deletions

File tree

  • native/spark-expr/src/conversion_funcs

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

Lines changed: 86 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -929,9 +929,7 @@ fn cast_array(
929929
(Boolean, Decimal128(precision, scale)) => {
930930
cast_boolean_to_decimal(&array, *precision, *scale)
931931
}
932-
(Int8 | Int16 | Int32 | Int64, Timestamp(_, _)) => {
933-
cast_int_to_timestamp(&array, cast_options)
934-
}
932+
(Int8 | Int16 | Int32 | Int64, Timestamp(_, tz)) => cast_int_to_timestamp(&array, tz),
935933
_ if cast_options.is_adapting_schema
936934
|| is_datafusion_spark_compatible(from_type, to_type) =>
937935
{
@@ -952,7 +950,7 @@ fn cast_array(
952950

953951
fn cast_int_to_timestamp(
954952
array_ref: &ArrayRef,
955-
cast_options: &SparkCastOptions,
953+
target_tz: &Option<Arc<str>>,
956954
) -> SparkResult<ArrayRef> {
957955
// Input is seconds since epoch, multiply by MICROS_PER_SECOND to get microseconds.
958956
let mut builder = TimestampMicrosecondBuilder::with_capacity(array_ref.len());
@@ -970,9 +968,7 @@ fn cast_int_to_timestamp(
970968
}
971969
}
972970

973-
// input tz is always defined or set to UTC on spark side
974-
let tz: Arc<str> = Arc::from(cast_options.timezone.as_str());
975-
Ok(Arc::new(builder.finish().with_timezone(tz)) as ArrayRef)
971+
Ok(Arc::new(builder.finish().with_timezone_opt(target_tz.clone())) as ArrayRef)
976972
}
977973

978974
fn cast_date_to_timestamp(
@@ -3564,97 +3560,91 @@ mod tests {
35643560

35653561
#[test]
35663562
fn test_cast_int_to_timestamp() {
3567-
let timezones = [
3568-
"UTC",
3569-
"America/New_York",
3570-
"America/Los_Angeles",
3571-
"Europe/London",
3572-
"Asia/Tokyo",
3573-
"Australia/Sydney",
3563+
let timezones: [Option<Arc<str>>; 6] = [
3564+
Some(Arc::from("UTC")),
3565+
Some(Arc::from("America/New_York")),
3566+
Some(Arc::from("America/Los_Angeles")),
3567+
Some(Arc::from("Europe/London")),
3568+
Some(Arc::from("Asia/Tokyo")),
3569+
Some(Arc::from("Australia/Sydney")),
35743570
];
35753571

3576-
let eval_modes = [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try];
3577-
3578-
for tz in timezones {
3579-
for eval_mode in &eval_modes {
3580-
let cast_options = SparkCastOptions::new(*eval_mode, tz, false);
3581-
3582-
let int8_array: ArrayRef = Arc::new(Int8Array::from(vec![
3583-
Some(0),
3584-
Some(1),
3585-
Some(-1),
3586-
Some(127),
3587-
Some(-128),
3588-
None,
3589-
]));
3590-
3591-
let result = cast_int_to_timestamp(&int8_array, &cast_options).unwrap();
3592-
let ts_array = result.as_primitive::<TimestampMicrosecondType>();
3593-
3594-
assert_eq!(ts_array.value(0), 0);
3595-
assert_eq!(ts_array.value(1), 1_000_000);
3596-
assert_eq!(ts_array.value(2), -1_000_000);
3597-
assert_eq!(ts_array.value(3), 127_000_000);
3598-
assert_eq!(ts_array.value(4), -128_000_000);
3599-
assert!(ts_array.is_null(5));
3600-
assert_eq!(ts_array.timezone(), Some(tz));
3601-
3602-
let int16_array: ArrayRef = Arc::new(Int16Array::from(vec![
3603-
Some(0),
3604-
Some(1),
3605-
Some(-1),
3606-
Some(32767),
3607-
Some(-32768),
3608-
None,
3609-
]));
3610-
3611-
let result = cast_int_to_timestamp(&int16_array, &cast_options).unwrap();
3612-
let ts_array = result.as_primitive::<TimestampMicrosecondType>();
3613-
3614-
assert_eq!(ts_array.value(0), 0);
3615-
assert_eq!(ts_array.value(1), 1_000_000);
3616-
assert_eq!(ts_array.value(2), -1_000_000);
3617-
assert_eq!(ts_array.value(3), 32_767_000_000_i64);
3618-
assert_eq!(ts_array.value(4), -32_768_000_000_i64);
3619-
assert!(ts_array.is_null(5));
3620-
assert_eq!(ts_array.timezone(), Some(tz));
3621-
3622-
let int32_array: ArrayRef = Arc::new(Int32Array::from(vec![
3623-
Some(0),
3624-
Some(1),
3625-
Some(-1),
3626-
Some(1704067200),
3627-
None,
3628-
]));
3629-
3630-
let result = cast_int_to_timestamp(&int32_array, &cast_options).unwrap();
3631-
let ts_array = result.as_primitive::<TimestampMicrosecondType>();
3632-
3633-
assert_eq!(ts_array.value(0), 0);
3634-
assert_eq!(ts_array.value(1), 1_000_000);
3635-
assert_eq!(ts_array.value(2), -1_000_000);
3636-
assert_eq!(ts_array.value(3), 1_704_067_200_000_000_i64);
3637-
assert!(ts_array.is_null(4));
3638-
assert_eq!(ts_array.timezone(), Some(tz));
3639-
3640-
let int64_array: ArrayRef = Arc::new(Int64Array::from(vec![
3641-
Some(0),
3642-
Some(1),
3643-
Some(-1),
3644-
Some(i64::MAX),
3645-
Some(i64::MIN),
3646-
]));
3647-
3648-
let result = cast_int_to_timestamp(&int64_array, &cast_options).unwrap();
3649-
let ts_array = result.as_primitive::<TimestampMicrosecondType>();
3650-
3651-
assert_eq!(ts_array.value(0), 0);
3652-
assert_eq!(ts_array.value(1), 1_000_000_i64);
3653-
assert_eq!(ts_array.value(2), -1_000_000_i64);
3654-
assert_eq!(ts_array.value(3), i64::MAX);
3655-
assert_eq!(ts_array.value(4), i64::MIN);
3656-
assert_eq!(ts_array.timezone(), Some(tz));
3657-
}
3572+
for tz in &timezones {
3573+
let int8_array: ArrayRef = Arc::new(Int8Array::from(vec![
3574+
Some(0),
3575+
Some(1),
3576+
Some(-1),
3577+
Some(127),
3578+
Some(-128),
3579+
None,
3580+
]));
3581+
3582+
let result = cast_int_to_timestamp(&int8_array, tz).unwrap();
3583+
let ts_array = result.as_primitive::<TimestampMicrosecondType>();
3584+
3585+
assert_eq!(ts_array.value(0), 0);
3586+
assert_eq!(ts_array.value(1), 1_000_000);
3587+
assert_eq!(ts_array.value(2), -1_000_000);
3588+
assert_eq!(ts_array.value(3), 127_000_000);
3589+
assert_eq!(ts_array.value(4), -128_000_000);
3590+
assert!(ts_array.is_null(5));
3591+
assert_eq!(ts_array.timezone(), tz.as_ref().map(|s| s.as_ref()));
3592+
3593+
let int16_array: ArrayRef = Arc::new(Int16Array::from(vec![
3594+
Some(0),
3595+
Some(1),
3596+
Some(-1),
3597+
Some(32767),
3598+
Some(-32768),
3599+
None,
3600+
]));
3601+
3602+
let result = cast_int_to_timestamp(&int16_array, tz).unwrap();
3603+
let ts_array = result.as_primitive::<TimestampMicrosecondType>();
3604+
3605+
assert_eq!(ts_array.value(0), 0);
3606+
assert_eq!(ts_array.value(1), 1_000_000);
3607+
assert_eq!(ts_array.value(2), -1_000_000);
3608+
assert_eq!(ts_array.value(3), 32_767_000_000_i64);
3609+
assert_eq!(ts_array.value(4), -32_768_000_000_i64);
3610+
assert!(ts_array.is_null(5));
3611+
assert_eq!(ts_array.timezone(), tz.as_ref().map(|s| s.as_ref()));
3612+
3613+
let int32_array: ArrayRef = Arc::new(Int32Array::from(vec![
3614+
Some(0),
3615+
Some(1),
3616+
Some(-1),
3617+
Some(1704067200),
3618+
None,
3619+
]));
3620+
3621+
let result = cast_int_to_timestamp(&int32_array, tz).unwrap();
3622+
let ts_array = result.as_primitive::<TimestampMicrosecondType>();
3623+
3624+
assert_eq!(ts_array.value(0), 0);
3625+
assert_eq!(ts_array.value(1), 1_000_000);
3626+
assert_eq!(ts_array.value(2), -1_000_000);
3627+
assert_eq!(ts_array.value(3), 1_704_067_200_000_000_i64);
3628+
assert!(ts_array.is_null(4));
3629+
assert_eq!(ts_array.timezone(), tz.as_ref().map(|s| s.as_ref()));
3630+
3631+
let int64_array: ArrayRef = Arc::new(Int64Array::from(vec![
3632+
Some(0),
3633+
Some(1),
3634+
Some(-1),
3635+
Some(i64::MAX),
3636+
Some(i64::MIN),
3637+
]));
3638+
3639+
let result = cast_int_to_timestamp(&int64_array, tz).unwrap();
3640+
let ts_array = result.as_primitive::<TimestampMicrosecondType>();
3641+
3642+
assert_eq!(ts_array.value(0), 0);
3643+
assert_eq!(ts_array.value(1), 1_000_000_i64);
3644+
assert_eq!(ts_array.value(2), -1_000_000_i64);
3645+
assert_eq!(ts_array.value(3), i64::MAX);
3646+
assert_eq!(ts_array.value(4), i64::MIN);
3647+
assert_eq!(ts_array.timezone(), tz.as_ref().map(|s| s.as_ref()));
36583648
}
36593649
}
36603650
}

0 commit comments

Comments
 (0)