Skip to content

Commit

Permalink
fix(arrow): Fix make_value/to_parts for MonthDayNano interval
Browse files Browse the repository at this point in the history
This implementation was aligned with https://github.com/cube-js/arrow-datafusion/blob/370f91fda4a7e9387dded4edaf2425c447e76fbd/datafusion/physical-expr/src/expressions/binary_distinct.rs#L306-L353

Can drop this after rebase on commit d4f038a " Update IntervalMonthDayNanoType::make_value() to conform to specifications (apache#2235)", first released in 20.0.0
  • Loading branch information
mcheshkov committed Aug 21, 2024
1 parent a3d3b4d commit c39a096
Showing 1 changed file with 53 additions and 6 deletions.
59 changes: 53 additions & 6 deletions arrow/src/datatypes/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -261,9 +261,10 @@ impl IntervalMonthDayNanoType {
days: i32,
nanos: i64,
) -> <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native {
let m = months as u128 & u32::MAX as u128;
let d = (days as u128 & u32::MAX as u128) << 32;
let n = (nanos as u128) << 64;
// casts thru u32/u64 to avoid sign extension
let m = (months as u32 as u128) << 96;
let d = (days as u32 as u128) << 64;
let n = nanos as u64 as u128;
(m | d | n) as <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native
}

Expand All @@ -275,9 +276,55 @@ impl IntervalMonthDayNanoType {
pub fn to_parts(
i: <IntervalMonthDayNanoType as ArrowPrimitiveType>::Native,
) -> (i32, i32, i64) {
let nanos = (i >> 64) as i64;
let days = (i >> 32) as i32;
let months = i as i32;
let i = i as u128;
let nanos = (i & (u64::MAX as u128)) as i64;
let days = ((i >> 64) & (u32::MAX as u128)) as i32;
let months = ((i >> 96) & (u32::MAX as u128)) as i32;
(months, days, nanos)
}
}

#[cfg(test)]
mod tests {
use super::{IntervalDayTimeType, IntervalMonthDayNanoType};

#[test]
fn interval_day_time_roundtrip() {
fn check(days: i32, millis: i32) {
let interval = IntervalDayTimeType::make_value(days, millis);
let parts = IntervalDayTimeType::to_parts(interval);
assert_eq!((days, millis), parts);
}

check(0, 0);

check(0, 1);
check(1, 0);
check(0, -1);
check(-1, 0);

check(1, 2);
check(-1, -2);
}

#[test]
fn interval_month_day_nano_roundtrip() {
fn check(months: i32, days: i32, nanos: i64) {
let interval = IntervalMonthDayNanoType::make_value(months, days, nanos);
let parts = IntervalMonthDayNanoType::to_parts(interval);
assert_eq!((months, days, nanos), parts);
}

check(0, 0, 0);

check(0, 0, 1);
check(0, 1, 0);
check(1, 0, 0);
check(0, 0, -1);
check(0, -1, 0);
check(-1, 0, 0);

check(1, 2, 3);
check(-1, -2, -3);
}
}

0 comments on commit c39a096

Please sign in to comment.