mirror of
https://github.com/scylladb/scylladb.git
synced 2026-05-01 21:55:50 +00:00
ff5add4287
Delaying the BTI encoding of partition keys is a good idea, because most of the time they don't have to be encoded. Usually the token alone is enough for indexing purposes. But for the translation of the `partition_key` part itself, there's no good reason to make it lazy, especially after we made the translation of clustering keys eager in a previous commit. Let's get rid of the `std::generator` and convert all cells of the partition key in one go.
328 lines
14 KiB
C++
328 lines
14 KiB
C++
/*
|
|
* Copyright (C) 2024-present ScyllaDB
|
|
*/
|
|
|
|
/*
|
|
* SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
|
|
*/
|
|
|
|
#define BOOST_TEST_MODULE core
|
|
|
|
#include <boost/test/unit_test.hpp>
|
|
#include "seastar/util/defer.hh"
|
|
#include "schema/schema_builder.hh"
|
|
#include "sstables/trie/bti_key_translation.hh"
|
|
#include "test/lib/key_utils.hh"
|
|
#include "test/lib/log.hh"
|
|
#include "test/lib/nondeterministic_choice_stack.hh"
|
|
#include <generator>
|
|
|
|
static std::vector<std::byte> linearize(comparable_bytes_iterator auto&& it) {
|
|
auto result = std::vector<std::byte>();
|
|
while (it != std::default_sentinel) {
|
|
result.append_range(*it);
|
|
++it;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// Generate an ordered list of various ring position views,
|
|
// which together should cover all kinds of various pairwise comparisons.
|
|
static std::generator<dht::ring_position_view> generate_rpvs(const schema& string_pair_pk_schema) {
|
|
constexpr auto alphabet = std::string_view("\x00\xff", 2);
|
|
constexpr auto max_len = 2;
|
|
auto all_strings = tests::generate_all_strings(alphabet, max_len);
|
|
co_yield dht::ring_position_view::min();
|
|
for (const auto& token : {dht::token::first(), dht::token::last()}) {
|
|
using token_bound = dht::ring_position_view::token_bound;
|
|
co_yield dht::ring_position_view(token, token_bound::start);
|
|
for (const auto& pk1 : all_strings)
|
|
for (const auto& pk2 : all_strings) {
|
|
std::vector<data_value> components;
|
|
components.emplace_back(pk1);
|
|
components.emplace_back(pk2);
|
|
auto pk = partition_key::from_deeply_exploded(string_pair_pk_schema, components);
|
|
for (int weight = -1; weight <= 1; ++weight) {
|
|
co_yield dht::ring_position_view(token, &pk, weight);
|
|
}
|
|
}
|
|
co_yield dht::ring_position_view(token, token_bound::end);
|
|
}
|
|
}
|
|
|
|
// Performs key conversion to BTI format for various keys,
|
|
// and checks that this encoding doesn't change the ordering of keys.
|
|
BOOST_AUTO_TEST_CASE(test_lazy_comparable_bytes_from_ring_position_preserves_order) {
|
|
auto s = schema_builder("ks", "t")
|
|
.with_column("pk1", utf8_type, column_kind::partition_key)
|
|
.with_column("pk2", utf8_type, column_kind::partition_key)
|
|
.build();
|
|
using encoding = sstables::trie::lazy_comparable_bytes_from_ring_position;
|
|
std::unique_ptr<encoding> prev;
|
|
for (const auto& rpv : generate_rpvs(*s)) {
|
|
if (!prev) {
|
|
prev = std::make_unique<encoding>(*s, rpv);
|
|
} else {
|
|
auto curr = std::make_unique<encoding>(*s, rpv);
|
|
auto prev_bytes = linearize(prev->begin());
|
|
auto curr_bytes = linearize(curr->begin());
|
|
testlog.debug("prev_bytes={}, curr_bytes={}", fmt_hex(prev_bytes), fmt_hex(curr_bytes));
|
|
SCYLLA_ASSERT(prev_bytes < curr_bytes);
|
|
std::swap(prev, curr);
|
|
}
|
|
}
|
|
auto prev_bytes = linearize(prev->begin());
|
|
SCYLLA_ASSERT(prev_bytes <= linearize(encoding(*s, dht::ring_position_view(dht::maximum_token(), nullptr, -1)).begin()));
|
|
SCYLLA_ASSERT(prev_bytes <= linearize(encoding(*s, dht::ring_position_view::max()).begin()));
|
|
}
|
|
|
|
// Tests lazy_comparable_bytes_from_ring_position::trim().
|
|
//
|
|
// Description:
|
|
// 1. Constructs some single arbitrary ring position `dk`.
|
|
// 2. For all legal possibilities:
|
|
// 2.1 Constructs `enc := lazy_comparable_bytes_from_ring_position(dk)`.
|
|
// 2.2 Reads N bytes from `enc`.
|
|
// 2.3 Optionally modifies some chosen byte on the way.
|
|
// 2.4 Trims `enc` to length M.
|
|
// 2.5 Checks that reading `enc` again gives a result consistent with the
|
|
// modification and the trim.
|
|
BOOST_AUTO_TEST_CASE(test_lazy_comparable_bytes_from_ring_position_trim) {
|
|
auto s = schema_builder("ks", "t")
|
|
.with_column("pk1", utf8_type, column_kind::partition_key)
|
|
.with_column("pk2", utf8_type, column_kind::partition_key)
|
|
.build();
|
|
auto pk = ({
|
|
std::vector<data_value> components;
|
|
components.emplace_back("ab");
|
|
components.emplace_back("cd");
|
|
partition_key::from_deeply_exploded(*s, components);
|
|
});
|
|
using encoding = sstables::trie::lazy_comparable_bytes_from_ring_position;
|
|
const auto dk = dht::decorated_key(dht::token::from_int64(42), std::move(pk));
|
|
const auto encoded = linearize(encoding(*s, dk).begin());
|
|
|
|
constexpr auto modification_byte = std::byte('z');
|
|
for (size_t view_position = 0; view_position <= encoded.size(); ++view_position)
|
|
for (size_t trim_position = 0; trim_position <= view_position; ++trim_position)
|
|
for (size_t modification_position = 0; modification_position <= trim_position; ++modification_position) {
|
|
auto enc = encoding(*s, std::move(dk));
|
|
{
|
|
size_t n_seen = 0;
|
|
auto it = enc.begin();
|
|
while (it != std::default_sentinel) {
|
|
auto fragment_start = n_seen;
|
|
n_seen += (*it).size_bytes();
|
|
if (fragment_start <= modification_position && modification_position < n_seen) {
|
|
(*it)[modification_position - fragment_start] = modification_byte;
|
|
}
|
|
if (view_position <= n_seen) {
|
|
break;
|
|
}
|
|
++it;
|
|
}
|
|
enc.trim(trim_position);
|
|
}
|
|
auto expected_encoded = encoded;
|
|
if (modification_position < trim_position) {
|
|
expected_encoded[modification_position] = modification_byte;
|
|
}
|
|
expected_encoded.resize(trim_position);
|
|
auto new_encoded = linearize(enc.begin());
|
|
testlog.debug("view_position={}, trim_position={}, modification_position={}", view_position, trim_position, modification_position);
|
|
testlog.debug("new_encoded={}, expected_encoded={}", fmt_hex(new_encoded), fmt_hex(expected_encoded));
|
|
SCYLLA_ASSERT(new_encoded == expected_encoded);
|
|
}
|
|
}
|
|
|
|
// Generate an ordered list of various clustering positions views,
|
|
// which together should cover all kinds of various pairwise comparisons.
|
|
static std::generator<position_in_partition_view> generate_pipvs(const schema& string_pair_ck_schema) {
|
|
constexpr auto alphabet = std::string_view("\x00\xff", 2);
|
|
constexpr auto max_len = 2;
|
|
auto all_strings = tests::generate_all_strings(alphabet, max_len);
|
|
co_yield position_in_partition_view::for_partition_start();
|
|
co_yield position_in_partition_view::before_all_clustered_rows();
|
|
for (const auto& ck1 : all_strings) {
|
|
std::vector<data_value> components;
|
|
components.emplace_back(ck1);
|
|
auto ckp = clustering_key_prefix::from_deeply_exploded(string_pair_ck_schema, components);
|
|
co_yield position_in_partition_view(ckp, bound_weight(-1));
|
|
for (const auto& ck2 : all_strings) {
|
|
components.emplace_back(ck2);
|
|
auto _pop = defer([&] { components.pop_back(); });
|
|
|
|
auto ck = clustering_key_prefix::from_deeply_exploded(string_pair_ck_schema, components);
|
|
for (int weight = -1; weight <= 1; ++weight) {
|
|
auto ret = position_in_partition_view(ck, bound_weight(weight));
|
|
testlog.debug("pipv={}", ret);
|
|
co_yield ret;
|
|
}
|
|
}
|
|
co_yield position_in_partition_view(ckp, bound_weight(1));
|
|
}
|
|
co_yield position_in_partition_view::after_all_clustered_rows();
|
|
co_yield position_in_partition_view::for_partition_end();
|
|
}
|
|
|
|
// Performs key conversion to BTI format for various keys,
|
|
// and checks that this encoding doesn't change the ordering of keys.
|
|
BOOST_AUTO_TEST_CASE(test_lazy_comparable_bytes_from_clustering_position_preserves_order) {
|
|
auto s = schema_builder("ks", "t")
|
|
.with_column("pk", long_type, column_kind::partition_key)
|
|
.with_column("ck1", utf8_type, column_kind::clustering_key)
|
|
.with_column("ck2", utf8_type, column_kind::clustering_key)
|
|
.build();
|
|
using encoding = sstables::trie::lazy_comparable_bytes_from_clustering_position;
|
|
std::unique_ptr<encoding> prev;
|
|
for (const auto& pipv : generate_pipvs(*s)) {
|
|
if (!prev) {
|
|
prev = std::make_unique<encoding>(*s, pipv);
|
|
} else {
|
|
auto curr = std::make_unique<encoding>(*s, pipv);
|
|
auto prev_bytes = linearize(prev->begin());
|
|
auto curr_bytes = linearize(curr->begin());
|
|
testlog.debug("prev_bytes={}, curr_bytes={}", fmt_hex(prev_bytes), fmt_hex(curr_bytes));
|
|
SCYLLA_ASSERT(prev_bytes < curr_bytes);
|
|
std::swap(prev, curr);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Tests lazy_comparable_bytes_from_ckp::trim().
|
|
// Just like test_lazy_comparable_bytes_from_ring_position_trim,
|
|
// but for clustering positions.
|
|
BOOST_AUTO_TEST_CASE(test_lazy_comparable_bytes_from_clustering_position_trim) {
|
|
auto s = schema_builder("ks", "t")
|
|
.with_column("pk", long_type, column_kind::partition_key)
|
|
.with_column("ck1", utf8_type, column_kind::clustering_key)
|
|
.with_column("ck2", utf8_type, column_kind::clustering_key)
|
|
.build();
|
|
auto ck = ({
|
|
std::vector<data_value> components;
|
|
components.emplace_back("ab");
|
|
components.emplace_back("cd");
|
|
clustering_key_prefix::from_deeply_exploded(*s, components);
|
|
});
|
|
using encoding = sstables::trie::lazy_comparable_bytes_from_clustering_position;
|
|
const auto pipv = position_in_partition_view(ck, bound_weight(0));
|
|
const auto encoded = linearize(encoding(*s, pipv).begin());
|
|
|
|
constexpr auto modification_byte = std::byte('z');
|
|
for (size_t view_position = 0; view_position <= encoded.size(); ++view_position)
|
|
for (size_t trim_position = 0; trim_position <= view_position; ++trim_position)
|
|
for (size_t modification_position = 0; modification_position <= trim_position; ++modification_position) {
|
|
auto enc = encoding(*s, pipv);
|
|
{
|
|
size_t n_seen = 0;
|
|
auto it = enc.begin();
|
|
while (it != std::default_sentinel) {
|
|
auto fragment_start = n_seen;
|
|
n_seen += (*it).size_bytes();
|
|
if (fragment_start <= modification_position && modification_position < n_seen) {
|
|
(*it)[modification_position - fragment_start] = modification_byte;
|
|
}
|
|
if (view_position <= n_seen) {
|
|
break;
|
|
}
|
|
++it;
|
|
}
|
|
enc.trim(trim_position);
|
|
}
|
|
auto expected_encoded = encoded;
|
|
if (modification_position < trim_position) {
|
|
expected_encoded[modification_position] = modification_byte;
|
|
}
|
|
expected_encoded.resize(trim_position);
|
|
auto new_encoded = linearize(enc.begin());
|
|
SCYLLA_ASSERT(new_encoded == expected_encoded);
|
|
}
|
|
}
|
|
|
|
// A pretty printer for a string split into several bytespans,
|
|
// as happens in a test below.
|
|
struct fmt_fragmented_buffer {
|
|
std::span<std::span<const std::byte>> frags;
|
|
};
|
|
template <>
|
|
struct fmt::formatter<fmt_fragmented_buffer> : fmt::formatter<fmt_hex> {
|
|
template <typename FormatContext>
|
|
auto format(const fmt_fragmented_buffer& s, FormatContext& ctx) const {
|
|
fmt::format_to(ctx.out(), "[");
|
|
const char* separator = "";
|
|
for (const auto& frag : s.frags) {
|
|
fmt::format_to(ctx.out(), "{}{}", separator, fmt_hex(frag));
|
|
separator = "|";
|
|
}
|
|
fmt::format_to(ctx.out(), "]");
|
|
return ctx.out();
|
|
}
|
|
};
|
|
|
|
// Tests the lcb_mismatch function.
|
|
// It generates all possible pairs of strings up to some complexity
|
|
// (alphabet size, max length), splits them into fragments in all possible ways
|
|
// (up to some max number of empty fragments),
|
|
// and checks that the lcb_mismatch function gives results consistent with
|
|
// std::ranges::mismatch.
|
|
BOOST_AUTO_TEST_CASE(test_lcb_mismatch) {
|
|
// The set of all strings up to the given complexity.
|
|
auto all_strings = tests::generate_all_strings("ab", 3);
|
|
nondeterministic_choice_stack ncs;
|
|
do {
|
|
// Choose a pair of strings to compare, nondeterministically.
|
|
auto a = std::as_bytes(std::span(all_strings[ncs.choose_up_to(all_strings.size() - 1)]));
|
|
auto b = std::as_bytes(std::span(all_strings[ncs.choose_up_to(all_strings.size() - 1)]));
|
|
auto expected_mismatch = std::ranges::mismatch(a, b);
|
|
uint64_t expected_mismatch_offset = expected_mismatch.in2 - b.begin();
|
|
auto expected_mismatch_ptr = &*expected_mismatch.in2;
|
|
|
|
using const_bytes = sstables::trie::const_bytes;
|
|
// Split the byte span into fragments, nondeterministically.
|
|
auto gen_fragments = [&ncs] (const_bytes sp) -> std::generator<const_bytes> {
|
|
int empty_spans_in_a_row = 0;
|
|
if (sp.empty()) {
|
|
co_yield sp;
|
|
co_return;
|
|
}
|
|
while (sp.size() > 0) {
|
|
auto next_fragment_size = empty_spans_in_a_row < 2
|
|
? ncs.choose_up_to(sp.size())
|
|
: 1 + ncs.choose_up_to(sp.size() - 1);
|
|
empty_spans_in_a_row = (next_fragment_size == 0) ? empty_spans_in_a_row + 1 : 0;
|
|
co_yield sp.subspan(0, next_fragment_size);
|
|
sp = sp.subspan(next_fragment_size);
|
|
}
|
|
};
|
|
auto [mismatch_offset, mismatch_ptr] = std::invoke([&] {
|
|
// Extracting the full generator to a vector is expensive,
|
|
// so do it only if there's a need for nice debug logs.
|
|
constexpr bool debug = false;
|
|
if constexpr (debug) {
|
|
// This branch is effectively the same as the other branch, but it nicely logs the split into fragments.
|
|
auto fragments_a = std::ranges::to<std::vector>(gen_fragments(a));
|
|
auto fragments_b = std::ranges::to<std::vector>(gen_fragments(b));
|
|
testlog.debug("a={}, b={}",
|
|
fmt_fragmented_buffer{fragments_a},
|
|
fmt_fragmented_buffer{fragments_b});
|
|
auto generator_from_range = [&](auto& r) -> std::generator<std::ranges::range_value_t<decltype(r)>> {
|
|
for (auto e : r) {
|
|
co_yield e;
|
|
}
|
|
};
|
|
return sstables::trie::lcb_mismatch(
|
|
generator_from_range(fragments_a).begin(),
|
|
generator_from_range(fragments_b).begin()
|
|
);
|
|
} else {
|
|
return sstables::trie::lcb_mismatch(
|
|
gen_fragments(a).begin(),
|
|
gen_fragments(b).begin()
|
|
);
|
|
}
|
|
});
|
|
SCYLLA_ASSERT(mismatch_offset == expected_mismatch_offset);
|
|
SCYLLA_ASSERT(mismatch_ptr == expected_mismatch_ptr);
|
|
} while (ncs.rewind());
|
|
}
|