This is an automated email from the ASF dual-hosted git repository.
jayzhan211 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 980667fc52 Add interval multiplication by i64 (#10336)
980667fc52 is described below
commit 980667fc52fc4cce49489f3cd2eedf832ec0022a
Author: Peter Lee <[email protected]>
AuthorDate: Thu Jul 16 21:54:10 2026 +0800
Add interval multiplication by i64 (#10336)
# Which issue does this PR close?
Part of https://github.com/apache/arrow-rs/issues/9030.
# Rationale for this change
As part of updating Datafusion to better support date / time math
operations (https://github.com/apache/datafusion/issues/19022 among
others) it was uncovered that `arrow-rs` does not support interval *
number ops that should be valid.
For example:
interval '1 second' * 900 → 00:15:00
interval '1 day' * 21 → 21 days
this pr only include the interval multiply integer support because div
and floating is too complex in case of
https://github.com/apache/arrow-rs/pull/6906#discussion_r1894338923
# What changes are included in this PR?
Allow:
- (Int64, Interval(YearMonth))
- (Int64, Interval(DayTime))
- (Int64, Interval(MonthDayNano)
and their swap to multiple pair by pair.
# Are these changes tested?
- `cargo test -p arrow-arith`
# Are there any user-facing changes?
yes, `arithmetic_op` now accept
1. (Interval(YearMonth), Int64)
2. (Interval(DayTime), Int64)
3. (Interval(MonthDayNano), Int64)
4. (Int64, Interval(YearMonth))
5. (Int64, Interval(DayTime))
6. (Int64, Interval(MonthDayNano))
as its (lhs, rhs) pair type.
---
arrow-arith/src/numeric.rs | 117 +++++++++++++++++++++++++++++++++++++++++++++
1 file changed, 117 insertions(+)
diff --git a/arrow-arith/src/numeric.rs b/arrow-arith/src/numeric.rs
index 7e0bb2e7aa..cc94006fae 100644
--- a/arrow-arith/src/numeric.rs
+++ b/arrow-arith/src/numeric.rs
@@ -243,6 +243,12 @@ fn arithmetic_op(op: Op, lhs: &dyn Datum, rhs: &dyn Datum)
-> Result<ArrayRef, A
(Duration(Millisecond), Duration(Millisecond)) =>
duration_op::<DurationMillisecondType>(op, l, l_scalar, r, r_scalar),
(Duration(Microsecond), Duration(Microsecond)) =>
duration_op::<DurationMicrosecondType>(op, l, l_scalar, r, r_scalar),
(Duration(Nanosecond), Duration(Nanosecond)) =>
duration_op::<DurationNanosecondType>(op, l, l_scalar, r, r_scalar),
+ (Interval(YearMonth), Int64) if matches!(op, Op::Mul) =>
interval_mul_op::<IntervalYearMonthType>(l, l_scalar, r, r_scalar),
+ (Interval(DayTime), Int64) if matches!(op, Op::Mul) =>
interval_mul_op::<IntervalDayTimeType>(l, l_scalar, r, r_scalar),
+ (Interval(MonthDayNano), Int64) if matches!(op, Op::Mul) =>
interval_mul_op::<IntervalMonthDayNanoType>(l, l_scalar, r, r_scalar),
+ (Int64, Interval(YearMonth)) if matches!(op, Op::Mul) =>
interval_mul_op::<IntervalYearMonthType>(r, r_scalar, l, l_scalar),
+ (Int64, Interval(DayTime)) if matches!(op, Op::Mul) =>
interval_mul_op::<IntervalDayTimeType>(r, r_scalar, l, l_scalar),
+ (Int64, Interval(MonthDayNano)) if matches!(op, Op::Mul) =>
interval_mul_op::<IntervalMonthDayNanoType>(r, r_scalar, l, l_scalar),
(Interval(YearMonth), Interval(YearMonth)) =>
interval_op::<IntervalYearMonthType>(op, l, l_scalar, r, r_scalar),
(Interval(DayTime), Interval(DayTime)) =>
interval_op::<IntervalDayTimeType>(op, l, l_scalar, r, r_scalar),
(Interval(MonthDayNano), Interval(MonthDayNano)) =>
interval_op::<IntervalMonthDayNanoType>(op, l, l_scalar, r, r_scalar),
@@ -621,6 +627,14 @@ date!(Date64Type);
trait IntervalOp: ArrowPrimitiveType {
fn add(left: Self::Native, right: Self::Native) -> Result<Self::Native,
ArrowError>;
fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native,
ArrowError>;
+ fn mul_i64(left: Self::Native, right: i64) -> Result<Self::Native,
ArrowError>;
+}
+
+fn mul_i32_i64(left: i32, right: i64) -> Result<i32, ArrowError> {
+ let value = i64::from(left).mul_checked(right)?;
+ i32::try_from(value).map_err(|_| {
+ ArrowError::ArithmeticOverflow(format!("Overflow happened on: {left} *
{right}"))
+ })
}
impl IntervalOp for IntervalYearMonthType {
@@ -631,6 +645,10 @@ impl IntervalOp for IntervalYearMonthType {
fn sub(left: Self::Native, right: Self::Native) -> Result<Self::Native,
ArrowError> {
left.sub_checked(right)
}
+
+ fn mul_i64(left: Self::Native, right: i64) -> Result<Self::Native,
ArrowError> {
+ mul_i32_i64(left, right)
+ }
}
impl IntervalOp for IntervalDayTimeType {
@@ -649,6 +667,14 @@ impl IntervalOp for IntervalDayTimeType {
let ms = l_ms.sub_checked(r_ms)?;
Ok(Self::make_value(days, ms))
}
+
+ fn mul_i64(left: Self::Native, right: i64) -> Result<Self::Native,
ArrowError> {
+ let (days, ms) = Self::to_parts(left);
+ Ok(Self::make_value(
+ mul_i32_i64(days, right)?,
+ mul_i32_i64(ms, right)?,
+ ))
+ }
}
impl IntervalOp for IntervalMonthDayNanoType {
@@ -669,6 +695,33 @@ impl IntervalOp for IntervalMonthDayNanoType {
let nanos = l_nanos.sub_checked(r_nanos)?;
Ok(Self::make_value(months, days, nanos))
}
+
+ fn mul_i64(left: Self::Native, right: i64) -> Result<Self::Native,
ArrowError> {
+ let (months, days, nanos) = Self::to_parts(left);
+ Ok(Self::make_value(
+ mul_i32_i64(months, right)?,
+ mul_i32_i64(days, right)?,
+ nanos.mul_checked(right)?,
+ ))
+ }
+}
+
+fn interval_mul_op<T: IntervalOp>(
+ interval: &dyn Array,
+ interval_scalar: bool,
+ factor: &dyn Array,
+ factor_scalar: bool,
+) -> Result<ArrayRef, ArrowError> {
+ let interval = interval.as_primitive::<T>();
+ let factor = factor.as_primitive::<Int64Type>();
+ Ok(try_op_ref!(
+ T,
+ interval,
+ interval_scalar,
+ factor,
+ factor_scalar,
+ T::mul_i64(interval, factor)
+ ))
}
/// Perform arithmetic operation on an interval array
@@ -1555,6 +1608,70 @@ mod tests {
);
}
+ #[test]
+ fn test_interval_mul_i64() {
+ let interval = IntervalYearMonthArray::from(vec![16, 5, 0]);
+ let factor = Int64Array::from(vec![3, -2, i64::MAX]);
+ let expected = IntervalYearMonthArray::from(vec![48, -10, 0]);
+ assert_eq!(mul(&interval, &factor).unwrap().as_ref(), &expected);
+ assert_eq!(mul(&factor, &interval).unwrap().as_ref(), &expected);
+
+ let interval = IntervalDayTimeArray::from(vec![
+ Some(IntervalDayTimeType::make_value(10, 2 * 60 * 60 * 1000)),
+ None,
+ ]);
+ let factor = Int64Array::new_scalar(3);
+ let expected = IntervalDayTimeArray::from(vec![
+ Some(IntervalDayTimeType::make_value(30, 6 * 60 * 60 * 1000)),
+ None,
+ ]);
+ assert_eq!(mul(&interval, &factor).unwrap().as_ref(), &expected);
+ assert_eq!(mul(&factor, &interval).unwrap().as_ref(), &expected);
+
+ let null_factor = Scalar::new(Int64Array::new_null(1));
+ let expected = IntervalDayTimeArray::new_null(interval.len());
+ assert_eq!(mul(&interval, &null_factor).unwrap().as_ref(), &expected);
+ assert_eq!(mul(&null_factor, &interval).unwrap().as_ref(), &expected);
+
+ let interval =
IntervalMonthDayNanoArray::new_scalar(IntervalMonthDayNanoType::make_value(
+ 12,
+ 15,
+ 5_000_000_000,
+ ));
+ let factor = Int64Array::from(vec![2, 0, -1]);
+ let expected = IntervalMonthDayNanoArray::from(vec![
+ IntervalMonthDayNanoType::make_value(24, 30, 10_000_000_000),
+ IntervalMonthDayNanoType::make_value(0, 0, 0),
+ IntervalMonthDayNanoType::make_value(-12, -15, -5_000_000_000),
+ ]);
+ assert_eq!(mul(&interval, &factor).unwrap().as_ref(), &expected);
+ assert_eq!(mul(&factor, &interval).unwrap().as_ref(), &expected);
+
+ let float_factor = Float64Array::new_scalar(2.);
+ assert!(mul(&interval, &float_factor).is_err());
+ assert!(mul_wrapping(&factor, &interval).is_err());
+ }
+
+ #[test]
+ fn test_interval_mul_i64_overflow() {
+ let interval = IntervalYearMonthArray::from(vec![i32::MAX]);
+ let factor = Int64Array::from(vec![2]);
+ assert_eq!(
+ mul(&interval, &factor).unwrap_err().to_string(),
+ "Arithmetic overflow: Overflow happened on: 2147483647 * 2"
+ );
+
+ let interval =
IntervalMonthDayNanoArray::from(vec![IntervalMonthDayNanoType::make_value(
+ 0,
+ 0,
+ i64::MAX,
+ )]);
+ assert_eq!(
+ mul(&interval, &factor).unwrap_err().to_string(),
+ "Arithmetic overflow: Overflow happened on: 9223372036854775807 *
2"
+ );
+ }
+
fn test_duration_impl<T: ArrowPrimitiveType<Native = i64>>() {
let a = PrimitiveArray::<T>::new(vec![1000, 4394, -3944].into(), None);
let b = PrimitiveArray::<T>::new(vec![4, -5, -243].into(), None);
