Fix NaN handling in map_subset

velox/docs/functions/presto/map.rst

@@ -87,7 +87,8 @@ Map Functions
 
 .. function:: map_subset(map(K,V), array(k)) -> map(K,V)
 
-    Constructs a map from those entries of ``map`` for which the key is in the array given::
+    Constructs a map from those entries of ``map`` for which the key is in the array given
+    For keys containing REAL and DOUBLE, NANs (Not-a-Number) are considered equal. ::
 
         SELECT map_subset(MAP(ARRAY[1,2], ARRAY['a','b']), ARRAY[10]); -- {}
         SELECT map_subset(MAP(ARRAY[1,2], ARRAY['a','b']), ARRAY[1]); -- {1->'a'}

velox/functions/prestosql/MapSubset.h

@@ -17,6 +17,7 @@
 
 #include "velox/expression/ComplexViewTypes.h"
 #include "velox/functions/Udf.h"
+#include "velox/type/FloatingPointUtil.h"
 
 namespace facebook::velox::functions {
 
@@ -84,7 +85,7 @@ struct MapSubsetPrimitiveFunction {
   }
 
   bool constantSearchKeys_{false};
-  folly::F14FastSet<arg_type<Key>> searchKeys_;
+  util::floating_point::HashSetNaNAware<arg_type<Key>> searchKeys_;
 };
 
 /// Fast path for constant string keys: map_subset(m, array['a', 'b', 'c']).

velox/functions/prestosql/tests/MapSubsetTest.cpp

@@ -21,7 +21,67 @@ using namespace facebook::velox::test;
 namespace facebook::velox::functions {
 namespace {
 
-class MapSubsetTest : public test::FunctionBaseTest {};
+class MapSubsetTest : public test::FunctionBaseTest {
+ public:
+  template <typename T>
+  void testFloatNaNs() {
+    static const auto kNaN = std::numeric_limits<T>::quiet_NaN();
+    static const auto kSNaN = std::numeric_limits<T>::signaling_NaN();
+
+    // Case 1: Non-constant search keys.
+    auto data = makeRowVector(
+        {makeMapVectorFromJson<T, int32_t>({
+             "{1:10, NaN:20, 3:null, 4:40, 5:50, 6:60}",
+             "{NaN:20}",
+         }),
+         makeArrayVector<T>({{1, kNaN, 5}, {kSNaN, 3}})});
+
+    auto expected = makeMapVectorFromJson<T, int32_t>({
+        "{1:10, NaN:20, 5:50}",
+        "{NaN:20}",
+    });
+    auto result = evaluate("map_subset(c0, c1)", data);
+    assertEqualVectors(expected, result);
+
+    // Case 2: Constant search keys.
+    data = makeRowVector(
+        {makeMapVectorFromJson<T, int32_t>({
+             "{1:10, NaN:20, 3:null, 4:40, 5:50, 6:60}",
+             "{NaN:20}",
+         }),
+         BaseVector::wrapInConstant(2, 0, makeArrayVector<T>({{1, kNaN, 5}}))});
+    expected = makeMapVectorFromJson<T, int32_t>({
+        "{1:10, NaN:20, 5:50}",
+        "{NaN:20}",
+    });
+    result = evaluate("map_subset(c0, c1)", data);
+    assertEqualVectors(expected, result);
+
+    // Case 3: Map with Complex type as key.
+    // Map: { [{1, NaN,3}: 1, {4, 5}: 2], [{NaN, 3}: 3, {1, 2}: 4] }
+    data = makeRowVector({
+        makeMapVector(
+            {0, 2},
+            makeArrayVector<T>({{1, kNaN, 3}, {4, 5}, {kSNaN, 3}, {1, 2}}),
+            makeFlatVector<int32_t>({1, 2, 3, 4})),
+        makeNestedArrayVectorFromJson<T>({
+            "[[1, NaN, 3], [4, 5]]",
+            "[[1, 2, 3], [NaN, 3]]",
+        }),
+    });
+    expected = makeMapVector(
+        {0, 2},
+        makeArrayVectorFromJson<T>({
+            "[1, NaN, 3]",
+            "[4, 5]",
+            "[NaN, 3]",
+        }),
+        makeFlatVector<int32_t>({1, 2, 3}));
+
+    result = evaluate("map_subset(c0, c1)", data);
+    assertEqualVectors(expected, result);
+  }
+};
 
 TEST_F(MapSubsetTest, bigintKey) {
   auto data = makeRowVector({
@@ -133,5 +193,10 @@ TEST_F(MapSubsetTest, arrayKey) {
   assertEqualVectors(expected, result);
 }
 
+TEST_F(MapSubsetTest, floatNaNs) {
+  testFloatNaNs<float>();
+  testFloatNaNs<double>();
+}
+
 } // namespace
 } // namespace facebook::velox::functions

velox/vector/tests/utils/VectorMaker.h

@@ -822,6 +822,7 @@ class VectorMaker {
       folly::json::serialization_opts options;
       options.convert_int_keys = true;
       options.allow_non_string_keys = true;
+      options.allow_nan_inf = true;
       folly::dynamic mapObject = folly::parseJson(jsonMap, options);
       if (mapObject.isNull()) {
         // Null map.