From 871b01fb1d84bd8254182e0bd27657f2d36af3ab Mon Sep 17 00:00:00 2001
From: rui-mo <rui.mo@intel.com>
Date: Mon, 11 Sep 2023 10:05:27 +0800
Subject: [PATCH] Fix

---
 velox/common/memory/SharedArbitrator.cpp                 | 9 +++++----
 velox/functions/sparksql/DecimalVectorFunctions.cpp      | 3 ++-
 velox/functions/sparksql/aggregates/AverageAggregate.cpp | 3 ++-
 .../functions/sparksql/aggregates/DecimalSumAggregate.h  | 2 +-
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/velox/common/memory/SharedArbitrator.cpp b/velox/common/memory/SharedArbitrator.cpp
index f807b18e6348..f2c5088edd8b 100644
--- a/velox/common/memory/SharedArbitrator.cpp
+++ b/velox/common/memory/SharedArbitrator.cpp
@@ -424,10 +424,11 @@ uint64_t SharedArbitrator::reclaim(
     const uint64_t newCapacity = pool->capacity();
     VELOX_CHECK_GE(oldCapacity, newCapacity);
     reclaimedBytes = oldCapacity - newCapacity;
-  numReclaimedBytes_ += reclaimedBytes - freedBytes;
-  numShrunkBytes_ += freedBytes;
-  reclaimTimeUs_ += reclaimDurationUs;
-  return reclaimedBytes;
+    numReclaimedBytes_ += reclaimedBytes - freedBytes;
+    numShrunkBytes_ += freedBytes;
+    reclaimTimeUs_ += reclaimDurationUs;
+    return reclaimedBytes;
+  }
 }
 
 void SharedArbitrator::abort(
diff --git a/velox/functions/sparksql/DecimalVectorFunctions.cpp b/velox/functions/sparksql/DecimalVectorFunctions.cpp
index 708a05a450ef..4c5f12581fea 100644
--- a/velox/functions/sparksql/DecimalVectorFunctions.cpp
+++ b/velox/functions/sparksql/DecimalVectorFunctions.cpp
@@ -235,8 +235,9 @@ class DecimalRoundFunction : public exec::VectorFunction {
       uint8_t rScale,
       int32_t scale) const {
     if (scale >= 0) {
+      bool isOverflow = false;
       auto rescaledValue = DecimalUtil::rescaleWithRoundUp<A, R>(
-          a, aPrecision, aScale, rPrecision, rScale);
+          a, aPrecision, aScale, rPrecision, rScale, isOverflow);
       VELOX_DCHECK(rescaledValue.has_value());
       r = rescaledValue.value();
     } else {
diff --git a/velox/functions/sparksql/aggregates/AverageAggregate.cpp b/velox/functions/sparksql/aggregates/AverageAggregate.cpp
index 9f24b7e42bd7..1c8518aa3538 100644
--- a/velox/functions/sparksql/aggregates/AverageAggregate.cpp
+++ b/velox/functions/sparksql/aggregates/AverageAggregate.cpp
@@ -280,8 +280,9 @@ class DecimalAverageAggregate : public DecimalAggregate<TInputType> {
 
     DecimalUtil::divideWithRoundUp<int128_t, int128_t, int128_t>(
         avg, validSum.value(), countDecimal, false, sumRescale, 0);
+    bool isOverflow = false;
     return DecimalUtil::rescaleWithRoundUp<int128_t, TResultType>(
-        avg, avgPrecision, avgScale, resultPrecision, resultScale);
+        avg, avgPrecision, avgScale, resultPrecision, resultScale, isOverflow);
   }
 
  private:
diff --git a/velox/functions/sparksql/aggregates/DecimalSumAggregate.h b/velox/functions/sparksql/aggregates/DecimalSumAggregate.h
index d53c9d9fce57..a53211fd4ff7 100644
--- a/velox/functions/sparksql/aggregates/DecimalSumAggregate.h
+++ b/velox/functions/sparksql/aggregates/DecimalSumAggregate.h
@@ -74,7 +74,7 @@ class DecimalSumAggregate : public exec::Aggregate {
 
     auto [resultPrecision, resultScale] =
         getDecimalPrecisionScale(*sumType_.get());
-    overflow = !DecimalUtil::valueInRangeWithPrecision(sum, resultPrecision);
+    overflow = !DecimalUtil::valueInPrecisionRange(sum, resultPrecision);
     return sum;
   }