Adding support for prefix queries by maheshrbapatu · Pull Request #1870 · kvcache-ai/Mooncake

maheshrbapatu · 2026-04-12T02:34:22Z

Description

Introduce a compressed Patricia-style radix tree over string keys to support efficient prefix-based enumeration and aggregation.

The trie is implemented as an independent data structure and is not yet integrated with MasterService or the existing hash-sharded metadata index.

Supports:

insert/remove of full string keys with automatic node cleanup
listing keys and lightweight leaf data by prefix
handling prefixes ending mid-edge (compressed trie semantics)
prefix continuation listing for hierarchical browsing
subtree aggregation (count, bytes) via maintained statistics

Uses coarse-grained locking to provide thread-safe concurrent reads and writes.

Stores lightweight per-terminal-node summary data (TrieLeafData) to keep scope decoupled from full metadata lifecycle.

Includes unit tests covering:

insert/remove and prefix query correctness
edge cases (partial edge splits, prefix-as-key, mid-edge traversal)
subtree statistic invariants (count/bytes consistency)
cleanup behavior after deletes
concurrent access patterns

Module

Type of Change

How Has This Been Tested?

Added a bunch of new UTs.

Checklist

I have performed a self-review of my own code.
I have formatted my own code using ./scripts/code_format.sh before submitting.
[] I have updated the documentation.
I have added tests to prove my changes are effective.

Ref #1732

gemini-code-assist

Code Review

This pull request introduces a PrefixIndex class, implemented as a compressed radix tree, to support efficient prefix-based operations. However, the PR contains critical issues: the CMakeLists.txt files have been severely misformatted with spaces in file paths, which will break the build. Additionally, the PrefixIndex implementation contains unreachable code in the insertion and removal methods, and the recursive DFS functions should be optimized to reduce unnecessary string allocations.

gemini-code-assist · 2026-04-12T02:41:09Z

+void PrefixIndex::dfs_collect_keys(const std::shared_ptr<RadixNode>& node,
+                                   const std::string& accumulated,
+                                   std::vector<std::string>& out) const {
+    if (node->is_terminal()) {
+        out.push_back(accumulated);
+    }
+    for (const auto& [ch, child] : node->children) {
+        dfs_collect_keys(child, accumulated + child->edge_label, out);
+    }


In dfs_collect_keys, the string concatenation accumulated + child->edge_label creates a new std::string object for every node in the trie. For large result sets, this leads to excessive allocations. Consider passing a std::string& buffer and using append/resize to manage the path efficiently during DFS.

References

Avoid unnecessary string allocations in performance-critical paths or recursive functions. ^(link)

Copilot

Pull request overview

This PR adds a new compressed Patricia-style radix tree (PrefixIndex) to Mooncake Store to support efficient prefix-based key enumeration and prefix-scoped aggregations (count/bytes), along with comprehensive unit tests.

Changes:

Introduces PrefixIndex (compressed radix trie) with insert/remove/clear, prefix listing, continuation listing, and subtree count/bytes aggregation.
Adds extensive unit tests covering split/merge edge cases, upsert semantics, cleanup after deletes, and concurrent access.
Wires the new implementation and tests into the Mooncake Store build (library + test target).

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
mooncake-store/include/prefix_index.h	Public API + data structure definitions and semantics for the prefix index
mooncake-store/src/prefix_index.cpp	Implementation of the compressed radix tree with locking and maintained subtree stats
mooncake-store/src/CMakeLists.txt	Adds `prefix_index.cpp` to the store library build
mooncake-store/tests/prefix_index_test.cpp	New unit test suite for correctness, invariants, and concurrency
mooncake-store/tests/CMakeLists.txt	Registers the new `prefix_index_test` target

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-14T06:27:16Z

+ * list_prefix_continuations, count_by_prefix, bytes_by_prefix) take a
+ * shared lock.


The thread-safety comment lists the read operations that take a shared lock, but it omits list_entries_by_prefix() (and size()) even though they are also read-only and use the shared lock. Consider updating the comment to include them for accuracy.

