feat: Add mathematical operators for IntervalYearMonth type

facebookincubator · Dec 11, 2024 · 0afc30c · 0afc30c
1 parent 929affe
commit 0afc30c
Show file tree

Hide file tree

Showing 3 changed files with 185 additions and 111 deletions.
diff --git a/velox/functions/prestosql/Arithmetic.h b/velox/functions/prestosql/Arithmetic.h
@@ -38,6 +38,8 @@ inline constexpr int kMinRadix = 2;
 inline constexpr int kMaxRadix = 36;
 inline constexpr long kLongMax = std::numeric_limits<int64_t>::max();
 inline constexpr long kLongMin = std::numeric_limits<int64_t>::min();
+inline constexpr long kIntegerMax = std::numeric_limits<int32_t>::max();
+inline constexpr long kIntegerMin = std::numeric_limits<int32_t>::min();
 
 inline constexpr char digits[36] = {
     '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b',
@@ -73,7 +75,8 @@ struct MultiplyFunction {
   }
 };
 
-// Multiply function for IntervalDayTime * Double and Double * IntervalDayTime.
+// Multiply function for IntervalDayTime * Double, Double * IntervalDayTime,
+// IntervalYearMonth * Double and Double * IntervalYearMonth.
 template <typename T>
 struct IntervalMultiplyFunction {
   FOLLY_ALWAYS_INLINE double sanitizeInput(double d) {
@@ -84,25 +87,39 @@ struct IntervalMultiplyFunction {
   }
 
   template <
+      typename TResult,
       typename T1,
       typename T2,
       typename = std::enable_if_t<
           (std::is_same_v<T1, int64_t> && std::is_same_v<T2, double>) ||
-          (std::is_same_v<T1, double> && std::is_same_v<T2, int64_t>)>>
-  FOLLY_ALWAYS_INLINE void call(int64_t& result, T1 a, T2 b) {
+          (std::is_same_v<T1, double> && std::is_same_v<T2, int64_t>) ||
+          (std::is_same_v<T1, int32_t> && std::is_same_v<T2, double>) ||
+          (std::is_same_v<T1, double> && std::is_same_v<T2, int32_t>)>>
+  FOLLY_ALWAYS_INLINE void call(TResult& result, T1 a, T2 b) {
     double resultDouble;
     if constexpr (std::is_same_v<T1, double>) {
       resultDouble = sanitizeInput(a) * b;
     } else {
       resultDouble = sanitizeInput(b) * a;
     }
 
+    TResult min, max, maxCheck;
+    if constexpr (std::is_same_v<TResult, int64_t>) {
+      min = kLongMin;
+      max = kLongMax;
+      maxCheck = kMaxDoubleBelowInt64Max;
+    } else {
+      min = kIntegerMin;
+      max = kIntegerMax;
+      maxCheck = kIntegerMax;
+    }
+
     if LIKELY (
-        std::isfinite(resultDouble) && resultDouble >= kLongMin &&
-        resultDouble <= kMaxDoubleBelowInt64Max) {
-      result = int64_t(resultDouble);
+        std::isfinite(resultDouble) && resultDouble >= min &&
+        resultDouble <= maxCheck) {
+      result = static_cast<TResult>(resultDouble);
     } else {
-      result = resultDouble > 0 ? kLongMax : kLongMin;
+      result = resultDouble > 0 ? max : min;
     }
   }
 };
@@ -123,9 +140,14 @@ struct DivideFunction {
   }
 };
 
+// Divide function for IntervalDayTime / Double and IntervalYearMonth / Double.
 template <typename T>
 struct IntervalDivideFunction {
-  FOLLY_ALWAYS_INLINE void call(int64_t& result, int64_t a, double b)
+  template <
+      typename TResult,
+      typename = std::enable_if_t<
+          std::is_same_v<TResult, int64_t> || std::is_same_v<TResult, int32_t>>>
+  FOLLY_ALWAYS_INLINE void call(TResult& result, TResult a, double b)
 // Depend on compiler have correct behaviour for divide by zero
 #if defined(__has_feature)
 #if __has_feature(__address_sanitizer__)
@@ -134,17 +156,28 @@ struct IntervalDivideFunction {
 #endif
 #endif
   {
+    TResult min, max, maxCheck;
+    if constexpr (std::is_same_v<TResult, int64_t>) {
+      min = kLongMin;
+      max = kLongMax;
+      maxCheck = kMaxDoubleBelowInt64Max;
+    } else {
+      min = kIntegerMin;
+      max = kIntegerMax;
+      maxCheck = kIntegerMax;
+    }
+
     if UNLIKELY (a == 0 || std::isnan(b) || !std::isfinite(b)) {
       result = 0;
       return;
     }
     double resultDouble = a / b;
     if LIKELY (
-        std::isfinite(resultDouble) && resultDouble >= kLongMin &&
-        resultDouble <= kMaxDoubleBelowInt64Max) {
-      result = int64_t(resultDouble);
+        std::isfinite(resultDouble) && resultDouble >= min &&
+        resultDouble <= maxCheck) {
+      result = static_cast<TResult>(resultDouble);
     } else {
-      result = resultDouble > 0 ? kLongMax : kLongMin;
+      result = resultDouble > 0 ? max : min;
     }
   }
 };

diff --git a/velox/functions/prestosql/registration/MathematicalOperatorsRegistration.cpp b/velox/functions/prestosql/registration/MathematicalOperatorsRegistration.cpp
@@ -28,12 +28,22 @@ void registerMathOperators(const std::string& prefix = "") {
       IntervalDayTime,
       IntervalDayTime,
       IntervalDayTime>({prefix + "plus"});
+  registerFunction<
+      PlusFunction,
+      IntervalYearMonth,
+      IntervalYearMonth,
+      IntervalYearMonth>({prefix + "plus"});
   registerBinaryFloatingPoint<MinusFunction>({prefix + "minus"});
   registerFunction<
       MinusFunction,
       IntervalDayTime,
       IntervalDayTime,
       IntervalDayTime>({prefix + "minus"});
+  registerFunction<
+      MinusFunction,
+      IntervalYearMonth,
+      IntervalYearMonth,
+      IntervalYearMonth>({prefix + "minus"});
   registerBinaryFloatingPoint<MultiplyFunction>({prefix + "multiply"});
   registerFunction<MultiplyFunction, IntervalDayTime, IntervalDayTime, int64_t>(
       {prefix + "multiply"});
@@ -49,12 +59,37 @@ void registerMathOperators(const std::string& prefix = "") {
       IntervalDayTime,
       double,
       IntervalDayTime>({prefix + "multiply"});
+  registerFunction<
+      MultiplyFunction,
+      IntervalYearMonth,
+      IntervalYearMonth,
+      int32_t>({prefix + "multiply"});
+  registerFunction<
+      MultiplyFunction,
+      IntervalYearMonth,
+      int32_t,
+      IntervalYearMonth>({prefix + "multiply"});
+  registerFunction<
+      IntervalMultiplyFunction,
+      IntervalYearMonth,
+      IntervalYearMonth,
+      double>({prefix + "multiply"});
+  registerFunction<
+      IntervalMultiplyFunction,
+      IntervalYearMonth,
+      double,
+      IntervalYearMonth>({prefix + "multiply"});
   registerBinaryFloatingPoint<DivideFunction>({prefix + "divide"});
   registerFunction<
       IntervalDivideFunction,
       IntervalDayTime,
       IntervalDayTime,
       double>({prefix + "divide"});
+  registerFunction<
+      IntervalDivideFunction,
+      IntervalYearMonth,
+      IntervalYearMonth,
+      double>({prefix + "divide"});
   registerBinaryFloatingPoint<ModulusFunction>({prefix + "mod"});
 }