feat: interval add timestamp

apache · Mar 6, 2023 · 8e65992 · 8e65992
1 parent 17b2f11
commit 8e65992
Showing 1 changed file with 38 additions and 0 deletions.
diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs
@@ -516,18 +516,34 @@ macro_rules! impl_op {
                 let value = seconds_add(*ts_s, $RHS, get_sign!($OPERATION))?;
                 Ok(ScalarValue::TimestampSecond(Some(value), zone.clone()))
             }
+            (_, ScalarValue::TimestampSecond(Some(ts_s), zone)) => {
+                let value = seconds_add(*ts_s, $LHS, get_sign!($OPERATION))?;
+                Ok(ScalarValue::TimestampSecond(Some(value), zone.clone()))
+            }
             (ScalarValue::TimestampMillisecond(Some(ts_ms), zone), _) => {
                 let value = milliseconds_add(*ts_ms, $RHS, get_sign!($OPERATION))?;
                 Ok(ScalarValue::TimestampMillisecond(Some(value), zone.clone()))
             }
+            (_, ScalarValue::TimestampMillisecond(Some(ts_ms), zone)) => {
+                let value = milliseconds_add(*ts_ms, $LHS, get_sign!($OPERATION))?;
+                Ok(ScalarValue::TimestampMillisecond(Some(value), zone.clone()))
+            }
             (ScalarValue::TimestampMicrosecond(Some(ts_us), zone), _) => {
                 let value = microseconds_add(*ts_us, $RHS, get_sign!($OPERATION))?;
                 Ok(ScalarValue::TimestampMicrosecond(Some(value), zone.clone()))
             }
+            (_, ScalarValue::TimestampMicrosecond(Some(ts_us), zone)) => {
+                let value = microseconds_add(*ts_us, $LHS, get_sign!($OPERATION))?;
+                Ok(ScalarValue::TimestampMicrosecond(Some(value), zone.clone()))
+            }
             (ScalarValue::TimestampNanosecond(Some(ts_ns), zone), _) => {
                 let value = nanoseconds_add(*ts_ns, $RHS, get_sign!($OPERATION))?;
                 Ok(ScalarValue::TimestampNanosecond(Some(value), zone.clone()))
             }
+            (_, ScalarValue::TimestampNanosecond(Some(ts_ns), zone)) => {
+                let value = nanoseconds_add(*ts_ns, $LHS, get_sign!($OPERATION))?;
+                Ok(ScalarValue::TimestampNanosecond(Some(value), zone.clone()))
+            }
             _ => Err(DataFusionError::Internal(format!(
                 "Operator {} is not implemented for types {:?} and {:?}",
                 stringify!($OPERATION),
@@ -2911,6 +2927,28 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_interval_add_timestamp() -> Result<()> {
+        let interval = ScalarValue::IntervalMonthDayNano(Some(123));
+        let timestamp = ScalarValue::TimestampNanosecond(Some(123), None);
+        let result = interval.add(&timestamp)?;
+        let expect = timestamp.add(&interval)?;
+        assert_eq!(result, expect);
+
+        let interval = ScalarValue::IntervalYearMonth(Some(123));
+        let timestamp = ScalarValue::TimestampNanosecond(Some(123), None);
+        let result = interval.add(&timestamp)?;
+        let expect = timestamp.add(&interval)?;
+        assert_eq!(result, expect);
+
+        let interval = ScalarValue::IntervalDayTime(Some(123));
+        let timestamp = ScalarValue::TimestampNanosecond(Some(123), None);
+        let result = interval.add(&timestamp)?;
+        let expect = timestamp.add(&interval)?;
+        assert_eq!(result, expect);
+        Ok(())
+    }
+
     #[test]
     fn scalar_decimal_test() -> Result<()> {
         let decimal_value = ScalarValue::Decimal128(Some(123), 10, 1);