fix(index): reuse cached FTS document lengths#7830
Conversation
github-actions
Bot
added
A-index
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: QUIET Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughDocSet now supports owned or shared Arrow-backed token-count storage. Lazy token statistics use a cached snapshot and reuse token-only views, while full row-id materialization and concurrent initialization are covered by updated tests. ChangesDocSet token storage
Lazy token statistics and view reuse
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Caller
participant DeferredDocSet
participant NUM_TOKEN_COL
participant DocSet
Caller->>DeferredDocSet: ensure_num_tokens_loaded()
DeferredDocSet->>NUM_TOKEN_COL: load token counts
NUM_TOKEN_COL-->>DeferredDocSet: Arrow token buffer
DeferredDocSet->>DocSet: create cached token-only view
DocSet-->>Caller: shared DocSet
Caller->>DeferredDocSet: ensure_loaded()
DeferredDocSet->>DocSet: materialize row-id-bearing view
DocSet-->>Caller: full DocSet
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance-index/src/scalar/inverted/lazy_docset.rs`:
- Around line 269-279: The total_tokens initialization path must preserve
fragment-reuse remapping instead of summing raw column values. Update the
closure around total_tokens and the relevant DocSet::from_columns flow to
compute totals through from_columns when a remapper is present, matching the
subsequently loaded full DocSet. Add a regression test that initializes
statistics before loading the full DocSet and verifies the remapped corpus total
remains consistent.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 6d1b15d6-7609-41f1-b98a-7737e7b69542
📒 Files selected for processing (2)
rust/lance-index/src/scalar/inverted/index.rsrust/lance-index/src/scalar/inverted/lazy_docset.rs
| let total_tokens = self | ||
| .total_tokens | ||
| .get_or_try_init(|| async { | ||
| if let Some(full) = self.full.get() { | ||
| return Result::Ok(full.total_tokens_num()); | ||
| } | ||
| let col = self.num_tokens_column().await?; | ||
| Result::Ok(col.values().iter().map(|&n| n as u64).sum()) | ||
| }) | ||
| .await?; | ||
| Ok(*total_tokens) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Preserve remapped totals when statistics initialize first.
For legacy indexes with a fragment-reuse remapper, this caches the raw column sum, while DocSet::from_columns can remove remapped-out rows. The cached corpus total then differs from the subsequently loaded full DocSet, making BM25 statistics initialization-order dependent.
Compute the total through from_columns for that case and add a stats-before-full regression test.
Proposed fix
let col = self.num_tokens_column().await?;
+ if self.is_legacy && self.frag_reuse_index.is_some() {
+ let row_ids = self.row_ids_column().await?;
+ let docs = DocSet::from_columns(
+ row_ids.as_ref(),
+ col.as_ref(),
+ self.is_legacy,
+ self.frag_reuse_index.clone(),
+ )?;
+ return Result::Ok(docs.total_tokens_num());
+ }
Result::Ok(col.values().iter().map(|&n| n as u64).sum())As per coding guidelines, “Every bugfix and feature must have corresponding tests.”
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| let total_tokens = self | |
| .total_tokens | |
| .get_or_try_init(|| async { | |
| if let Some(full) = self.full.get() { | |
| return Result::Ok(full.total_tokens_num()); | |
| } | |
| let col = self.num_tokens_column().await?; | |
| Result::Ok(col.values().iter().map(|&n| n as u64).sum()) | |
| }) | |
| .await?; | |
| Ok(*total_tokens) | |
| let total_tokens = self | |
| .total_tokens | |
| .get_or_try_init(|| async { | |
| if let Some(full) = self.full.get() { | |
| return Result::Ok(full.total_tokens_num()); | |
| } | |
| let col = self.num_tokens_column().await?; | |
| if self.is_legacy && self.frag_reuse_index.is_some() { | |
| let row_ids = self.row_ids_column().await?; | |
| let docs = DocSet::from_columns( | |
| row_ids.as_ref(), | |
| col.as_ref(), | |
| self.is_legacy, | |
| self.frag_reuse_index.clone(), | |
| )?; | |
| return Result::Ok(docs.total_tokens_num()); | |
| } | |
| Result::Ok(col.values().iter().map(|&n| n as u64).sum()) | |
| }) | |
| .await?; | |
| Ok(*total_tokens) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance-index/src/scalar/inverted/lazy_docset.rs` around lines 269 - 279,
The total_tokens initialization path must preserve fragment-reuse remapping
instead of summing raw column values. Update the closure around total_tokens and
the relevant DocSet::from_columns flow to compute totals through from_columns
when a remapper is present, matching the subsequently loaded full DocSet. Add a
regression test that initializes statistics before loading the full DocSet and
verifies the remapped corpus total remains consistent.
Source: Coding guidelines
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance-index/src/scalar/inverted/index.rs (1)
7246-7292: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winAdd retained-memory assertions for shared token storage.
These tests cover sharing and scoring, but not the changed
NumTokens::memory_size/DeepSizeOfbehavior. Add a sliced-buffer test proving the shared Arrow allocation is charged once and owned storage still uses vector capacity.As per coding guidelines, “Every bugfix and feature must have corresponding tests.” <coding_guidelines>
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/scalar/inverted/index.rs` around lines 7246 - 7292, The existing tests for DocSet shared and owned token storage do not verify memory accounting. Extend test_num_tokens_only_reuses_sliced_arrow_storage with assertions that NumTokens::memory_size or DeepSizeOf charges the retained sliced Arrow allocation only once, and extend test_cached_num_tokens_uses_supplied_total_and_full_stays_owned to verify owned storage reports its vector capacity; use the existing storage instances and memory-sizing APIs without changing production behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@rust/lance-index/src/scalar/inverted/index.rs`:
- Around line 7246-7292: The existing tests for DocSet shared and owned token
storage do not verify memory accounting. Extend
test_num_tokens_only_reuses_sliced_arrow_storage with assertions that
NumTokens::memory_size or DeepSizeOf charges the retained sliced Arrow
allocation only once, and extend
test_cached_num_tokens_uses_supplied_total_and_full_stays_owned to verify owned
storage reports its vector capacity; use the existing storage instances and
memory-sizing APIs without changing production behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: f7a93baf-3fdd-4666-9087-2ae0ca64ff03
📒 Files selected for processing (2)
rust/lance-index/src/scalar/inverted/index.rsrust/lance-index/src/scalar/inverted/lazy_docset.rs
## Why FTS corpus stats already read and cache each partition's Arrow document-length column and compute its total token count. The first WAND match then copied that entire column into a new Vec and summed it again while building the num-tokens-only DocSet. On large partitions this introduced O(rows) anonymous-page work for every PE/index partition even when no index parts were loaded. ## What Make the num-tokens-only DocSet retain a shared ScalarBuffer view and accept the previously cached total. Full DocSet materialization remains owned for masks, fragment reuse, remapping, and row-id semantics. Total initialization is single-flight, and retained-memory accounting deduplicates the shared Arrow storage. This removes the confirmed first-touch copy and repeated sum without changing WAND behavior or broadening prewarm or cache scope. (cherry picked from commit aed29d5)
2 participants
Why
FTS corpus stats already read and cache each partition's Arrow document-length column and compute its total token count. The first WAND match then copied that entire column into a new Vec and summed it again while building the num-tokens-only DocSet. On large partitions this introduced O(rows) anonymous-page work for every PE/index partition even when no index parts were loaded.
What
Make the num-tokens-only DocSet retain a shared ScalarBuffer view and accept the previously cached total. Full DocSet materialization remains owned for masks, fragment reuse, remapping, and row-id semantics. Total initialization is single-flight, and retained-memory accounting deduplicates the shared Arrow storage.
This removes the confirmed first-touch copy and repeated sum without changing WAND behavior or broadening prewarm or cache scope.