Suggested change

* list_prefix_continuations, count_by_prefix, bytes_by_prefix) take a

* shared lock.

* list_entries_by_prefix, list_prefix_continuations, count_by_prefix,

* bytes_by_prefix, size) take a shared lock.

Copilot · 2026-04-14T06:27:17Z

+void PrefixIndex::dfs_collect_keys(const std::shared_ptr<RadixNode>& node,
+                                   const std::string& accumulated,
+                                   std::vector<std::string>& out) const {
+    if (node->is_terminal()) {
+        out.push_back(accumulated);
+    }
+    for (const auto& [ch, child] : node->children) {
+        dfs_collect_keys(child, accumulated + child->edge_label, out);
+    }
+}
+
+void PrefixIndex::dfs_collect_entries(
+    const std::shared_ptr<RadixNode>& node, const std::string& accumulated,
+    std::vector<std::pair<std::string, TrieLeafData>>& out) const {
+    if (node->is_terminal()) {
+        out.emplace_back(accumulated, node->leaf_data.value());
+    }
+    for (const auto& [ch, child] : node->children) {
+        dfs_collect_entries(child, accumulated + child->edge_label, out);
+    }
+}
+


dfs_collect_keys() builds accumulated + child->edge_label at every recursion, which creates a new std::string (allocation/copy) per edge visited. For large subtrees this can be a noticeable overhead during prefix enumeration; consider using a single mutable buffer (append before recursion, resize/pop after) to avoid the per-node concatenation allocations.

Suggested change

void PrefixIndex::dfs_collect_keys(const std::shared_ptr<RadixNode>& node,

const std::string& accumulated,

std::vector<std::string>& out) const {

if (node->is_terminal()) {

out.push_back(accumulated);

}

for (const auto& [ch, child] : node->children) {

dfs_collect_keys(child, accumulated + child->edge_label, out);

}

}

void PrefixIndex::dfs_collect_entries(

const std::shared_ptr<RadixNode>& node, const std::string& accumulated,

std::vector<std::pair<std::string, TrieLeafData>>& out) const {

if (node->is_terminal()) {

out.emplace_back(accumulated, node->leaf_data.value());

}

for (const auto& [ch, child] : node->children) {

dfs_collect_entries(child, accumulated + child->edge_label, out);

}

}

namespace {

void dfs_collect_keys_impl(const std::shared_ptr<RadixNode>& node,

std::string& accumulated,

std::vector<std::string>& out) {

if (node->is_terminal()) {

out.push_back(accumulated);

}

for (const auto& [ch, child] : node->children) {

const size_t original_size = accumulated.size();

accumulated += child->edge_label;

dfs_collect_keys_impl(child, accumulated, out);

accumulated.resize(original_size);

}

}

void dfs_collect_entries_impl(

const std::shared_ptr<RadixNode>& node, std::string& accumulated,

std::vector<std::pair<std::string, TrieLeafData>>& out) {

if (node->is_terminal()) {

out.emplace_back(accumulated, node->leaf_data.value());

}

for (const auto& [ch, child] : node->children) {

const size_t original_size = accumulated.size();

accumulated += child->edge_label;

dfs_collect_entries_impl(child, accumulated, out);

accumulated.resize(original_size);

}

}

} // namespace

void PrefixIndex::dfs_collect_keys(const std::shared_ptr<RadixNode>& node,

const std::string& accumulated,

std::vector<std::string>& out) const {

std::string buffer = accumulated;

dfs_collect_keys_impl(node, buffer, out);

}

void PrefixIndex::dfs_collect_entries(

const std::shared_ptr<RadixNode>& node, const std::string& accumulated,

std::vector<std::pair<std::string, TrieLeafData>>& out) const {

std::string buffer = accumulated;

dfs_collect_entries_impl(node, buffer, out);

}

Copilot · 2026-04-14T06:27:17Z

