Skip to content

Commit

Permalink
Cache the hash map in SubscriptUtil if a single map is reused with co…
Browse files Browse the repository at this point in the history 
…mplex keys (#10414)

Summary:
Pull Request resolved: #10414

Today in SubscriptUtil, if the map passed in has primitive keys and the same MapVector is passed in
multiple times, we cache the hash maps for optimized lookups across batches.  If the map has
complex keys, and base MapVector has a single value, we construct a local hash map to optimize the
lookups within a batch.

This change merges the two approaches for maps with Complex keys, so if we see the same
Vector passed in multiple times where the base MapVector has a single value, we cache the hash
map so we don't need to reconstruct it for every batch.

In some cases we've seen this can significantly speed up Presto queries, particularly because the
cost of hashing complex types to construct the hash map can be fairly high.

We could go a step further and, like for maps with primitive keys, cache maps with complex keys
regardless of the number of values in the base MapVector, but I haven't seen any cases that would
benefit from this yet so I'm not sure what the tradeoffs in terms of memory and construction cost
would look like, so extending the original optimizations seems like a good starting point.

Differential Revision: D59474490
  • Loading branch information
@facebook-github-bot
Kevin Wilfong authored and facebook-github-bot committed Aug 14, 2024
1 parent 94763e7 commit 4676df5
Show file tree
Hide file tree
Showing 5 changed files with 275 additions and 69 deletions.
77 changes: 34 additions & 43 deletions velox/functions/lib/SubscriptUtil.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,19 +58,19 @@ struct SimpleType<TypeKind::VARCHAR> {
template <TypeKind kind>
VectorPtr applyMapTyped(
bool triggerCaching,
std::shared_ptr<LookupTableBase>& cachedLookupTablePtr,
std::shared_ptr<detail::LookupTableBase>& cachedLookupTablePtr,
const SelectivityVector& rows,
const VectorPtr& mapArg,
const VectorPtr& indexArg,
exec::EvalCtx& context) {
static constexpr vector_size_t kMinCachedMapSize = 100;
using TKey = typename TypeTraits<kind>::NativeType;

LookupTable<kind>* typedLookupTable = nullptr;
detail::LookupTable<kind>* typedLookupTable = nullptr;
if (triggerCaching) {
if (!cachedLookupTablePtr) {
cachedLookupTablePtr =
std::make_shared<LookupTable<kind>>(*context.pool());
std::make_shared<detail::LookupTable<kind>>(*context.pool());
}

typedLookupTable = cachedLookupTablePtr->typedTable<kind>();
Expand Down Expand Up @@ -178,39 +178,13 @@ VectorPtr applyMapTyped(
nullsBuilder.build(), indices, rows.end(), baseMap->mapValues());
}

// A flat vector of map keys, an index into that vector and an index into
// the original map keys vector that may have encodings.
struct MapKey {
const BaseVector* baseVector;
const vector_size_t baseIndex;
const vector_size_t index;

size_t hash() const {
return baseVector->hashValueAt(baseIndex);
}

bool operator==(const MapKey& other) const {
return baseVector->equalValueAt(
other.baseVector, baseIndex, other.baseIndex);
}

bool operator<(const MapKey& other) const {
return baseVector->compare(other.baseVector, baseIndex, other.baseIndex) <
0;
}
};

struct MapKeyHasher {
size_t operator()(const MapKey& key) const {
return key.hash();
}
};

VectorPtr applyMapComplexType(
const SelectivityVector& rows,
const VectorPtr& mapArg,
const VectorPtr& indexArg,
exec::EvalCtx& context) {
exec::EvalCtx& context,
bool triggerCaching,
std::shared_ptr<detail::ComplexKeyHashMap>& cachedHashMap) {
auto* pool = context.pool();

// Use indices with the mapValues wrapped in a dictionary vector.
Expand Down Expand Up @@ -247,18 +221,32 @@ VectorPtr applyMapComplexType(
// Fast path for the case of a single map. It may be constant or dictionary
// encoded. Use hash table for quick search.
if (baseMap->size() == 1) {
folly::F14FastSet<MapKey, MapKeyHasher> set;
auto numKeys = rawSizes[0];
set.reserve(numKeys * 1.3);
for (auto i = 0; i < numKeys; ++i) {
set.insert(MapKey{mapKeysBase, mapKeysIndices[i], i});
detail::ComplexKeyHashMap hashMap{detail::MapKeyAllocator(*pool)};
detail::ComplexKeyHashMap* hashMapPtr = &hashMap;

if (triggerCaching) {
if (cachedHashMap == nullptr) {
cachedHashMap = std::make_shared<detail::ComplexKeyHashMap>(
detail::MapKeyAllocator(*pool));
}

hashMapPtr = cachedHashMap.get();
}

if (hashMapPtr->empty()) {
auto numKeys = rawSizes[0];
hashMapPtr->reserve(numKeys * 1.3);
for (auto i = 0; i < numKeys; ++i) {
hashMapPtr->insert(detail::MapKey{mapKeysBase, mapKeysIndices[i], i});
}
}

rows.applyToSelected([&](vector_size_t row) {
VELOX_CHECK_EQ(0, mapIndices[row]);

auto searchIndex = searchIndices[row];
auto it = set.find(MapKey{searchBase, searchIndex, row});
if (it != set.end()) {
auto it = hashMapPtr->find(detail::MapKey{searchBase, searchIndex, row});
if (it != hashMapPtr->end()) {
rawIndices[row] = it->index;
} else {
nullsBuilder.setNull(row);
Expand Down Expand Up @@ -302,6 +290,8 @@ VectorPtr applyMapComplexType(

} // namespace

namespace detail {

VectorPtr MapSubscript::applyMap(
const SelectivityVector& rows,
std::vector<VectorPtr>& args,
Expand All @@ -312,20 +302,20 @@ VectorPtr MapSubscript::applyMap(
// Ensure map key type and second argument are the same.
VELOX_CHECK(mapArg->type()->childAt(0)->equivalent(*indexArg->type()));

bool triggerCaching = shouldTriggerCaching(mapArg);
if (indexArg->type()->isPrimitiveType()) {
bool triggerCaching = shouldTriggerCaching(mapArg);

return VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH(
applyMapTyped,
indexArg->typeKind(),
triggerCaching,
const_cast<std::shared_ptr<LookupTableBase>&>(lookupTable_),
primitiveKeyLookupTable_,
rows,
mapArg,
indexArg,
context);
} else {
return applyMapComplexType(rows, mapArg, indexArg, context);
return applyMapComplexType(
rows, mapArg, indexArg, context, triggerCaching, complexKeyHashMap_);
}
}

Expand Down Expand Up @@ -369,5 +359,6 @@ const std::exception_ptr& negativeSubscriptError() {
static std::exception_ptr error = makeNegativeSubscriptError();
return error;
}
} // namespace detail

} // namespace facebook::velox::functions
75 changes: 61 additions & 14 deletions velox/functions/lib/SubscriptUtil.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

namespace facebook::velox::functions {

namespace detail {
// Below functions return a stock instance of each of the possible errors in
// SubscriptImpl
const std::exception_ptr& zeroSubscriptError();
Expand Down Expand Up @@ -96,6 +97,42 @@ class LookupTable : public LookupTableBase {
std::unique_ptr<outer_map_t> map_;
};

// A flat vector of map keys, an index into that vector and an index into
// the original map keys vector that may have encodings.
struct MapKey {
const BaseVector* baseVector;
const vector_size_t baseIndex;
const vector_size_t index;

size_t hash() const {
return baseVector->hashValueAt(baseIndex);
}

bool operator==(const MapKey& other) const {
return baseVector->equalValueAt(
other.baseVector, baseIndex, other.baseIndex);
}

bool operator<(const MapKey& other) const {
return baseVector->compare(other.baseVector, baseIndex, other.baseIndex) <
0;
}
};

struct MapKeyHasher {
size_t operator()(const MapKey& key) const {
return key.hash();
}
};

using MapKeyAllocator = memory::StlAllocator<detail::MapKey>;

using ComplexKeyHashMap = folly::F14FastSet<
detail::MapKey,
detail::MapKeyHasher,
folly::f14::DefaultKeyEqual<detail::MapKey>,
MapKeyAllocator>;

class MapSubscript {
public:
explicit MapSubscript(bool allowCaching) : allowCaching_(allowCaching) {}
Expand All @@ -109,8 +146,12 @@ class MapSubscript {
return allowCaching_;
}

auto& lookupTable() const {
return lookupTable_;
auto& primitiveKeyLookupTable() const {
return primitiveKeyLookupTable_;
}

auto& complexKeyHashMap() const {
return complexKeyHashMap_;
}

auto& firstSeenMap() const {
Expand All @@ -123,9 +164,8 @@ class MapSubscript {
return false;
}

if (!mapArg->type()->childAt(0)->isPrimitiveType() &&
!!mapArg->type()->childAt(0)->isBoolean()) {
// Disable caching if the key type is not primitive or is boolean.
if (mapArg->type()->childAt(0)->isBoolean()) {
// Disable caching if the key type is boolean.
allowCaching_ = false;
return false;
}
Expand All @@ -141,7 +181,8 @@ class MapSubscript {

// Disable caching forever.
allowCaching_ = false;
lookupTable_.reset();
primitiveKeyLookupTable_.reset();
complexKeyHashMap_.reset();
firstSeenMap_.reset();
return false;
}
Expand All @@ -155,9 +196,15 @@ class MapSubscript {
// seen again then it was not modified.
mutable VectorPtr firstSeenMap_;

// Materialized cached version of firstSeenMap_ used to optimize the lookup.
mutable std::shared_ptr<LookupTableBase> lookupTable_;
// Materialized cached version of firstSeenMap_ when the keys are primitive
// used to optimize the lookup.
mutable std::shared_ptr<LookupTableBase> primitiveKeyLookupTable_;

// Materialized cached version of firstSeenMap_ when the keys are complex used
// to optimize the lookup.
mutable std::shared_ptr<ComplexKeyHashMap> complexKeyHashMap_;
};
} // namespace detail

/// Generic subscript/element_at implementation for both array and map data
/// types.
Expand All @@ -179,7 +226,7 @@ template <
class SubscriptImpl : public exec::Subscript {
public:
explicit SubscriptImpl(bool allowCaching)
: mapSubscript_(MapSubscript(allowCaching)) {}
: mapSubscript_(detail::MapSubscript(allowCaching)) {}

void apply(
const SelectivityVector& rows,
Expand Down Expand Up @@ -286,7 +333,7 @@ class SubscriptImpl : public exec::Subscript {
const auto adjustedIndex =
adjustIndex(decodedIndices->valueAt<I>(0), isZeroSubscriptError);
if (isZeroSubscriptError) {
context.setErrors(rows, zeroSubscriptError());
context.setErrors(rows, detail::zeroSubscriptError());
allFailed = true;
}

Expand All @@ -307,7 +354,7 @@ class SubscriptImpl : public exec::Subscript {
const auto adjustedIndex =
adjustIndex(originalIndex, isZeroSubscriptError);
if (isZeroSubscriptError) {
context.setVeloxExceptionError(row, zeroSubscriptError());
context.setVeloxExceptionError(row, detail::zeroSubscriptError());
return;
}
const auto elementIndex = getIndex(
Expand Down Expand Up @@ -372,7 +419,7 @@ class SubscriptImpl : public exec::Subscript {
index += arraySize;
}
} else {
context.setVeloxExceptionError(row, negativeSubscriptError());
context.setVeloxExceptionError(row, detail::negativeSubscriptError());
return -1;
}
}
Expand All @@ -383,7 +430,7 @@ class SubscriptImpl : public exec::Subscript {
if constexpr (allowOutOfBound) {
return -1;
} else {
context.setVeloxExceptionError(row, badSubscriptError());
context.setVeloxExceptionError(row, detail::badSubscriptError());
return -1;
}
}
Expand All @@ -394,7 +441,7 @@ class SubscriptImpl : public exec::Subscript {
}

private:
MapSubscript mapSubscript_;
detail::MapSubscript mapSubscript_;
};

} // namespace facebook::velox::functions
21 changes: 20 additions & 1 deletion velox/functions/prestosql/benchmarks/MapSubscriptCachingBenchmark.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ extern void registerSubscriptFunction(

int main(int argc, char** argv) {
folly::Init init(&argc, &argv);
memory::MemoryManager::testingSetInstance({});

ExpressionBenchmarkBuilder benchmarkBuilder;
facebook::velox::functions::prestosql::registerAllScalarFunctions();
Expand All @@ -53,7 +54,8 @@ int main(int argc, char** argv) {
VectorFuzzer::Options options;
options.vectorSize = 1000;
options.containerLength = mapLength;
options.complexElementsMaxSize = 10000000000;
// Make sure it's big enough for nested complex types.
options.complexElementsMaxSize = baseVectorSize * mapLength * mapLength;
options.containerVariableLength = false;

VectorFuzzer fuzzer(options, pool);
Expand Down Expand Up @@ -112,6 +114,23 @@ int main(int argc, char** argv) {
createSetsForType(INTEGER());
createSetsForType(VARCHAR());

// For complex types, caching only applies if the Vector has a single constant
// value, so we only run with a baseVectorSize of 1. Also, due to the
// cost of the cardinality explosion from having nested complex types, we
// limit the number of iterations to 100.
auto createSetsForComplexType = [&](const auto& keyType) {
createSet(MAP(keyType, INTEGER()), 10, 1, 100);

createSet(MAP(keyType, INTEGER()), 100, 1, 100);

createSet(MAP(keyType, INTEGER()), 1000, 1, 100);

createSet(MAP(keyType, INTEGER()), 10000, 1, 100);
};

createSetsForComplexType(ARRAY(BIGINT()));
createSetsForComplexType(MAP(INTEGER(), VARCHAR()));

benchmarkBuilder.registerBenchmarks();
benchmarkBuilder.testBenchmarks();

Expand Down
Loading

0 comments on commit 4676df5

Please sign in to comment.