Skip to content

Commit

Permalink
feat: add support for casting Duration/Interval to Int64Array (#1196)
Browse files Browse the repository at this point in the history
* feat: add support for casting Duration to Int64Array

* feat: cast from Interval to Int64
  • Loading branch information
e-dard authored Jan 19, 2022
1 parent d68c4ae commit 799330b
Showing 1 changed file with 57 additions and 1 deletion.
58 changes: 57 additions & 1 deletion arrow/src/compute/kernels/cast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,14 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
(Timestamp(_, _), Timestamp(_, _) | Date32 | Date64) => true,
// date64 to timestamp might not make sense,
(Int64, Duration(_)) => true,
(Duration(_), Int64) => true,
(Interval(from_type), Int64) => {
match from_type{
IntervalUnit::YearMonth => true,
IntervalUnit::DayTime => true,
IntervalUnit::MonthDayNano => false, // Native type is i128
}
}
(_, _) => false,
}
}
Expand Down Expand Up @@ -358,7 +366,6 @@ macro_rules! cast_decimal_to_float {
/// * To or from `StructArray`
/// * List to primitive
/// * Utf8 to boolean
/// * Interval and duration
pub fn cast_with_options(
array: &ArrayRef,
to_type: &DataType,
Expand Down Expand Up @@ -1111,6 +1118,17 @@ pub fn cast_with_options(
}
}
}
(Duration(_), Int64) => cast_array_data::<Int64Type>(array, to_type.clone()),
(Interval(from_type), Int64) => match from_type {
IntervalUnit::YearMonth => {
cast_numeric_arrays::<IntervalYearMonthType, Int64Type>(array)
}
IntervalUnit::DayTime => cast_array_data::<Int64Type>(array, to_type.clone()),
IntervalUnit::MonthDayNano => Err(ArrowError::CastError(format!(
"Casting from {:?} to {:?} not supported",
from_type, to_type,
))),
},
(_, _) => Err(ArrowError::CastError(format!(
"Casting from {:?} to {:?} not supported",
from_type, to_type,
Expand Down Expand Up @@ -2765,6 +2783,44 @@ mod tests {
assert!(c.is_null(2));
}

#[test]
fn test_cast_duration_to_i64() {
let base = vec![5, 6, 7, 8, 100000000];

let duration_arrays = vec![
Arc::new(DurationNanosecondArray::from(base.clone())) as ArrayRef,
Arc::new(DurationMicrosecondArray::from(base.clone())) as ArrayRef,
Arc::new(DurationMillisecondArray::from(base.clone())) as ArrayRef,
Arc::new(DurationSecondArray::from(base.clone())) as ArrayRef,
];

for arr in duration_arrays {
assert!(can_cast_types(arr.data_type(), &DataType::Int64));
let result = cast(&arr, &DataType::Int64).unwrap();
let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
assert_eq!(base.as_slice(), result.values());
}
}

#[test]
fn test_cast_interval_to_i64() {
let base = vec![5, 6, 7, 8];

let interval_arrays = vec![
Arc::new(IntervalDayTimeArray::from(base.clone())) as ArrayRef,
Arc::new(IntervalYearMonthArray::from(
base.iter().map(|x| *x as i32).collect::<Vec<i32>>(),
)) as ArrayRef,
];

for arr in interval_arrays {
assert!(can_cast_types(arr.data_type(), &DataType::Int64));
let result = cast(&arr, &DataType::Int64).unwrap();
let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
assert_eq!(base.as_slice(), result.values());
}
}

#[test]
fn test_cast_to_strings() {
let a = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
Expand Down

0 comments on commit 799330b

Please sign in to comment.