Adding long decimal support for arbitrary() function (#9217)

Summary: Delivers #7311 Pull Request resolved: #9217 Reviewed By: amitkdutta Differential Revision: D56662231 Pulled By: kagamiori fbshipit-source-id: 630f8ecc4b65eb6baa7ec8a6f430640d15d4835d
facebookincubator · Apr 29, 2024 · 84acb4c · 84acb4c
1 parent 37f4700
commit 84acb4c
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 0 deletions.
diff --git a/velox/functions/prestosql/aggregates/ArbitraryAggregate.cpp b/velox/functions/prestosql/aggregates/ArbitraryAggregate.cpp
@@ -39,6 +39,10 @@ class ArbitraryAggregate : public SimpleNumericAggregate<T, T, T> {
     return sizeof(T);
   }
 
+  int32_t accumulatorAlignmentSize() const override {
+    return 1;
+  }
+
   void extractValues(char** groups, int32_t numGroups, VectorPtr* result)
       override {
     BaseAggregate::doExtractValues(groups, numGroups, result, [&](char* group) {
@@ -139,6 +143,14 @@ class ArbitraryAggregate : public SimpleNumericAggregate<T, T, T> {
   }
 };
 
+/// Override 'accumulatorAlignmentSize' for UnscaledLongDecimal values as it
+/// uses int128_t type. Some CPUs don't support misaligned access to int128_t
+/// type.
+template <>
+inline int32_t ArbitraryAggregate<int128_t>::accumulatorAlignmentSize() const {
+  return static_cast<int32_t>(sizeof(int128_t));
+}
+
 // Arbitrary for non-numeric types. We always keep the first (non-NULL) element
 // seen. Arbitrary (x) will produce partial and final aggregations of type x.
 class NonNumericArbitrary : public exec::Aggregate {
@@ -301,6 +313,11 @@ void registerArbitraryAggregate(
             return std::make_unique<ArbitraryAggregate<int32_t>>(inputType);
           case TypeKind::BIGINT:
             return std::make_unique<ArbitraryAggregate<int64_t>>(inputType);
+          case TypeKind::HUGEINT:
+            if (inputType->isLongDecimal()) {
+              return std::make_unique<ArbitraryAggregate<int128_t>>(inputType);
+            }
+            VELOX_NYI();
           case TypeKind::REAL:
             return std::make_unique<ArbitraryAggregate<float>>(inputType);
           case TypeKind::DOUBLE:

diff --git a/velox/functions/prestosql/aggregates/tests/ArbitraryTest.cpp b/velox/functions/prestosql/aggregates/tests/ArbitraryTest.cpp
@@ -351,6 +351,60 @@ TEST_F(ArbitraryTest, interval) {
   testAggregations({data}, {}, {"arbitrary(c2)"}, "SELECT null");
 }
 
+TEST_F(ArbitraryTest, longDecimal) {
+  auto data = makeRowVector({// Grouping key.
+                             makeFlatVector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
+                             makeNullableFlatVector<int128_t>(
+                                 {HugeInt::build(10, 100),
+                                  HugeInt::build(10, 100),
+                                  HugeInt::build(10, 200),
+                                  HugeInt::build(10, 200),
+                                  std::nullopt,
+                                  std::nullopt,
+                                  std::nullopt,
+                                  HugeInt::build(10, 400)},
+                                 DECIMAL(38, 8))});
+
+  auto expectedResult = makeRowVector({
+      makeFlatVector<int64_t>({1, 2, 3, 4}),
+      makeNullableFlatVector<int128_t>(
+          {HugeInt::build(10, 100),
+           HugeInt::build(10, 200),
+           std::nullopt,
+           HugeInt::build(10, 400)},
+          DECIMAL(38, 8)),
+  });
+
+  testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});
+}
+
+TEST_F(ArbitraryTest, shortDecimal) {
+  auto data = makeRowVector({// Grouping key.
+                             makeFlatVector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
+                             makeNullableFlatVector<int64_t>(
+                                 {10000000000000000,
+                                  10000000000000000,
+                                  20000000000000000,
+                                  20000000000000000,
+                                  std::nullopt,
+                                  std::nullopt,
+                                  std::nullopt,
+                                  40000000000000000},
+                                 DECIMAL(15, 2))});
+
+  auto expectedResult = makeRowVector({
+      makeFlatVector<int64_t>({1, 2, 3, 4}),
+      makeNullableFlatVector<int64_t>(
+          {10000000000000000,
+           20000000000000000,
+           std::nullopt,
+           40000000000000000},
+          DECIMAL(15, 2)),
+  });
+
+  testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});
+}
+
 class ArbitraryWindowTest : public WindowTestBase {};
 
 TEST_F(ArbitraryWindowTest, basic) {