From 4f1140247e2f5059abd8f71b2ae54e3d1dcf1be2 Mon Sep 17 00:00:00 2001
From: Amit Dutta <adutta@meta.com>
Date: Sun, 6 Oct 2024 10:05:39 -0700
Subject: [PATCH] Fix array_normalize to support int types. (#11173)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/11173

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.

Reviewed By: xiaoxmeng, kgpai

Differential Revision: D63931080

fbshipit-source-id: e58f1e0630b0a493d81ae18d16f9080e392e16c3
---
 .../ArrayFunctionsRegistration.cpp            |  5 +++++
 .../prestosql/tests/ArrayNormalizeTest.cpp    | 22 ++++++++++++++-----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/velox/functions/prestosql/registration/ArrayFunctionsRegistration.cpp b/velox/functions/prestosql/registration/ArrayFunctionsRegistration.cpp
index 506e047ba8b1..f3183f4d1577 100644
--- 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
index 5f40f84010f3..5493df0cc1eb 100644
--- a/velox/functions/prestosql/tests/ArrayNormalizeTest.cpp
+++ b/velox/functions/prestosql/tests/ArrayNormalizeTest.cpp
@@ -55,15 +55,17 @@ 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");
@@ -71,8 +73,9 @@ class ArrayNormalizeTest : public FunctionBaseTest {
 
   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>();
 }