Fix Spark array_union on NaN (7086)

velox/functions/sparksql/ArrayFunctions.h

@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include "velox/expression/ComplexViewTypes.h"
+#include "velox/functions/Udf.h"
+#include "velox/functions/lib/CheckedArithmetic.h"
+#include "velox/type/Conversions.h"
+
+namespace facebook::velox::functions::sparksql {
+
+/// This class implements the array union function.
+///
+/// DEFINITION:
+/// array_union(x, y) → array
+/// Returns an array of the elements in the union of x and y, without
+/// duplicates.
+template <typename T>
+struct ArrayUnionFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T)
+
+  // Fast path for primitives.
+  template <typename Out, typename In>
+  void call(Out& out, const In& inputArray1, const In& inputArray2) {
+    folly::F14FastSet<typename In::element_t> elementSet;
+    bool nullAdded = false;
+    bool nanAdded = false;
+    auto addItems = [&](auto& inputArray) {
+      for (const auto& item : inputArray) {
+        if (item.has_value()) {
+          if constexpr (
+              std::is_same_v<In, arg_type<Array<float>>> ||
+              std::is_same_v<In, arg_type<Array<double>>>) {
+            bool isNaN = std::isnan(item.value());
+            if ((!nanAdded || !isNaN) &&
+                elementSet.insert(item.value()).second) {
+              auto& newItem = out.add_item();
+              newItem = item.value();
+            }
+            if (!nanAdded && isNaN) {
+              nanAdded = true;
+            }
+          } else if (elementSet.insert(item.value()).second) {
+            auto& newItem = out.add_item();
+            newItem = item.value();
+          }
+        } else if (!nullAdded) {
+          nullAdded = true;
+          out.add_null();
+        }
+      }
+    };
+    addItems(inputArray1);
+    addItems(inputArray2);
+  }
+};
+} // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/Register.cpp

@@ -24,6 +24,7 @@
 #include "velox/functions/prestosql/JsonFunctions.h"
 #include "velox/functions/prestosql/StringFunctions.h"
 #include "velox/functions/sparksql/ArrayMinMaxFunction.h"
+#include "velox/functions/sparksql/ArrayFunctions.h"
 #include "velox/functions/sparksql/ArraySort.h"
 #include "velox/functions/sparksql/Bitwise.h"
 #include "velox/functions/sparksql/DateTimeFunctions.h"
@@ -117,6 +118,12 @@ void registerExpressionGeneralFunctions(const std::string& prefix) {
 }
 } // namespace
 
+template <typename T>
+inline void registerArrayUnionFunctions(const std::string& prefix) {
+  registerFunction<sparksql::ArrayUnionFunction, Array<T>, Array<T>, Array<T>>(
+      {prefix + "array_union"});
+}
+
 void registerFunctions(const std::string& prefix) {
   registerAllSpecialFormGeneralFunctions();
 
@@ -284,6 +291,18 @@ void registerFunctions(const std::string& prefix) {
   registerArrayMinMaxFunctions(prefix);
   // Register expression general functions.
   registerExpressionGeneralFunctions(prefix);
+
+  registerArrayUnionFunctions<int8_t>(prefix);
+  registerArrayUnionFunctions<int16_t>(prefix);
+  registerArrayUnionFunctions<int32_t>(prefix);
+  registerArrayUnionFunctions<int64_t>(prefix);
+  registerArrayUnionFunctions<int128_t>(prefix);
+  registerArrayUnionFunctions<float>(prefix);
+  registerArrayUnionFunctions<double>(prefix);
+  registerArrayUnionFunctions<bool>(prefix);
+  registerArrayUnionFunctions<Timestamp>(prefix);
+  registerArrayUnionFunctions<Date>(prefix);
+  registerArrayUnionFunctions<Varbinary>(prefix);
 }
 
 } // namespace sparksql