Skip to content

Commit

Permalink
feat: Support casting Date32 to Timestamp
Browse files Browse the repository at this point in the history
  • Loading branch information
MazterQyou committed Dec 15, 2022
1 parent 51f62a7 commit 096ef28
Showing 1 changed file with 40 additions and 0 deletions.
40 changes: 40 additions & 0 deletions arrow/src/compute/kernels/cast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
(Timestamp(_, _), Int64) => true,
(Int64, Timestamp(_, _)) => true,
(Timestamp(_, _), Timestamp(_, _) | Date32 | Date64) => true,
(Date32, Timestamp(_, _)) => true,
// date64 to timestamp might not make sense,
(Int64, Duration(_)) => true,
(Duration(_), Int64) => true,
Expand Down Expand Up @@ -1161,6 +1162,32 @@ pub fn cast_with_options(

Ok(Arc::new(b.finish()) as ArrayRef)
}
(Date32, Timestamp(to_unit, _)) => {
let time_array = Int32Array::from(array.data().clone());
let time_array = numeric_cast::<Int32Type, Int64Type>(&time_array);
let to_size = time_unit_multiple(to_unit) * SECONDS_IN_DAY;
let converted =
multiply(&time_array, &Int64Array::from(vec![to_size; array.len()]))?;
let array_ref = Arc::new(converted) as ArrayRef;
use TimeUnit::*;
match to_unit {
Second => {
cast_array_data::<TimestampSecondType>(&array_ref, to_type.clone())
}
Millisecond => cast_array_data::<TimestampMillisecondType>(
&array_ref,
to_type.clone(),
),
Microsecond => cast_array_data::<TimestampMicrosecondType>(
&array_ref,
to_type.clone(),
),
Nanosecond => cast_array_data::<TimestampNanosecondType>(
&array_ref,
to_type.clone(),
),
}
}
(Timestamp(from_unit, _), Date64) => {
let from_size = time_unit_multiple(from_unit);
let to_size = MILLISECONDS;
Expand Down Expand Up @@ -2970,6 +2997,19 @@ mod tests {
assert!(c.is_null(2));
}

#[test]
fn test_cast_date32_to_timestamp() {
let a = Date32Array::from(vec![10000, 17890]);
let array = Arc::new(a) as ArrayRef;
let b = cast(&array, &DataType::Timestamp(TimeUnit::Nanosecond, None)).unwrap();
let c = b
.as_any()
.downcast_ref::<TimestampNanosecondArray>()
.unwrap();
assert_eq!(864000000000000000, c.value(0));
assert_eq!(1545696000000000000, c.value(1));
}

#[test]
fn test_cast_timestamp_to_date64() {
let a = TimestampMillisecondArray::from_opt_vec(
Expand Down

0 comments on commit 096ef28

Please sign in to comment.