tustvold commented on code in PR #4546:
URL: https://github.com/apache/arrow-rs/pull/4546#discussion_r1273920965


##########
arrow-arith/src/numeric.rs:
##########
@@ -341,50 +344,201 @@ fn float_op<T: ArrowPrimitiveType>(
     Ok(Arc::new(array))
 }
 
-/// Arithmetic trait for timestamp arrays
-trait TimestampOp: ArrowTimestampType {
-    type Duration: ArrowPrimitiveType<Native = i64>;
+/// Add the given number of months to the given datetime.
+///
+/// Returns `None` when it will result in overflow.
+fn add_months_datetime<Tz: TimeZone>(
+    dt: DateTime<Tz>,
+    months: i32,
+) -> Option<DateTime<Tz>> {
+    match months.cmp(&0) {
+        Ordering::Equal => Some(dt),
+        Ordering::Greater => dt.checked_add_months(Months::new(months as u32)),
+        Ordering::Less => 
dt.checked_sub_months(Months::new(months.unsigned_abs())),
+    }
+}
 
-    fn add_year_month(timestamp: i64, delta: i32) -> Result<i64, ArrowError>;
-    fn add_day_time(timestamp: i64, delta: i64) -> Result<i64, ArrowError>;
-    fn add_month_day_nano(timestamp: i64, delta: i128) -> Result<i64, 
ArrowError>;
+/// Add the given number of days to the given datetime.
+///
+/// Returns `None` when it will result in overflow.
+fn add_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> 
Option<DateTime<Tz>> {
+    match days.cmp(&0) {
+        Ordering::Equal => Some(dt),
+        Ordering::Greater => dt.checked_add_days(Days::new(days as u64)),
+        Ordering::Less => dt.checked_sub_days(Days::new(days.unsigned_abs() as 
u64)),
+    }
+}
 
-    fn sub_year_month(timestamp: i64, delta: i32) -> Result<i64, ArrowError>;
-    fn sub_day_time(timestamp: i64, delta: i64) -> Result<i64, ArrowError>;
-    fn sub_month_day_nano(timestamp: i64, delta: i128) -> Result<i64, 
ArrowError>;
+/// Substract the given number of months to the given datetime.
+///
+/// Returns `None` when it will result in overflow.
+fn sub_months_datetime<Tz: TimeZone>(
+    dt: DateTime<Tz>,
+    months: i32,
+) -> Option<DateTime<Tz>> {
+    match months.cmp(&0) {
+        Ordering::Equal => Some(dt),
+        Ordering::Greater => dt.checked_sub_months(Months::new(months as u32)),
+        Ordering::Less => 
dt.checked_add_months(Months::new(months.unsigned_abs())),
+    }
 }
 
-macro_rules! timestamp {
-    ($t:ty, $d:ty) => {
-        impl TimestampOp for $t {
-            type Duration = $d;
+/// Substract the given number of days to the given datetime.
+///
+/// Returns `None` when it will result in overflow.
+fn sub_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> 
Option<DateTime<Tz>> {
+    match days.cmp(&0) {
+        Ordering::Equal => Some(dt),
+        Ordering::Greater => dt.checked_sub_days(Days::new(days as u64)),
+        Ordering::Less => dt.checked_add_days(Days::new(days.unsigned_abs() as 
u64)),
+    }
+}
 
-            fn add_year_month(left: i64, right: i32) -> Result<i64, 
ArrowError> {
-                Self::add_year_months(left, right)
-            }
+/// Arithmetic trait for timestamp arrays
+pub trait TimestampOp: ArrowTimestampType + Sized {

Review Comment:
   I'm not a fan of making this trait public, as it is intended as an internal 
implementation detail of how the kernel chooses to implement arithmetic. I 
could definitely see it evolving in future to support vectorisation or some 
other extensions.
   
   I think I would prefer to just expose the dyn kernels



-- 
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]

Reply via email to