From 2f4496516379b91d8b6d8fc8ce75c2686ebae679 Mon Sep 17 00:00:00 2001
From: Amit Dutta <adutta@meta.com>
Date: Tue, 15 Oct 2024 12:00:23 -0700
Subject: [PATCH] Throw array_normalize when registered with int types to match
 Presto java behavior.

Differential Revision: D64308532
---
 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..23a1bb3368e1 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 its 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 int types. 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