From 47a9328574522451239aa12a8067cfa7b7054e96 Mon Sep 17 00:00:00 2001
From: Amit Dutta <adutta@meta.com>
Date: Tue, 15 Oct 2024 16:09:59 -0700
Subject: [PATCH] Throw array_normalize when registered with int types to match
 Presto java behavior. (#11263)

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

Reviewed By: kevinwilfong

Differential Revision: D64308532

fbshipit-source-id: 2bd3fcd3cc16c58527ba6619f099558d2136a9da
---
 velox/functions/prestosql/ArrayFunctions.h        | 15 +++++++++++++++
 .../prestosql/tests/ArrayNormalizeTest.cpp        | 15 ++++++++-------
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/velox/functions/prestosql/ArrayFunctions.h b/velox/functions/prestosql/ArrayFunctions.h
index adc96d1b7025..61dc8757090f 100644
--- a/velox/functions/prestosql/ArrayFunctions.h
+++ b/velox/functions/prestosql/ArrayFunctions.h
@@ -593,6 +593,21 @@ struct ArrayNormalizeFunction {
     VELOX_USER_CHECK_GE(
         p, 0, "array_normalize only supports non-negative p: {}", p);
 
+    // Ideally, we should not register this function with int types. However,
+    // in Presto, during plan conversion it's possible to create this function
+    // with int types. In that case, we want to have default NULL behavior for
+    // int types, which can be achieved by registering the function and failing
+    // it for non-null int values. Example: SELECT array_normalize(x, 2) from
+    // (VALUES NULL) t(x) creates a plan with `array_normalize :=
+    // array_normalize(CAST(field AS array(integer)), INTEGER'2') (1:37)` which
+    // ideally should fail saying the function is not registered with int types.
+    // But it falls back to default NULL behavior and returns NULL in Presto.
+    if constexpr (
+        std::is_same_v<T, int8_t> || std::is_same_v<T, int16_t> ||
+        std::is_same_v<T, int32_t> || std::is_same_v<T, int64_t>) {
+      VELOX_UNSUPPORTED("array_normalize only supports double and float types");
+    }
+
     // If the input array is empty, then the empty result should be returned,
     // same as Presto.
     if (inputArray.size() == 0) {
diff --git a/velox/functions/prestosql/tests/ArrayNormalizeTest.cpp b/velox/functions/prestosql/tests/ArrayNormalizeTest.cpp
index 5493df0cc1eb..3d9b8f28fb58 100644
--- a/velox/functions/prestosql/tests/ArrayNormalizeTest.cpp
+++ b/velox/functions/prestosql/tests/ArrayNormalizeTest.cpp
@@ -162,19 +162,16 @@ 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>();
 }
@@ -185,27 +182,31 @@ 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>();
 }
 
+TEST_F(ArrayNormalizeTest, throwForIntTypes) {
+  std::string errorMessage =
+      "array_normalize only supports double and float types";
+  VELOX_ASSERT_THROW(testArrayWithElementsEqualZero<int>(), errorMessage);
+  VELOX_ASSERT_THROW(testArrayWithPEqualZero<int>(), errorMessage);
+  VELOX_ASSERT_THROW(testArrayWithPEqualOne<int>(), errorMessage);
+}
+
 } // namespace