Support UNKNOWN type in Map function (#9410)

Summary: Currently the key type of `MAP` is specified as `knownTypeVariable` which refuses to bind `UNKNOWN` type, while in Presto `MAP(UNKNOWN, UNKNOWN)` is valid. Differential Revision: D55902535
facebookincubator · Apr 9, 2024 · 7fd7828 · 7fd7828
1 parent 3558eb7
commit 7fd7828
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/velox/functions/prestosql/Map.cpp b/velox/functions/prestosql/Map.cpp
@@ -278,7 +278,7 @@ class MapFunction : public exec::VectorFunction {
   static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
     // array(K), array(V) -> map(K,V)
     return {exec::FunctionSignatureBuilder()
-                .knownTypeVariable("K")
+                .typeVariable("K")
                 .typeVariable("V")
                 .returnType("map(K,V)")
                 .argumentType("array(K)")

diff --git a/velox/functions/prestosql/tests/MapTest.cpp b/velox/functions/prestosql/tests/MapTest.cpp
@@ -420,4 +420,21 @@ TEST_F(MapTest, nestedNullInKeys) {
       "map key cannot be indeterminate");
 }
 
+TEST_F(MapTest, unknownType) {
+  // MAP(ARRAY[], ARRAY[])
+  auto emptyArrayVector = makeArrayVector<UnknownValue>({{}});
+  auto expectedMap = makeMapVector<UnknownValue, UnknownValue>({{}});
+  auto result = evaluate(
+      "map(c0, c1)", makeRowVector({emptyArrayVector, emptyArrayVector}));
+  assertEqualVectors(expectedMap, result);
+
+  // MAP(ARRAY[null], ARRAY[null])
+  auto elementVector = makeNullableFlatVector<UnknownValue>({std::nullopt});
+  auto nullArrayVector = makeArrayVector({0}, elementVector);
+  VELOX_ASSERT_THROW(
+      evaluate(
+          "map(c0, c1)", makeRowVector({nullArrayVector, nullArrayVector})),
+      "map key cannot be null");
+}
+
 } // namespace