Skip to content

Commit

Permalink
feat: Add MAP() signature for MapFunction.
Browse files Browse the repository at this point in the history
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
  • Loading branch information
Sergey Pershin authored and facebook-github-bot committed Jan 17, 2025
1 parent dc6c6e3 commit 4b6348d
Show file tree
Hide file tree
Showing 4 changed files with 57 additions and 7 deletions.
28 changes: 21 additions & 7 deletions velox/functions/prestosql/Map.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down Expand Up @@ -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:
Expand Down
11 changes: 11 additions & 0 deletions velox/functions/prestosql/tests/MapTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
19 changes: 19 additions & 0 deletions velox/vector/ComplexVector.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
6 changes: 6 additions & 0 deletions velox/vector/ComplexVector.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down

0 comments on commit 4b6348d

Please sign in to comment.