Match array min/max behavior with presto. (facebookincubator#7128)

Summary: Based on prestodb/presto#21065 Reviewed By: mbasmanova Differential Revision: D50386287
amitkdutta · Oct 21, 2023 · 103a919 · 103a919
1 parent d87866b
commit 103a919
Show file tree

Hide file tree

Showing 4 changed files with 106 additions and 9 deletions.
diff --git a/velox/docs/functions/presto/array.rst b/velox/docs/functions/presto/array.rst
@@ -93,21 +93,31 @@ Array Functions
 
 .. function:: array_max(array(E)) -> E
 
-    Returns the maximum value of input array. ::
+    Returns the maximum value of input array.
+    Returns NaN if E is REAL or DOUBLE and array contains a NaN value.
+    Returns NULL if array doesn't contain a NaN value, but contains a NULL value. ::
 
         SELECT array_max(ARRAY [1, 2, 3]); -- 3
         SELECT array_max(ARRAY [-1, -2, -2]); -- -1
         SELECT array_max(ARRAY [-1, -2, NULL]); -- NULL
         SELECT array_max(ARRAY []); -- NULL
+        SELECT array_max(ARRAY[NULL, nan()]); -- NaN
+        SELECT array_max(ARRAY[{-1, -2, -3, nan()]); -- NaN
+        SELECT array_max(ARRAY[-0.0001, NULL, -0.0003, nan()]); -- NaN
 
 .. function:: array_min(array(E)) -> E
 
-    Returns the minimum value of input array. ::
+    Returns the minimum value of input array.
+    Returns NaN if E is REAL or DOUBLE and array contains a NaN value.
+    Returns NULL if array doesn't contain a NaN value, but contains a NULL value. ::
 
         SELECT array_min(ARRAY [1, 2, 3]); -- 1
         SELECT array_min(ARRAY [-1, -2, -2]); -- -2
         SELECT array_min(ARRAY [-1, -2, NULL]); -- NULL
         SELECT array_min(ARRAY []); -- NULL
+        SELECT array_min(ARRAY[NULL, nan()]); -- NaN
+        SELECT array_min(ARRAY[{-1, -2, -3, nan()]); -- NaN
+        SELECT array_min(ARRAY[-0.0001, NULL, -0.0003, nan()]); -- NaN
 
 .. function:: array_normalize(array(E), E) -> array(E)
 

diff --git a/velox/functions/prestosql/ArrayFunctions.h b/velox/functions/prestosql/ArrayFunctions.h
@@ -52,13 +52,71 @@ struct ArrayMinMaxFunction {
     }
   }
 
+  template <typename TReturn, typename TInput>
+  bool callForFloatOrDouble(TReturn& out, const TInput& array) {
+    bool hasNull = false;
+    auto it = array.begin();
+
+    // Find the first non-null item (if any)
+    while (it != array.end()) {
+      if (it->has_value()) {
+        break;
+      }
+
+      hasNull = true;
+      ++it;
+    }
+
+    // Return false if end of array is reached without finding a non-null item.
+    if (it == array.end()) {
+      return false;
+    }
+
+    // If first non-null item is NAN, return immediately.
+    auto currentValue = it->value();
+    if (std::isnan(currentValue)) {
+      assign(out, currentValue);
+      return true;
+    }
+
+    ++it;
+    while (it != array.end()) {
+      if (it->has_value()) {
+        auto newValue = it->value();
+        if (std::isnan(newValue)) {
+          assign(out, newValue);
+          return true;
+        }
+        update(currentValue, newValue);
+      } else {
+        hasNull = true;
+      }
+      ++it;
+    }
+
+    // If we found a null, return false. Note that, if we found
+    // a NAN, the function will return at earlier stage as soon as
+    // a NAN is observed.
+    if (hasNull) {
+      return false;
+    }
+
+    assign(out, currentValue);
+    return true;
+  }
+
   template <typename TReturn, typename TInput>
   FOLLY_ALWAYS_INLINE bool call(TReturn& out, const TInput& array) {
     // Result is null if array is empty.
     if (array.size() == 0) {
       return false;
     }
 
+    if constexpr (
+        std::is_same_v<TReturn, float> || std::is_same_v<TReturn, double>) {
+      return callForFloatOrDouble(out, array);
+    }
+
     if (!array.mayHaveNulls()) {
       // Input array does not have nulls.
       auto currentValue = *array[0];

diff --git a/velox/functions/prestosql/tests/ArrayMaxTest.cpp b/velox/functions/prestosql/tests/ArrayMaxTest.cpp
@@ -145,11 +145,24 @@ TEST_F(ArrayMaxTest, longVarcharNoNulls) {
 
 // Test documented example.
 TEST_F(ArrayMaxTest, docs) {
-  auto input = makeNullableArrayVector<int32_t>(
+  auto input1 = makeNullableArrayVector<int32_t>(
       {{1, 2, 3}, {-1, -2, -2}, {-1, -2, std::nullopt}, {}});
-  auto expected =
+  auto expected1 =
       makeNullableFlatVector<int32_t>({3, -1, std::nullopt, std::nullopt});
-  testArrayMax(input, expected);
+  testArrayMax(input1, expected1);
+
+  auto input2 = makeNullableArrayVector<double>(
+      {{std::nullopt, std::numeric_limits<double>::quiet_NaN()},
+       {-1, -2, -3, std::numeric_limits<double>::quiet_NaN()},
+       {-0.0001,
+        std::nullopt,
+        -0.0003,
+        std::numeric_limits<double>::quiet_NaN()}});
+  auto expected2 = makeNullableFlatVector<double>(
+      {std::numeric_limits<double>::quiet_NaN(),
+       std::numeric_limits<double>::quiet_NaN(),
+       std::numeric_limits<double>::quiet_NaN()});
+  testArrayMax(input2, expected2);
 }
 
 template <typename Type>
@@ -238,6 +251,7 @@ class ArrayMaxFloatingPointTest : public FunctionBaseTest {
           -0.0003,
           std::numeric_limits<T>::max()},
          {},
+         {std::nullopt},
          {std::numeric_limits<T>::min(), 1.1, 1.22222, 1.33, std::nullopt},
          {-0.00001, -0.0002, 0.0001}});
 
@@ -247,6 +261,7 @@ class ArrayMaxFloatingPointTest : public FunctionBaseTest {
          std::numeric_limits<T>::max(),
          std::nullopt,
          std::nullopt,
+         std::nullopt,
          0.0001});
     testArrayMax(input, expected);
   }

diff --git a/velox/functions/prestosql/tests/ArrayMinTest.cpp b/velox/functions/prestosql/tests/ArrayMinTest.cpp
@@ -36,11 +36,25 @@ class ArrayMinTest : public FunctionBaseTest {
   }
 
   void testDocExample() {
-    auto arrayVector = makeNullableArrayVector<int64_t>(
+    auto input1 = makeNullableArrayVector<int64_t>(
         {{1, 2, 3}, {-1, -2, -2}, {-1, -2, std::nullopt}, {}});
-    auto expected =
+    auto expected1 =
         makeNullableFlatVector<int64_t>({1, -2, std::nullopt, std::nullopt});
-    testExpr<int64_t>(expected, "array_min(C0)", {arrayVector});
+    testExpr<int64_t>(expected1, "array_min(C0)", {input1});
+
+    auto input2 = makeNullableArrayVector<float>(
+        {{std::nullopt, std::numeric_limits<double>::quiet_NaN()},
+         {-1, -2, -3, std::numeric_limits<float>::quiet_NaN()},
+         {-0.0001,
+          std::nullopt,
+          -0.0003,
+          std::numeric_limits<float>::quiet_NaN()}});
+    auto expected2 = makeNullableFlatVector<float>(
+        {std::numeric_limits<float>::quiet_NaN(),
+         std::numeric_limits<float>::quiet_NaN(),
+         std::numeric_limits<float>::quiet_NaN()});
+
+    testExpr<float>(expected2, "array_min(C0)", {input2});
   }
 
   template <typename T>
@@ -145,7 +159,7 @@ class ArrayMinTest : public FunctionBaseTest {
 };
 
 } // namespace
-TEST_F(ArrayMinTest, docArrays) {
+TEST_F(ArrayMinTest, docs) {
   testDocExample();
 }