+void PrefixIndex::dfs_collect_keys(const std::shared_ptr<RadixNode>& node,
+                                   const std::string& accumulated,
+                                   std::vector<std::string>& out) const {
+    if (node->is_terminal()) {
+        out.push_back(accumulated);
+    }
+    for (const auto& [ch, child] : node->children) {
+        dfs_collect_keys(child, accumulated + child->edge_label, out);
+    }
+}
+
+void PrefixIndex::dfs_collect_entries(
+    const std::shared_ptr<RadixNode>& node, const std::string& accumulated,
+    std::vector<std::pair<std::string, TrieLeafData>>& out) const {
+    if (node->is_terminal()) {
+        out.emplace_back(accumulated, node->leaf_data.value());
+    }
+    for (const auto& [ch, child] : node->children) {
+        dfs_collect_entries(child, accumulated + child->edge_label, out);
+    }
+}
+


dfs_collect_entries() also uses accumulated + child->edge_label on each recursion, which allocates/copies a new string per visited edge. Consider switching to an append/backtrack approach (shared mutable path string) so listing large prefixes doesn't incur O(nodes) intermediate string allocations beyond the final returned keys.

Suggested change

void PrefixIndex::dfs_collect_keys(const std::shared_ptr<RadixNode>& node,

const std::string& accumulated,

std::vector<std::string>& out) const {

if (node->is_terminal()) {

out.push_back(accumulated);

}

for (const auto& [ch, child] : node->children) {

dfs_collect_keys(child, accumulated + child->edge_label, out);

}

}

void PrefixIndex::dfs_collect_entries(

const std::shared_ptr<RadixNode>& node, const std::string& accumulated,

std::vector<std::pair<std::string, TrieLeafData>>& out) const {

if (node->is_terminal()) {

out.emplace_back(accumulated, node->leaf_data.value());

}

for (const auto& [ch, child] : node->children) {

dfs_collect_entries(child, accumulated + child->edge_label, out);

}

}

namespace {

void dfs_collect_keys_impl(

const std::shared_ptr<PrefixIndex::RadixNode>& node, std::string& path,

std::vector<std::string>& out) {

if (node->is_terminal()) {

out.push_back(path);

}

for (const auto& [ch, child] : node->children) {

const size_t original_size = path.size();

path += child->edge_label;

dfs_collect_keys_impl(child, path, out);

path.resize(original_size);

}

}

void dfs_collect_entries_impl(

const std::shared_ptr<PrefixIndex::RadixNode>& node, std::string& path,

std::vector<std::pair<std::string, TrieLeafData>>& out) {

if (node->is_terminal()) {

out.emplace_back(path, node->leaf_data.value());

}

for (const auto& [ch, child] : node->children) {

const size_t original_size = path.size();

path += child->edge_label;

dfs_collect_entries_impl(child, path, out);

path.resize(original_size);

}

}

} // namespace

void PrefixIndex::dfs_collect_keys(const std::shared_ptr<RadixNode>& node,

const std::string& accumulated,

std::vector<std::string>& out) const {

std::string path = accumulated;

dfs_collect_keys_impl(node, path, out);

}

void PrefixIndex::dfs_collect_entries(

const std::shared_ptr<RadixNode>& node, const std::string& accumulated,

std::vector<std::pair<std::string, TrieLeafData>>& out) const {

std::string path = accumulated;

dfs_collect_entries_impl(node, path, out);

}

Copilot · 2026-04-14T06:27:17Z

+     * returned string `s` is a continuation: `prefix + s` is a prefix of
+     * at least one stored key (or equals a stored key when `s` is the
+     * sole terminal continuation). Results are sorted lexicographically.


The list_prefix_continuations() doc says a continuation s may be such that prefix + s “equals a stored key when s is the sole terminal continuation”, which would imply returning an empty-string continuation for an exact terminal-leaf match. The current implementation/tests never return an empty continuation (exact terminal leaf returns an empty vector), so the comment should be updated to match the actual API semantics (continuations are non-empty suffixes).

Suggested change

* returned string `s` is a continuation: `prefix + s` is a prefix of

