sparksql decimal avg add some check (facebookincubator#14)

rui-mo · Mar 22, 2023 · a91568a · a91568a
1 parent 9ddcf41
commit a91568a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/velox/functions/sparksql/aggregates/DecimalAvgAggregate.h b/velox/functions/sparksql/aggregates/DecimalAvgAggregate.h
@@ -319,7 +319,9 @@ class DecimalAverageAggregate : public exec::Aggregate {
 
         // Spark use DECIMAL(20, 0) to represent long value, but we need to
         // create a SHORT_DECIMAL to get SHORT_DECIMAL result
+        DCHECK(DecimalUtil::numDigits(accumulator->count) <= 18);
         auto countDecimal = UnscaledShortDecimal(accumulator->count);
+        DCHECK(accumulator->sum <= LONG_MAX);
         auto longUnscaledSum = (int64_t)accumulator->sum;
         DecimalUtil::divideWithRoundUp<
             UnscaledShortDecimal,
@@ -338,6 +340,7 @@ class DecimalAverageAggregate : public exec::Aggregate {
       if constexpr (std::is_same_v<TResultType, UnscaledLongDecimal>) {
         // Spark use DECIMAL(20, 0) to represent long value
         auto countDecimal = UnscaledLongDecimal(accumulator->count);
+        DCHECK(accumulator->sum <= LONG_MAX);
         auto longUnscaledSum = (int64_t)accumulator->sum;
         DecimalUtil::divideWithRoundUp<
             UnscaledLongDecimal,
@@ -355,7 +358,9 @@ class DecimalAverageAggregate : public exec::Aggregate {
 
         // Spark use DECIMAL(20, 0) to represent long value, but we need to
         // create a SHORT_DECIMAL to get SHORT_DECIMAL result
+        DCHECK(DecimalUtil::numDigits(accumulator->count) <= 18);
         auto countDecimal = UnscaledShortDecimal(accumulator->count);
+        DCHECK(accumulator->sum <= LONG_MAX);
         auto longUnscaledSum = (int64_t)accumulator->sum;
         DecimalUtil::divideWithRoundUp<
             UnscaledShortDecimal,

diff --git a/velox/type/DecimalUtil.h b/velox/type/DecimalUtil.h
@@ -454,6 +454,18 @@ class DecimalUtil {
     return (unscaled * sig) / DecimalUtil::kPowersOfTen[scale];
   }
 
+  template <class T>
+  inline static int numDigits(T number)
+  {
+    int digits = 0;
+    if (number < 0) digits = 1; // remove this line if '-' counts as a digit
+    while (number) {
+      number /= 10;
+      digits++;
+    }
+    return digits;
+  }
+
   static constexpr double double10pow[] = {
       1.0e0,  1.0e1,  1.0e2,  1.0e3,  1.0e4,  1.0e5,  1.0e6,  1.0e7,
       1.0e8,  1.0e9,  1.0e10, 1.0e11, 1.0e12, 1.0e13, 1.0e14, 1.0e15,