tustvold commented on code in PR #4493:
URL: https://github.com/apache/arrow-rs/pull/4493#discussion_r1258294631
##########
arrow-arith/src/numeric.rs:
##########
@@ -670,3 +682,490 @@ fn decimal_op<T: DecimalType>(
Ok(Arc::new(array))
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::temporal_conversions::{as_date, as_datetime};
+ use chrono::{DateTime, NaiveDate};
+
+ #[test]
+ fn test_integer() {
+ let a = Int32Array::from(vec![4, 3, 5, -6, 100]);
+ let b = Int32Array::from(vec![6, 2, 5, -7, 3]);
+ let result = add(&a, &b).unwrap();
+ assert_eq!(
+ result.as_ref(),
+ &Int32Array::from(vec![10, 5, 10, -13, 103])
+ );
+ let result = sub(&a, &b).unwrap();
+ assert_eq!(result.as_ref(), &Int32Array::from(vec![-2, 1, 0, 1, 97]));
+ let result = div(&a, &b).unwrap();
+ assert_eq!(result.as_ref(), &Int32Array::from(vec![0, 1, 1, 0, 33]));
+ let result = mul(&a, &b).unwrap();
+ assert_eq!(result.as_ref(), &Int32Array::from(vec![24, 6, 25, 42,
300]));
+ let result = rem(&a, &b).unwrap();
+ assert_eq!(result.as_ref(), &Int32Array::from(vec![4, 1, 0, -6, 1]));
+
+ let a = Int8Array::from(vec![Some(2), None, Some(45)]);
+ let b = Int8Array::from(vec![Some(5), Some(3), None]);
+ let result = add(&a, &b).unwrap();
+ assert_eq!(result.as_ref(), &Int8Array::from(vec![Some(7), None,
None]));
+
+ let a = UInt8Array::from(vec![56, 5, 3]);
+ let b = UInt8Array::from(vec![200, 2, 5]);
+ let err = add(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Compute error: Overflow happened on: 56 + 200");
+ let result = add_wrapping(&a, &b).unwrap();
+ assert_eq!(result.as_ref(), &UInt8Array::from(vec![0, 7, 8]));
+
+ let a = UInt8Array::from(vec![34, 5, 3]);
+ let b = UInt8Array::from(vec![200, 2, 5]);
+ let err = sub(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Compute error: Overflow happened on: 34 - 200");
+ let result = sub_wrapping(&a, &b).unwrap();
+ assert_eq!(result.as_ref(), &UInt8Array::from(vec![90, 3, 254]));
+
+ let a = UInt8Array::from(vec![34, 5, 3]);
+ let b = UInt8Array::from(vec![200, 2, 5]);
+ let err = mul(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Compute error: Overflow happened on: 34 * 200");
+ let result = mul_wrapping(&a, &b).unwrap();
+ assert_eq!(result.as_ref(), &UInt8Array::from(vec![144, 10, 15]));
+
+ let a = Int16Array::from(vec![i16::MIN]);
+ let b = Int16Array::from(vec![-1]);
+ let err = div(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Compute error: Overflow happened on: -32768 / -1");
+
+ let a = Int16Array::from(vec![21]);
+ let b = Int16Array::from(vec![0]);
+ let err = div(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Divide by zero error");
+ }
+
+ #[test]
+ fn test_float() {
+ let a = Float32Array::from(vec![1., f32::MAX, 6., -4., -1., 0.]);
+ let b = Float32Array::from(vec![1., f32::MAX, f32::MAX, -3., 45., 0.]);
+ let result = add(&a, &b).unwrap();
+ assert_eq!(
+ result.as_ref(),
+ &Float32Array::from(vec![2., f32::INFINITY, f32::MAX, -7., 44.0,
0.])
+ );
+
+ let result = sub(&a, &b).unwrap();
+ assert_eq!(
+ result.as_ref(),
+ &Float32Array::from(vec![0., 0., f32::MIN, -1., -46., 0.])
+ );
+
+ let result = mul(&a, &b).unwrap();
+ assert_eq!(
+ result.as_ref(),
+ &Float32Array::from(vec![1., f32::INFINITY, f32::INFINITY, 12.,
-45., 0.])
+ );
+
+ let result = div(&a, &b).unwrap();
+ let r = result.as_primitive::<Float32Type>();
+ assert_eq!(r.value(0), 1.);
+ assert_eq!(r.value(1), 1.);
+ assert!(r.value(2) < f32::EPSILON);
+ assert_eq!(r.value(3), -4. / -3.);
+ assert!(r.value(5).is_nan());
+
+ let result = rem(&a, &b).unwrap();
+ let r = result.as_primitive::<Float32Type>();
+ assert_eq!(&r.values()[..5], &[0., 0., 6., -1., -1.]);
+ assert!(r.value(5).is_nan());
+ }
+
+ #[test]
+ fn test_decimal() {
+ // 0.015 7.842 -0.577 0.334 -0.078 0.003
+ let a = Decimal128Array::from(vec![15, 0, -577, 334, -78, 3])
+ .with_precision_and_scale(12, 3)
+ .unwrap();
+
+ // 5.4 0 -35.6 0.3 0.6 7.45
+ let b = Decimal128Array::from(vec![54, 34, -356, 3, 6, 745])
+ .with_precision_and_scale(12, 1)
+ .unwrap();
+
+ let result = add(&a, &b).unwrap();
+ assert_eq!(result.data_type(), &DataType::Decimal128(15, 3));
+ assert_eq!(
+ result.as_primitive::<Decimal128Type>().values(),
+ &[5415, 3400, -36177, 634, 522, 74503]
+ );
+
+ let result = sub(&a, &b).unwrap();
+ assert_eq!(result.data_type(), &DataType::Decimal128(15, 3));
+ assert_eq!(
+ result.as_primitive::<Decimal128Type>().values(),
+ &[-5385, -3400, 35023, 34, -678, -74497]
+ );
+
+ let result = mul(&a, &b).unwrap();
+ assert_eq!(result.data_type(), &DataType::Decimal128(25, 4));
+ assert_eq!(
+ result.as_primitive::<Decimal128Type>().values(),
+ &[810, 0, 205412, 1002, -468, 2235]
+ );
+
+ let result = div(&a, &b).unwrap();
+ assert_eq!(result.data_type(), &DataType::Decimal128(17, 7));
+ assert_eq!(
+ result.as_primitive::<Decimal128Type>().values(),
+ &[27777, 0, 162078, 11133333, -1300000, 402]
+ );
+
+ let result = rem(&a, &b).unwrap();
+ assert_eq!(result.data_type(), &DataType::Decimal128(12, 3));
+ assert_eq!(
+ result.as_primitive::<Decimal128Type>().values(),
+ &[15, 0, -577, 34, -78, 3]
+ );
+
+ let a = Decimal128Array::from(vec![1])
+ .with_precision_and_scale(3, 3)
+ .unwrap();
+ let b = Decimal128Array::from(vec![1])
+ .with_precision_and_scale(37, 37)
+ .unwrap();
+ let err = mul(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Invalid argument error: Output scale of Decimal128(3,
3) * Decimal128(37, 37) would exceed max scale of 38");
+
+ let a = Decimal128Array::from(vec![1])
+ .with_precision_and_scale(3, -2)
+ .unwrap();
+ let err = add(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Compute error: Overflow happened on: 10 ^ 39");
+
+ let a = Decimal128Array::from(vec![10])
+ .with_precision_and_scale(3, -1)
+ .unwrap();
+ let err = add(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Compute error: Overflow happened on: 10 *
100000000000000000000000000000000000000");
+
+ let b = Decimal128Array::from(vec![0])
+ .with_precision_and_scale(1, 1)
+ .unwrap();
+ let err = div(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Divide by zero error");
+ let err = rem(&a, &b).unwrap_err().to_string();
+ assert_eq!(err, "Divide by zero error");
+ }
+
+ fn test_timestamp_impl<T: TimestampOp>() {
+ let a = PrimitiveArray::<T>::new(vec![2000000, 434030324,
53943340].into(), None);
+ let b = PrimitiveArray::<T>::new(vec![329593, 59349, 694994].into(),
None);
+
+ let result = sub(&a, &b).unwrap();
+ assert_eq!(
+ result.as_primitive::<T::Duration>().values(),
+ &[1670407, 433970975, 53248346]
+ );
+
+ let r2 = add(&b, &result.as_ref()).unwrap();
+ assert_eq!(r2.as_ref(), &a);
+
+ let r3 = add(&result.as_ref(), &b).unwrap();
+ assert_eq!(r3.as_ref(), &a);
+
+ let format_array = |x: &dyn Array| -> Vec<String> {
+ x.as_primitive::<T>()
+ .values()
+ .into_iter()
+ .map(|x| as_datetime::<T>(*x).unwrap().to_string())
+ .collect()
+ };
+
+ let values = vec![
+ "1970-01-01T00:00:00Z",
+ "2010-04-01T04:00:20Z",
+ "1960-01-30T04:23:20Z",
+ ]
+ .into_iter()
+ .map(|x| {
+
T::make_value(DateTime::parse_from_rfc3339(x).unwrap().naive_utc()).unwrap()
+ })
+ .collect();
+
+ let a = PrimitiveArray::<T>::new(values, None);
+ let b = IntervalYearMonthArray::from(vec![
+ IntervalYearMonthType::make_value(5, 34),
+ IntervalYearMonthType::make_value(-2, 4),
Review Comment:
Currently the overflow behaviour is a little broken, #4456 tracks fixing
this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
