/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <folly/compression/QuotientMultiSet.h>
#include <random>
#include <folly/Format.h>
#include <folly/Random.h>
#include <folly/container/Enumerate.h>
#include <folly/io/IOBufQueue.h>
#include <folly/portability/GTest.h>
#if FOLLY_QUOTIENT_MULTI_SET_SUPPORTED
namespace {
class QuotientMultiSetTest : public ::testing::Test {
protected:
static constexpr uint64_t kBlockSize = folly::QuotientMultiSet<>::kBlockSize;
void SetUp() override { rng.seed(folly::randomNumberSeed()); }
void buildAndValidate(
std::vector<uint64_t>& keys, uint64_t keyBits, double loadFactor) {
// Elements must be added in ascending order.
std::sort(keys.begin(), keys.end());
folly::QuotientMultiSetBuilder builder(keyBits, keys.size(), loadFactor);
folly::IOBufQueue buff;
for (const auto& iter : folly::enumerate(keys)) {
if (builder.insert(*iter)) {
// Set payload to relative position of the first key in block.
builder.setBlockPayload(iter.index);
}
if (builder.numReadyBlocks() >= 1) {
builder.flush(buff);
}
}
builder.close(buff);
auto spBuf = buff.move();
folly::StringPiece data(spBuf->coalesce());
folly::QuotientMultiSet<> reader(data);
size_t index = 0;
folly::QuotientMultiSet<>::Iterator iter(&reader);
while (index < keys.size()) {
uint64_t start = index;
uint64_t& key = keys[start];
auto debugInfo = [&] {
return folly::sformat("Index: {} Key: {}", index, key);
};
while (index < keys.size() && keys[index] == key) {
index++;
}
auto range = reader.equalRange(key);
size_t pos = reader.getBlockPayload(range.begin / kBlockSize);
EXPECT_EQ(start, pos + range.begin % kBlockSize) << debugInfo();
pos = reader.getBlockPayload((range.end - 1) / kBlockSize);
EXPECT_EQ(index - 1, pos + (range.end - 1) % kBlockSize) << debugInfo();
EXPECT_TRUE(iter.skipTo(key));
EXPECT_EQ(key, iter.key()) << debugInfo();
EXPECT_EQ(range.begin, iter.pos()) << debugInfo();
EXPECT_EQ(
start,
reader.getBlockPayload(iter.pos() / kBlockSize) +
iter.pos() % kBlockSize)
<< debugInfo();
if (start < keys.size() - 1) {
EXPECT_TRUE(iter.next());
EXPECT_EQ(keys[start + 1], iter.key());
}
// Verify getting keys not in keys returns false.
uint64_t prevKey = (start == 0 ? 0 : keys[start - 1]);
if (prevKey + 1 < key) {
uint64_t missedKey = folly::Random::rand64(prevKey + 1, key, rng);
EXPECT_FALSE(reader.equalRange(missedKey)) << key;
iter.skipTo(missedKey);
EXPECT_EQ(key, iter.key()) << debugInfo();
EXPECT_EQ(range.begin, iter.pos()) << debugInfo();
EXPECT_EQ(
start,
reader.getBlockPayload(iter.pos() / kBlockSize) +
iter.pos() % kBlockSize)
<< debugInfo();
if (start < keys.size() - 1) {
EXPECT_TRUE(iter.next());
EXPECT_EQ(keys[start + 1], iter.key());
}
}
}
const auto maxKey = folly::qms_detail::maxValue(keyBits);
if (keys.back() < maxKey) {
uint64_t key = folly::Random::rand64(keys.back(), maxKey, rng) + 1;
EXPECT_FALSE(reader.equalRange(key)) << keys.back() << " " << key;
EXPECT_FALSE(iter.done());
EXPECT_FALSE(iter.skipTo(key));
EXPECT_TRUE(iter.done());
}
folly::QuotientMultiSet<>::Iterator nextIter(&reader);
for (const auto key : keys) {
EXPECT_TRUE(nextIter.next());
EXPECT_EQ(key, nextIter.key());
}
EXPECT_FALSE(nextIter.next());
EXPECT_TRUE(nextIter.done());
if (maxKey + 1 != 0) {
folly::QuotientMultiSet<>::Iterator skipToEndIter(&reader);
EXPECT_FALSE(skipToEndIter.done());
EXPECT_FALSE(skipToEndIter.skipTo(maxKey + 1));
EXPECT_TRUE(skipToEndIter.done());
}
}
std::mt19937 rng;
};
} // namespace
TEST_F(QuotientMultiSetTest, Simple) {
std::vector<uint64_t> keys = {
100, 1000, 1 << 14, 10 << 14, 0, 10, 100 << 14, 1000 << 14};
buildAndValidate(keys, 32, 0.95);
}
TEST_F(QuotientMultiSetTest, Empty) {
folly::QuotientMultiSetBuilder builder(32, 1024);
folly::IOBufQueue buff;
builder.close(buff);
auto spBuf = buff.move();
folly::StringPiece data(spBuf->coalesce());
folly::QuotientMultiSet<> reader(data);
for (size_t idx = 0; idx < 1024; idx++) {
uint64_t key = folly::Random::rand32(rng);
EXPECT_FALSE(reader.equalRange(key));
}
}
TEST_F(QuotientMultiSetTest, ZeroKeyBits) {
std::vector<uint64_t> keys(67, 0);
buildAndValidate(keys, 0, 0.95);
}
TEST_F(QuotientMultiSetTest, Uniform) {
constexpr auto kLoadFactor =
folly::QuotientMultiSetBuilder::kDefaultMaxLoadFactor;
constexpr uint64_t kAvgSize = 1 << 16;
auto randSize = [&](uint64_t avgSize) {
return folly::Random::rand64(avgSize / 2, avgSize * 3 / 2, rng);
};
std::vector<std::tuple<int, uint64_t, double>> testCases = {
{1, randSize(1 << 9), kLoadFactor},
{8, randSize(1 << 10), kLoadFactor},
{9, randSize(1 << 11), kLoadFactor},
{12, randSize(kAvgSize), kLoadFactor},
{32, randSize(kAvgSize), kLoadFactor},
{48, randSize(kAvgSize), kLoadFactor},
{64, randSize(kAvgSize), kLoadFactor},
{32, randSize(kAvgSize), 1}, // Full
{12, 3800, kLoadFactor}, // Almost full
{64, randSize(16), kLoadFactor}, // Sparse, long keys.
};
for (const auto& testCase : testCases) {
const auto& [keyBits, size, loadFactor] = testCase;
SCOPED_TRACE(folly::sformat(
"Key bits: {} Size: {} Load factor: {}", keyBits, size, loadFactor));
std::vector<uint64_t> keys;
for (uint64_t idx = 0; idx < size; idx++) {
keys.emplace_back(
folly::Random::rand64(rng) & folly::qms_detail::maxValue(keyBits));
}
buildAndValidate(keys, keyBits, loadFactor);
}
}
TEST_F(QuotientMultiSetTest, UniformDistributionFullLoadFactor) {
const uint64_t numElements = 1 << 16;
std::vector<uint64_t> keys;
for (uint64_t idx = 0; idx < numElements; idx++) {
uint64_t key = folly::Random::rand32(idx << 16, (idx + 1) << 16, rng);
keys.emplace_back(key);
}
buildAndValidate(keys, 32, 1.0);
}
TEST_F(QuotientMultiSetTest, Overflow) {
const uint64_t numElements = 1 << 12;
std::vector<uint64_t> keys;
for (uint64_t idx = 0; idx < numElements; idx++) {
keys.emplace_back(idx);
keys.emplace_back(idx);
keys.emplace_back(idx);
}
buildAndValidate(keys, 12, 0.95);
}
TEST_F(QuotientMultiSetTest, RandomLengthRuns) {
const uint64_t numElements = 1 << 16;
std::vector<uint64_t> keys;
for (uint64_t idx = 0; idx < (numElements >> 4); idx++) {
uint64_t key = folly::Random::rand32(rng);
uint64_t length = folly::Random::rand32(0, 10, rng);
for (uint64_t k = 0; k < length; k++) {
keys.emplace_back(key + k);
}
}
buildAndValidate(keys, 32, 0.95);
}
TEST_F(QuotientMultiSetTest, RunAcrossBlocks) {
const uint64_t numElements = 1 << 10;
std::vector<uint64_t> keys;
// Add keys with cluster size 137.
for (uint64_t idx = 0; idx < (numElements >> 4); idx++) {
uint64_t key = folly::Random::rand32(rng);
for (uint64_t k = 0; k < 136; k++) {
key += k;
keys.emplace_back(key);
}
}
buildAndValidate(keys, 32, 0.95);
}
TEST_F(QuotientMultiSetTest, PackAtHeadSlots) {
const uint64_t numElements = 1 << 12;
std::vector<uint64_t> keys;
for (uint64_t idx = 0; idx < numElements; idx++) {
uint64_t key = folly::Random::rand32(idx << 8, (idx + 1) << 8, rng);
keys.emplace_back(key);
}
buildAndValidate(keys, 32, 0.95);
}
TEST_F(QuotientMultiSetTest, PackAtTailSlots) {
const uint64_t numElements = 1 << 12;
std::vector<uint64_t> keys;
uint64_t key = (1 << 30);
for (uint64_t idx = 0; idx < numElements; idx++) {
keys.emplace_back(key + idx);
}
buildAndValidate(keys, 32, 0.95);
}
TEST_F(QuotientMultiSetTest, KeysOnlyInHeadAndTail) {
const uint64_t numElements = 1 << 11;
std::vector<uint64_t> keys;
for (uint64_t idx = 0; idx < numElements; idx++) {
keys.emplace_back(idx);
}
uint64_t key = (1 << 30);
for (uint64_t idx = 0; idx < numElements; idx++) {
keys.emplace_back(key + idx);
}
buildAndValidate(keys, 32, 0.95);
}
TEST_F(QuotientMultiSetTest, RunendRightBeforeFirstOccupiedRunend) {
std::vector<uint64_t> keys;
// 60 ranges [0, 67] with occupied slot 20.
for (size_t idx = 0; idx < 68; idx++) {
keys.push_back(60);
}
// 60 ranges [68, 68] with occupied slot 66.
for (size_t idx = 0; idx < 1; idx++) {
keys.push_back(200);
}
// 60 ranges [69, 88] with occupied slot 83.
for (size_t idx = 0; idx < 20; idx++) {
keys.push_back(250);
}
buildAndValidate(keys, 8, 0.95);
}
#endif // FOLLY_QUOTIENT_MULTI_SET_SUPPORTED
Codebase Browser
folly