Fix array_normalize to support int types.

Summary: Presto java supports array normalize with int type (https://github.com/prestodb/presto/blob/832b071b3c74f65a6fe8386dc296b61f17192a7a/presto-main/src/main/java/com/facebook/presto/operator/scalar/ArrayNormalizeFunction.java#L56-L64). To support similar caess, registering array normalize with int types and extended tests. Differential Revision: D63931080
amitkdutta · Oct 5, 2024 · 61a0955 · 61a0955
1 parent ddf16c8
commit 61a0955
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
diff --git a/velox/functions/prestosql/registration/ArrayFunctionsRegistration.cpp b/velox/functions/prestosql/registration/ArrayFunctionsRegistration.cpp
@@ -308,6 +308,11 @@ void registerArrayFunctions(const std::string& prefix) {
   registerArrayFrequencyFunctions<Date>(prefix);
   registerArrayFrequencyFunctions<Varchar>(prefix);
 
+  registerArrayNormalizeFunctions<int8_t>(prefix);
+  registerArrayNormalizeFunctions<int16_t>(prefix);
+  registerArrayNormalizeFunctions<int32_t>(prefix);
+  registerArrayNormalizeFunctions<int64_t>(prefix);
+  registerArrayNormalizeFunctions<int128_t>(prefix);
   registerArrayNormalizeFunctions<float>(prefix);
   registerArrayNormalizeFunctions<double>(prefix);
 }

diff --git a/velox/functions/prestosql/tests/ArrayNormalizeTest.cpp b/velox/functions/prestosql/tests/ArrayNormalizeTest.cpp
@@ -55,24 +55,27 @@ class ArrayNormalizeTest : public FunctionBaseTest {
 
   template <typename T>
   void testArrayWithElementsEqualZero() {
-    auto input = makeArrayVector<T>({{0.0, -0.0, 0.0}});
-    auto p = makeConstant<T>(2.0, input->size());
+    auto input = makeArrayVector<T>(
+        {{static_cast<T>(0.0), static_cast<T>(-0.0), static_cast<T>(0.0)}});
+    auto p = makeConstant<T>(static_cast<T>(2.0), input->size());
     testExpr(input, {input, p});
   }
 
   template <typename T>
   void testArrayWithPLessThanZero() {
-    auto input = makeArrayVector<T>({{1.0, 2.0, 3.0}});
-    auto p = makeConstant<T>(-1.0, input->size());
+    auto input = makeArrayVector<T>(
+        {{static_cast<T>(1.0), static_cast<T>(2.0), static_cast<T>(3.0)}});
+    auto p = makeConstant<T>(static_cast<T>(-1.0), input->size());
     VELOX_ASSERT_THROW(
         testExpr(input, {input, p}),
         "array_normalize only supports non-negative p");
   }
 
   template <typename T>
   void testArrayWithPEqualZero() {
-    auto vector = makeArrayVector<T>({{1.0, -2.0, 3.0}});
-    auto p = makeConstant<T>(0.0, vector->size());
+    auto vector = makeArrayVector<T>(
+        {{static_cast<T>(1.0), static_cast<T>(-2.0), static_cast<T>(3.0)}});
+    auto p = makeConstant<T>(static_cast<T>(0.0), vector->size());
     testExpr(vector, {vector, p});
   }
 
@@ -159,16 +162,19 @@ class ArrayNormalizeTest : public FunctionBaseTest {
 };
 
 TEST_F(ArrayNormalizeTest, arrayWithElementsZero) {
+  testArrayWithElementsEqualZero<int>();
   testArrayWithElementsEqualZero<float>();
   testArrayWithElementsEqualZero<double>();
 }
 
 TEST_F(ArrayNormalizeTest, pLessThanZero) {
+  testArrayWithElementsEqualZero<int>();
   testArrayWithPLessThanZero<float>();
   testArrayWithPLessThanZero<double>();
 }
 
 TEST_F(ArrayNormalizeTest, pEqualsZero) {
+  testArrayWithPEqualZero<int>();
   testArrayWithPEqualZero<float>();
   testArrayWithPEqualZero<double>();
 }
@@ -179,21 +185,25 @@ TEST_F(ArrayNormalizeTest, pLessThanOne) {
 }
 
 TEST_F(ArrayNormalizeTest, pEqualsOne) {
+  testArrayWithPEqualOne<int>();
   testArrayWithPEqualOne<float>();
   testArrayWithPEqualOne<double>();
 }
 
 TEST_F(ArrayNormalizeTest, limits) {
+  testArrayWithPEqualOne<int>();
   testArrayWithTypeLimits<float>();
   testArrayWithTypeLimits<double>();
 }
 
 TEST_F(ArrayNormalizeTest, nullValues) {
+  testArrayWithPEqualOne<int>();
   testArrayWithNullValues<float>();
   testArrayWithNullValues<double>();
 }
 
 TEST_F(ArrayNormalizeTest, differentValues) {
+  testArrayWithPEqualOne<int>();
   testArrayWithDifferentValues<float>();
   testArrayWithDifferentValues<double>();
 }