From 4b6348d049907a8308f4a083369c873e392f7830 Mon Sep 17 00:00:00 2001
From: Sergey Pershin <spershin@meta.com>
Date: Fri, 17 Jan 2025 15:52:35 -0800
Subject: [PATCH] feat: Add MAP() signature for MapFunction.

Summary:
Some queries can fail with the following error:
Scalar function presto.default.map not registered with arguments: ()

MAP() is supported in Presto and works by creating an emtpy map of the output type.
Add the same capability to Velox.

Differential Revision: D68359107
---
 velox/functions/prestosql/Map.cpp           | 28 +++++++++++++++------
 velox/functions/prestosql/tests/MapTest.cpp | 11 ++++++++
 velox/vector/ComplexVector.cpp              | 19 ++++++++++++++
 velox/vector/ComplexVector.h                |  6 +++++
 4 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/velox/functions/prestosql/Map.cpp b/velox/functions/prestosql/Map.cpp
index fa24b4deb2eb8..05b40ad86cf51 100644
--- a/velox/functions/prestosql/Map.cpp
+++ b/velox/functions/prestosql/Map.cpp
@@ -32,6 +32,14 @@ class MapFunction : public exec::VectorFunction {
       const TypePtr& outputType,
       exec::EvalCtx& context,
       VectorPtr& result) const override {
+    // No arguments case (empty map).
+    if (args.empty()) {
+      auto emptyMapVector =
+          std::make_shared<MapVector>(context.pool(), outputType, rows);
+      context.moveOrCopyResult(emptyMapVector, rows, result);
+      return;
+    }
+
     VELOX_CHECK_EQ(args.size(), 2);
 
     auto keys = args[0];
@@ -277,13 +285,19 @@ class MapFunction : public exec::VectorFunction {
 
   static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
     // array(K), array(V) -> map(K,V)
-    return {exec::FunctionSignatureBuilder()
-                .typeVariable("K")
-                .typeVariable("V")
-                .returnType("map(K,V)")
-                .argumentType("array(K)")
-                .argumentType("array(V)")
-                .build()};
+    // () -> map()
+    return {
+        exec::FunctionSignatureBuilder()
+            .typeVariable("K")
+            .typeVariable("V")
+            .returnType("map(K,V)")
+            .argumentType("array(K)")
+            .argumentType("array(V)")
+            .build(),
+        exec::FunctionSignatureBuilder()
+            .returnType("map(unknown,unknown)")
+            .build(),
+    };
   }
 
  private:
diff --git a/velox/functions/prestosql/tests/MapTest.cpp b/velox/functions/prestosql/tests/MapTest.cpp
index 218f9e5877c99..7844e60bf2102 100644
--- a/velox/functions/prestosql/tests/MapTest.cpp
+++ b/velox/functions/prestosql/tests/MapTest.cpp
@@ -97,6 +97,17 @@ TEST_F(MapTest, noNulls) {
 
   auto result = evaluate("map(c0, c1)", makeRowVector({keys, values}));
   assertEqualVectors(expectedMap, result);
+
+  // Empty map constructor. We use unknown types here because no way to specify
+  // the output type for the map() call.
+  const TypePtr& emptyMapType =
+      MAP(CppToType<UnknownValue>::create(), CppToType<UnknownValue>::create());
+  SelectivityVector rows(size);
+  auto emptyMapVector =
+      std::make_shared<MapVector>(pool_.get(), emptyMapType, rows);
+
+  result = evaluate("map()", makeRowVector({keys, values}));
+  assertEqualVectors(emptyMapVector, result);
 }
 
 TEST_F(MapTest, someNulls) {
diff --git a/velox/vector/ComplexVector.cpp b/velox/vector/ComplexVector.cpp
index bc2295325fdcf..038f7a18ebf0b 100644
--- a/velox/vector/ComplexVector.cpp
+++ b/velox/vector/ComplexVector.cpp
@@ -1201,6 +1201,25 @@ void ArrayVector::copyRanges(
   copyRangesImpl(source, ranges, &elements_, nullptr);
 }
 
+MapVector::MapVector(
+    velox::memory::MemoryPool* pool,
+    std::shared_ptr<const Type> type,
+    const SelectivityVector& rows)
+    : ArrayVectorBase(
+          pool,
+          type,
+          VectorEncoding::Simple::MAP,
+          nullptr, // nulls
+          rows.end(),
+          0, // nullCount
+          allocateOffsets(rows.end(), pool),
+          allocateSizes(rows.end(), pool)),
+      keys_(BaseVector::create(type->childAt(0), 0, pool)),
+      values_(BaseVector::create(type->childAt(1), 0, pool)),
+      sortedKeys_{false} {
+  VELOX_CHECK_EQ(type->kind(), TypeKind::MAP);
+}
+
 bool MapVector::containsNullAt(vector_size_t idx) const {
   if (BaseVector::isNullAt(idx)) {
     return true;
diff --git a/velox/vector/ComplexVector.h b/velox/vector/ComplexVector.h
index e4fa2a3ccebc7..26a42bb513c11 100644
--- a/velox/vector/ComplexVector.h
+++ b/velox/vector/ComplexVector.h
@@ -591,6 +591,12 @@ class MapVector : public ArrayVectorBase {
         type->childAt(1)->toString());
   }
 
+  /// Constructor for creating an empty MapVector for the given rows.
+  MapVector(
+      velox::memory::MemoryPool* pool,
+      std::shared_ptr<const Type> type,
+      const SelectivityVector& rows);
+
   virtual ~MapVector() override {}
 
   bool containsNullAt(vector_size_t idx) const override;