fix: add support for Decimal128 and Decimal256 types in interval arithmetic

datafusion/expr-common/src/interval_arithmetic.rs

@@ -76,6 +76,14 @@ macro_rules! get_extreme_value {
             DataType::Interval(IntervalUnit::MonthDayNano) => {
                 ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano::$extreme))
             }
+            DataType::Decimal128(precision, scale) => {
+                ScalarValue::Decimal128(Some(i128::$extreme), *precision, *scale)
+            }
+            DataType::Decimal256(precision, scale) => ScalarValue::Decimal256(
+                Some(arrow::datatypes::i256::$extreme),
+                *precision,
+                *scale,
+            ),
             _ => unreachable!(),
         }
     };
@@ -1008,17 +1016,20 @@ fn handle_overflow<const UPPER: bool>(
     lhs: &ScalarValue,
     rhs: &ScalarValue,
 ) -> ScalarValue {
-    let zero = ScalarValue::new_zero(dt).unwrap();
+    let lhs_zero = ScalarValue::new_zero(&lhs.data_type()).unwrap();
+    let rhs_zero = ScalarValue::new_zero(&rhs.data_type()).unwrap();
     let positive_sign = match op {
         Operator::Multiply | Operator::Divide => {
-            lhs.lt(&zero) && rhs.lt(&zero) || lhs.gt(&zero) && rhs.gt(&zero)
+            lhs.lt(&lhs_zero) && rhs.lt(&rhs_zero)
+                || lhs.gt(&lhs_zero) && rhs.gt(&rhs_zero)
         }
-        Operator::Plus => lhs.ge(&zero),
+        Operator::Plus => lhs.ge(&lhs_zero),
         Operator::Minus => lhs.ge(rhs),
         _ => {
             unreachable!()
         }
     };
+
     match (UPPER, positive_sign) {
         (true, true) | (false, false) => ScalarValue::try_from(dt).unwrap(),
         (true, false) => {
@@ -1832,7 +1843,12 @@ impl NullableInterval {
 
 #[cfg(test)]
 mod tests {
-    use crate::interval_arithmetic::{next_value, prev_value, satisfy_greater, Interval};
+    use crate::{
+        interval_arithmetic::{
+            handle_overflow, next_value, prev_value, satisfy_greater, Interval,
+        },
+        operator::Operator,
+    };
 
     use arrow::datatypes::DataType;
     use datafusion_common::{Result, ScalarValue};
@@ -3108,6 +3124,70 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_overflow_handling() -> Result<()> {
+        // Test integer overflow handling:
+        let dt = DataType::Int32;
+        let op = Operator::Plus;
+        let lhs = ScalarValue::Int32(Some(i32::MAX));
+        let rhs = ScalarValue::Int32(Some(1));
+        let result = handle_overflow::<true>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Int32(None));
+        let result = handle_overflow::<false>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Int32(Some(i32::MAX)));
+
+        // Test float overflow handling:
+        let dt = DataType::Float32;
+        let op = Operator::Multiply;
+        let lhs = ScalarValue::Float32(Some(f32::MAX));
+        let rhs = ScalarValue::Float32(Some(2.0));
+        let result = handle_overflow::<true>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Float32(None));
+        let result = handle_overflow::<false>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Float32(Some(f32::MAX)));
+
+        // Test float underflow handling:
+        let lhs = ScalarValue::Float32(Some(f32::MIN));
+        let rhs = ScalarValue::Float32(Some(2.0));
+        let result = handle_overflow::<true>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Float32(Some(f32::MIN)));
+        let result = handle_overflow::<false>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Float32(None));
+
+        // Test integer underflow handling:
+        let dt = DataType::Int64;
+        let op = Operator::Minus;
+        let lhs = ScalarValue::Int64(Some(i64::MIN));
+        let rhs = ScalarValue::Int64(Some(1));
+        let result = handle_overflow::<true>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Int64(Some(i64::MIN)));
+        let result = handle_overflow::<false>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Int64(None));
+
+        // Test unsigned integer handling:
+        let dt = DataType::UInt32;
+        let op = Operator::Minus;
+        let lhs = ScalarValue::UInt32(Some(0));
+        let rhs = ScalarValue::UInt32(Some(1));
+        let result = handle_overflow::<true>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::UInt32(Some(0)));
+        let result = handle_overflow::<false>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::UInt32(None));
+
+        // Test decimal handling:
+        let dt = DataType::Decimal128(38, 35);
+        let op = Operator::Plus;
+        let lhs =
+            ScalarValue::Decimal128(Some(54321543215432154321543215432154321), 35, 35);
+        let rhs = ScalarValue::Decimal128(Some(10000), 20, 0);
+        let result = handle_overflow::<true>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Decimal128(None, 38, 35));
+        let result = handle_overflow::<false>(&dt, op, &lhs, &rhs);
+        assert_eq!(result, ScalarValue::Decimal128(Some(i128::MAX), 38, 35));
+
+        Ok(())
+    }
+
     #[test]
     fn test_cardinality_of_intervals() -> Result<()> {
         // In IEEE 754 standard for floating-point arithmetic, if we keep the sign and exponent fields same,