* at least one stored key (or equals a stored key when `s` is the

* sole terminal continuation). Results are sorted lexicographically.

* returned string `s` is a non-empty continuation: `prefix + s` is

* either a stored key or a prefix of at least one stored key.

* Exact terminal matches do not produce an empty-string continuation;

* they produce no result unless a longer extension also exists.

* Results are sorted lexicographically.

Copilot · 2026-04-14T06:27:18Z

+TEST(PrefixIndexTest, CountByPrefixIsO1) {
+    PrefixIndex idx;
+    for (int i = 0; i < 1000; ++i) {
+        idx.insert("prefix/" + std::to_string(i), 10, kClient1);
+    }
+
+    EXPECT_EQ(idx.count_by_prefix("prefix/"), 1000);
+    EXPECT_EQ(idx.count_by_prefix(""), 1000);
+}


The test name CountByPrefixIsO1 is misleading: PrefixIndex::count_by_prefix() still does an O(prefix_length) traversal before doing an O(1) counter read. Consider renaming the test (and/or its comment) to avoid encoding an incorrect complexity claim.

stmatengss · 2026-04-15T07:27:27Z

Great work! Thanks @maheshrbapatu, will review it ASAP. CC. @ykwd

maheshrbapatu added 2 commits April 11, 2026 22:28

Adding support for prefix queries

ebd11a9

fix styling issues

bab86f5

maheshrbapatu requested review from XucSh, YiXR, stmatengss and ykwd as code owners April 12, 2026 02:34

github-actions bot added run-ci Store labels Apr 12, 2026

maheshrbapatu linked an issue Apr 12, 2026 that may be closed by this pull request

[RFC]: Support Radix Tree Index in Store Master for Prefix Query #1732

Open

1 task

maheshrbapatu removed a link to an issue Apr 12, 2026

[RFC]: Support Radix Tree Index in Store Master for Prefix Query #1732

Open

1 task

gemini-code-assist bot reviewed Apr 12, 2026

View reviewed changes

maheshrbapatu added 3 commits April 11, 2026 22:59

fix cmake

f5f30e8

fix cmake

57b366e

removing dead code

c930022

stmatengss added the run-e2e-ci label Apr 13, 2026

github-actions bot removed the run-e2e-ci label Apr 13, 2026

stmatengss requested a review from Copilot April 14, 2026 06:22

Copilot started reviewing on behalf of stmatengss April 14, 2026 06:22 View session

Copilot AI reviewed Apr 14, 2026

View reviewed changes

ykwd mentioned this pull request Apr 15, 2026

[RFC]: Support Radix Tree Index in Store Master for Prefix Query #1732

Open

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Adding support for prefix queries#1870

Adding support for prefix queries#1870
maheshrbapatu wants to merge 5 commits intokvcache-ai:mainfrom
maheshrbapatu:prefix_support_phase1

maheshrbapatu commented Apr 12, 2026 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist bot Apr 12, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 14, 2026

Uh oh!

Copilot AI Apr 14, 2026

Uh oh!

Copilot AI Apr 14, 2026

Uh oh!

Copilot AI Apr 14, 2026

Uh oh!

Copilot AI Apr 14, 2026

Uh oh!

stmatengss commented Apr 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		* list_prefix_continuations, count_by_prefix, bytes_by_prefix) take a
		* shared lock.

-     * returned string `s` is a continuation: `prefix + s` is a prefix of
-     * at least one stored key (or equals a stored key when `s` is the
-     * sole terminal continuation). Results are sorted lexicographically.
+     * returned string `s` is a non-empty continuation: `prefix + s` is
+     * either a stored key or a prefix of at least one stored key.
+     * Exact terminal matches do not produce an empty-string continuation;
+     * they produce no result unless a longer extension also exists.
+     * Results are sorted lexicographically.

Conversation

maheshrbapatu commented Apr 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Module

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist bot Apr 12, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

stmatengss commented Apr 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

maheshrbapatu commented Apr 12, 2026 •

edited

Loading