jabberwocky238/rocksdb
1
0
Fork 0
mirror of https://github.com/facebook/rocksdb.git synced 2026-07-07 14:47:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

Support Pinnable Reads In SstFileReader (#14500)

Browse Source 
Summary:
Add `SstFileReader::Get` (single-key) and `SstFileReader::MultiGet` (PinnableSlice) overloads to enable zero-copy point lookups directly from SST files. The existing `MultiGet(std::string*)` is refactored to delegate to the new `MultiGet(PinnableSlice*)`, which writes results directly into caller-provided `PinnableSlice` values instead of copying through an intermediate buffer. The single-key `Get` uses `TableReader::Get` with a `GetContext` for efficient single-key lookups without the overhead of MultiGet's sorting and batching machinery.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/14500

Test Plan: - New unit tests

Reviewed By: xingbowang

Differential Revision: D97825648

Pulled By: joshkang97

fbshipit-source-id: 17f3edd59bbf4747d17309c44ef12f0d952ea4eb
This commit is contained in:
Josh Kang
2026-03-25 15:11:15 -07:00
committed by meta-codesync[bot]
parent c78144e452
commit f620a6d039
5 changed files with 185 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/rocksdb/sst_file_reader.h
+13
View File
@@ -35,6 +35,19 @@ class SstFileReader {
const std::vector<Slice>& keys,
std::vector<std::string>* values);
// MultiGet variant that returns PinnableSlice values, enabling zero-copy
// when the underlying TableReader supports pinning.
std::vector<Status> MultiGet(const ReadOptions& options,
const std::vector<Slice>& keys,
std::vector<PinnableSlice>* values);
// Point lookup a single key from the SST.
Status Get(const ReadOptions& options, const Slice& key, std::string* value);
// Point lookup variant that returns a PinnableSlice.
Status Get(const ReadOptions& options, const Slice& key,
PinnableSlice* value);
// Returns a new iterator over the table contents as a raw table iterator,
// a.k.a a `TableIterator`that iterates all point data entries in the table
// including logically invisible entries like delete entries.
table/sst_file_reader.cc
+77 -8
View File
@@ -83,12 +83,15 @@ Status SstFileReader::Open(const std::string& file_path) {
return s;
}
std::vector<Status> SstFileReader::MultiGet(const ReadOptions& roptions,
const std::vector<Slice>& keys,
std::vector<std::string>* values) {
std::vector<Status> SstFileReader::MultiGet(
const ReadOptions& roptions, const std::vector<Slice>& keys,
std::vector<PinnableSlice>* values) {
const auto num_keys = keys.size();
std::vector<Status> statuses(num_keys, Status::OK());
std::vector<PinnableSlice> pin_values(num_keys);
values->resize(num_keys);
for (size_t i = 0; i < num_keys; ++i) {
(*values)[i].Reset();
}
auto r = rep_.get();
const Comparator* user_comparator =
@@ -101,8 +104,7 @@ std::vector<Status> SstFileReader::MultiGet(const ReadOptions& roptions,
autovector<MergeContext, MultiGetContext::MAX_BATCH_SIZE> merge_ctx;
sorted_keys.resize(num_keys);
for (size_t i = 0; i < num_keys; ++i) {
PinnableSlice* val = &pin_values[i];
val->Reset();
PinnableSlice* val = &(*values)[i];
merge_ctx.emplace_back();
key_context.emplace_back(nullptr, keys[i], val, nullptr,
nullptr /* timestamp */, &statuses[i]);
@@ -139,14 +141,12 @@ std::vector<Status> SstFileReader::MultiGet(const ReadOptions& roptions,
r->moptions.prefix_extractor.get(),
false /* skip filters */);
values->resize(num_keys);
for (size_t i = 0; i < num_keys; ++i) {
get_ctx[i].ReportCounters();
if (statuses[i].ok()) {
switch (get_ctx[i].State()) {
case GetContext::kFound:
(*values)[i].assign(pin_values[i].data(), pin_values[i].size());
break;
case GetContext::kNotFound:
case GetContext::kDeleted:
@@ -166,6 +166,75 @@ std::vector<Status> SstFileReader::MultiGet(const ReadOptions& roptions,
return statuses;
}
std::vector<Status> SstFileReader::MultiGet(const ReadOptions& roptions,
const std::vector<Slice>& keys,
std::vector<std::string>* values) {
std::vector<PinnableSlice> pin_values;
std::vector<Status> statuses = MultiGet(roptions, keys, &pin_values);
values->resize(keys.size());
for (size_t i = 0; i < keys.size(); ++i) {
if (statuses[i].ok()) {
(*values)[i].assign(pin_values[i].data(), pin_values[i].size());
}
}
return statuses;
}
Status SstFileReader::Get(const ReadOptions& roptions, const Slice& key,
PinnableSlice* value) {
auto r = rep_.get();
value->Reset();
const Comparator* user_comparator =
r->ioptions.internal_comparator.user_comparator();
Statistics* statistics = r->ioptions.stats;
Status status;
MergeContext merge_context;
SequenceNumber max_covering_tombstone_seq = 0;
GetContext get_ctx(user_comparator, r->ioptions.merge_operator.get(),
nullptr /* logger */, statistics, GetContext::kNotFound,
key, value, nullptr /* columns */, nullptr /* timestamp */,
nullptr /* value_found */, &merge_context, true,
&max_covering_tombstone_seq, r->ioptions.clock);
status = r->table_reader->Get(
roptions, InternalKey(key, kMaxSequenceNumber, kTypeValue).Encode(),
&get_ctx, r->moptions.prefix_extractor.get(), false /* skip_filters */);
get_ctx.ReportCounters();
if (status.ok()) {
switch (get_ctx.State()) {
case GetContext::kFound:
break;
case GetContext::kNotFound:
case GetContext::kDeleted:
status = Status::NotFound();
break;
case GetContext::kMerge:
status = Status::MergeInProgress();
break;
case GetContext::kCorrupt:
case GetContext::kUnexpectedBlobIndex:
case GetContext::kMergeOperatorFailed:
status = Status::Corruption();
break;
}
}
return status;
}
Status SstFileReader::Get(const ReadOptions& roptions, const Slice& key,
std::string* value) {
PinnableSlice pin_value;
Status s = Get(roptions, key, &pin_value);
if (s.ok()) {
value->assign(pin_value.data(), pin_value.size());
}
return s;
}
Iterator* SstFileReader::NewIterator(const ReadOptions& roptions) {
assert(roptions.io_activity == Env::IOActivity::kUnknown);
auto r = rep_.get();
table/sst_file_reader_test.cc
+44 -40
View File
@@ -768,37 +768,33 @@ TEST_F(SstFileReaderTableIteratorTest, UserDefinedTimestampsEnabled) {
Close();
}
class SstFileReaderTableMultiGetTest : public DBTestBase {
class SstFileReaderTableGetTest : public DBTestBase,
public testing::WithParamInterface<bool> {
public:
SstFileReaderTableMultiGetTest()
: DBTestBase("sst_file_reader_table_multi_get_test",
SstFileReaderTableGetTest()
: DBTestBase("sst_file_reader_table_get_test",
/*env_do_fsync=*/false) {}
void VerifyTableEntry(Iterator* iter, const std::string& user_key,
ValueType value_type,
std::optional<std::string> expected_value,
bool backward_iteration = false) {
ASSERT_TRUE(iter->Valid());
ASSERT_TRUE(iter->status().ok());
ParsedInternalKey pikey;
ASSERT_OK(ParseInternalKey(iter->key(), &pikey, /*log_err_key=*/false));
ASSERT_EQ(pikey.user_key, user_key);
ASSERT_EQ(pikey.type, value_type);
if (expected_value.has_value()) {
ASSERT_EQ(iter->value(), expected_value.value());
}
if (!backward_iteration) {
iter->Next();
bool UseMultiGet() const { return GetParam(); }
std::vector<Status> DoGet(SstFileReader& reader,
const std::vector<Slice>& keys,
std::vector<std::string>* values) {
if (UseMultiGet()) {
return reader.MultiGet(ReadOptions(), keys, values);
} else {
iter->Prev();
values->resize(keys.size());
std::vector<Status> statuses(keys.size());
for (size_t i = 0; i < keys.size(); ++i) {
statuses[i] = reader.Get(ReadOptions(), keys[i], &(*values)[i]);
}
return statuses;
}
}
};
TEST_F(SstFileReaderTableMultiGetTest, Basic) {
TEST_P(SstFileReaderTableGetTest, Basic) {
Options options = CurrentOptions();
const Comparator* ucmp = BytewiseComparator();
options.comparator = ucmp;
options.disable_auto_compactions = true;
options.merge_operator = MergeOperators::CreateStringAppendOperator();
options.statistics = CreateDBStatistics();
@@ -823,7 +819,7 @@ TEST_F(SstFileReaderTableMultiGetTest, Basic) {
std::vector<LiveFileMetaData> files;
dbfull()->GetLiveFilesMetaData(&files);
ASSERT_TRUE(files.size() == 1);
ASSERT_EQ(files.size(), 1);
ASSERT_TRUE(files[0].level == 0);
std::string file_name = files[0].directory + "/" + files[0].relative_filename;
@@ -833,25 +829,30 @@ TEST_F(SstFileReaderTableMultiGetTest, Basic) {
ASSERT_OK(options.statistics->Reset());
std::vector<Slice> keys;
std::vector<Slice> keys = {"fo1", "foo", "baz",
"bar", "aaa", "zzz_not_in_sst"};
std::vector<std::string> values;
auto statuses = DoGet(reader, keys, &values);
keys.emplace_back("fo1");
keys.emplace_back("foo");
keys.emplace_back("baz");
keys.emplace_back("bar");
keys.emplace_back("aaa");
auto statuses = reader.MultiGet(ReadOptions(), keys, &values);
ASSERT_TRUE(statuses[0].IsNotFound())
<< "Failed: status=" << statuses[0].ToString() << " val=" << values[0];
ASSERT_TRUE(statuses[1].IsNotFound())
<< "Failed: status=" << statuses[0].ToString() << " val=" << values[1];
ASSERT_TRUE(statuses[2].ok());
ASSERT_TRUE(statuses[3].ok());
ASSERT_EQ("val3", values[2]);
ASSERT_EQ("val2", values[3]);
ASSERT_TRUE(statuses[4].ok());
ASSERT_EQ("val4,val5", values[4]);
// Non-existent key returns NotFound
ASSERT_TRUE(statuses[0].IsNotFound());
// Deleted key returns NotFound
ASSERT_TRUE(statuses[1].IsNotFound());
// Found keys
ASSERT_OK(statuses[2]);
ASSERT_EQ(values[2], "val3");
ASSERT_OK(statuses[3]);
ASSERT_EQ(values[3], "val2");
// Merged key
ASSERT_OK(statuses[4]);
ASSERT_EQ(values[4], "val4,val5");
// Bloom filter filtered key
ASSERT_TRUE(statuses[5].IsNotFound());
uint64_t cache_hits = options.statistics->getTickerCount(BLOCK_CACHE_HIT);
uint64_t cache_misses = options.statistics->getTickerCount(BLOCK_CACHE_MISS);
@@ -863,6 +864,9 @@ TEST_F(SstFileReaderTableMultiGetTest, Basic) {
Close();
}
INSTANTIATE_TEST_CASE_P(SingleAndMulti, SstFileReaderTableGetTest,
testing::Bool());
} // namespace ROCKSDB_NAMESPACE
int main(int argc, char** argv) {
table/table_test.cc
+50
View File
@@ -7410,6 +7410,56 @@ TEST_F(ExternalTableTest, PinnedGetTest) {
ASSERT_EQ(factory->last_reader()->pin_cleanup_count(), 4);
}
TEST_F(ExternalTableTest, SstReaderPinnableMultiGetTest) {
if (encrypted_env_) {
ROCKSDB_GTEST_SKIP("Test requires non-encrypted environment");
return;
}
Options options = GetDefaultOptions();
std::string dbname =
test::PerThreadDBPath("sst_reader_pinnable_multiget_test");
std::string sst_file = dbname + "/test.sst";
ASSERT_OK(options.env->CreateDirIfMissing(dbname));
std::unique_ptr<SstFileWriter> writer(
new SstFileWriter(EnvOptions(), options));
ASSERT_OK(writer->Open(sst_file));
ASSERT_OK(writer->Put("a", "val_a"));
ASSERT_OK(writer->Put("b", "val_b"));
ASSERT_OK(writer->Put("c", "val_c"));
ASSERT_OK(writer->Finish());
writer.reset();
std::unique_ptr<SstFileReader> reader(new SstFileReader(options));
ASSERT_OK(reader->Open(sst_file));
// Test PinnableSlice MultiGet
std::vector<Slice> keys = {"a", "b", "missing", "c"};
std::vector<PinnableSlice> values;
std::vector<Status> statuses = reader->MultiGet(ReadOptions(), keys, &values);
ASSERT_EQ(values.size(), keys.size());
ASSERT_EQ(statuses.size(), keys.size());
ASSERT_OK(statuses[0]);
ASSERT_EQ(values[0].ToString(), "val_a");
ASSERT_OK(statuses[1]);
ASSERT_EQ(values[1].ToString(), "val_b");
ASSERT_TRUE(statuses[2].IsNotFound());
ASSERT_OK(statuses[3]);
ASSERT_EQ(values[3].ToString(), "val_c");
// Verify std::string MultiGet wrapper still works
std::vector<std::string> str_values;
statuses = reader->MultiGet(ReadOptions(), keys, &str_values);
ASSERT_EQ(str_values.size(), keys.size());
ASSERT_OK(statuses[0]);
ASSERT_EQ(str_values[0], "val_a");
ASSERT_OK(statuses[1]);
ASSERT_EQ(str_values[1], "val_b");
ASSERT_TRUE(statuses[2].IsNotFound());
ASSERT_OK(statuses[3]);
ASSERT_EQ(str_values[3], "val_c");
}
TEST_F(ExternalTableTest, ExternalFileChecksumTest) {
if (encrypted_env_) {
ROCKSDB_GTEST_SKIP("Test requires non-encrypted environment");
unreleased_history/public_api_changes/sst_file_reader_pinnable_multiget.md
+1
View File
@@ -0,0 +1 @@
Added `SstFileReader::Get` and `SstFileReader::MultiGet` overloads that accept `PinnableSlice`/`std::vector<PinnableSlice>*`, enabling zero-copy reads when the underlying `TableReader` supports pinning.