Skip to content

perf(fts): resolve deferred row_ids via the cached whole ROW_ID column#7897

Merged
Xuanwo merged 7 commits into
lance-format:mainfrom
LuQQiu:lu/fts_resolve_row_ids_fix
Jul 23, 2026
Merged

perf(fts): resolve deferred row_ids via the cached whole ROW_ID column#7897
Xuanwo merged 7 commits into
lance-format:mainfrom
LuQQiu:lu/fts_resolve_row_ids_fix

Conversation

@LuQQiu

@LuQQiu LuQQiu commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Problem

DeferredDocSet::resolve_row_ids's slow path opens a fresh docs-file reader and issues scattered single-row read_ranges (one d..d+1 range per candidate doc) on every call, caching nothing. Any partition whose DocSet is not materialized (e.g. no prewarm, select-all queries where docs_for_wand loads only num_tokens) pays file-open + random single-row reads per query, forever.

Under concurrency these per-row reads funnel through the shared scheduler/cache locks and effectively serialize the whole search path: in our benchmark the process showed hundreds of threads parked in futex_wait while CPU sat below 50%, and throughput was capped ~4 qps regardless of concurrency or warmup.

Fix

Two changes, addressing the review feedback on both work volume and memory accounting:

  1. Resolve row ids after the global top-k merge. The global heap carries (partition slot, doc_id) for deferred candidates; only the final survivors are resolved — at most limit lookups per query instead of up to partitions × limit, and only partitions holding survivors load their ROW_ID column at all.

  2. Store each partition's ROW_ID column as its own weighed index-cache entry (DocRowIdsKey → CachedDocRowIds), loaded through WeakLanceCache::get_or_insert_with_key: one sequential column read on first use, weighed at insert time so the cache accounts for the memory, independently evictable under pressure, and single-flight deduplicated on concurrent loads. The lazily populated row_ids_col OnceCell (invisible to cache accounting) is removed; no strong copy is retained on the index object.

Memory cost when resident: ~8 bytes/doc per partition (~2.5 MB per 312k-doc partition; 800 MB for a 100M-doc index), now visible to index_cache_size_bytes capacity accounting.

Benchmarks

Single-node bench (320-core host, 100M-doc 42-language dataset on local NVMe, tier-100-200 common words, 5-term match_any, k=100, 960 GiB index cache, distinct query words per query):

variant (c16, 300 s) qps mean latency
before 4.1 3.9 s
whole-column cache only 107.4 149 ms
late resolution + column cache (this PR) 107.7 (26x) 148 ms
late resolution + scattered per-survivor reads, no column cache 19.2 833 ms

The last row A/Bs the "late resolution without whole-column residency" hypothesis: resolving only the ≤k survivors via targeted scattered reads still caps at 19 qps — per-query file opens and random single-row reads dominate regardless of how few rows are read — so the cached whole-column load is what recovers throughput, and late resolution keeps its benefits on the work/memory side (fewer lookups, fewer partitions loading columns).

An ablation that skipped resolution entirely (returning raw doc ids) measured ~155 qps blended on the same run shape, so this change recovers most of the available headroom while returning correct row_ids.

Tests

  • test_bm25_search_many_partitions_resolves_exact_row_ids: 41 partitions (40 matching + 1 whose only token does not match); asserts the exact row_id set, one full-column ROW_ID read per resolving partition, zero scattered single-row reads, zero re-reads on subsequent queries, and no ROW_ID read for the non-matching partition.
  • test_bm25_search_resolves_only_topk_survivors_and_accounts_cache: with k=3 over 40 matching partitions, only the partitions holding final survivors (≤3) load their ROW_ID column, and each loaded column is present in the index cache as its own DocRowIds entry.

Follow-up in this PR: the cache entry becomes the column's only home

The first version of change 2 still had ensure_loaded() (prewarm, masked wand) materialize a full DocSet with its own owned row_ids copy — so a prewarmed partition stored the column twice, and the second copy was invisible to accounting and un-evictable. The follow-up commit removes that copy entirely:

  • DocSet.row_ids becomes Owned(Vec<u64>) | Shared(ScalarBuffer<u64>) (mirroring the existing NumTokens idiom). Shared is a zero-copy view of the cache entry and reports 0 to DeepSizeOf (the entry is weighed at insert).
  • ensure_loaded() — and therefore prewarm — loads the column into the DocRowIdsKey entry and keeps a resident DocSet of num_tokens + the derived reverse-lookup inv only. Nothing long-lived references the entry's allocation, so evicting the entry actually frees the memory; the next borrower reloads it single-flight (index files are immutable, so a reload is always equivalent). This is the same prewarm/cache contract posting lists already have.
  • The masked non-flat wand path (which must check mask.selected(row_id) inside the scoring loop) borrows the entry per query via a Shared view that drops when the query finishes.
  • Flat-shaped masks need inv instead: shape selection now uses the exact should_flat_search predicate Wand::search itself uses, so the two decisions cannot drift (a masked query scored without row_ids would silently skip mask filtering), plus a debug assert in flat_search.
  • Top-k resolution always reads through the entry, so a hot prewarmed column stays recently-used in the LRU instead of looking cold.
  • Legacy and frag-reuse partitions rewrite the mapping at load (sort / tombstone), so they keep their private owned copy and do not populate the raw-column entry. into_builder copies an owned DocSet out of the entry and never stashes it.

Additional tests: test_prewarm_fills_row_ids_cache_entry_without_a_second_copy (prewarm fills every partition's entry, the resident DocSet has no row_ids, and a warm query resolves with zero further docs-file reads) and test_row_ids_resolution_reloads_after_eviction (with a no-retention cache — the always-evicted worst case — every query reloads the column and stays correct, proving nothing depends on a pinned copy).

@github-actions github-actions Bot added A-index Vector index, linalg, tokenizer performance labels Jul 21, 2026
@coderabbitai

coderabbitai Bot commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

DeferredDocSet now caches the full ROW_ID column during deferred resolution. A BM25 regression test verifies exact row IDs, positive scores, and top-k behavior across matching and unrelated partitions.

Changes

Deferred row-id resolution

Layer / File(s) Summary
Cache ROW_ID column during resolution
rust/lance-index/src/scalar/inverted/lazy_docset.rs
resolve_row_ids loads the complete ROW_ID column once and maps deferred document IDs from the cached array.
Validate multi-partition BM25 results and access caching
rust/lance-index/src/scalar/inverted/index.rs
Adds counting wrappers and regression coverage for exact row IDs, positive scores, matching-partition reads, top-k limits, and reuse of cached row IDs.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: xuanwo, bubblecal

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 12.50% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly summarizes the main change: resolving deferred row_ids via the cached ROW_ID column.
Description check ✅ Passed The description directly matches the implemented performance fix, cache behavior, benchmarks, and tests.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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/index.rs`:
- Around line 8510-8612: Strengthen
test_bm25_search_many_partitions_resolves_exact_row_ids by instrumenting
docs-file ROW_ID reads and tracking per-partition access. Assert the first
bm25_search performs exactly one full-column read for each matching partition,
then run the second query and assert it performs no additional ROW_ID reads
while preserving the existing result assertions.
🪄 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: b6c37e72-e010-4bb5-a3b5-c99391c5573f

📥 Commits

Reviewing files that changed from the base of the PR and between 4696387 and 5d1d6b1.

📒 Files selected for processing (2)
  • rust/lance-index/src/scalar/inverted/index.rs
  • rust/lance-index/src/scalar/inverted/lazy_docset.rs

Comment thread rust/lance-index/src/scalar/inverted/index.rs Outdated
@LuQQiu
LuQQiu force-pushed the lu/fts_resolve_row_ids_fix branch from 5d1d6b1 to 84b5cad Compare July 21, 2026 22:41
@LuQQiu

LuQQiu commented Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

summary of the change:

  1. Before: every query re-opened the partition's docs file and did scattered single-row reads to map each candidate's doc_id -> row_id. Nothing was cached, so the same partition paid this cost again on every query.
  2. Now: the ROW_ID column is loaded once per partition (one sequential read) and kept in memory as part of the index entry in the index cache. All later resolves are pure in-memory lookups.

Memory cost: 8 bytes per doc. For a 100M-doc index that is ~800 MB total (~2.5 MB per 312k-doc partition), and it is loaded lazily — only for partitions that actually produce search candidates.

There are some future improvements

  • Index cache size of index partition is calculated at insert time, so row id size may not be correctly calculated if it's loaded lazily, can consider have separate index cache entry.

resolve_row_ids' slow path opened a fresh docs-file reader and issued
scattered single-row read_ranges on every call, caching nothing - so any
partition whose DocSet is not materialized pays a file open plus random
single-row reads per query, forever. Under concurrency these per-row reads
funnel through the shared scheduler/cache locks and serialize the whole
search path.

Use the existing row_ids_column() loader instead: one sequential read of
the ROW_ID column, cached in the row_ids_col OnceCell, after which every
resolve for the partition is an in-memory lookup. Costs ~8 bytes per doc
resident (~2.5MB per 312k-doc partition), the same lazily-populated
accounting semantics as num_tokens_col.

Measured on a 320-core node, 100M-doc 42-language dataset, tier-100-200
5-term match_any, k=100, 960GiB index cache, single-node bench:
- c1: resolve stage 237ms -> 0.1ms; qps 2.5 -> 3.2; mean 406 -> 316ms
- c16, 300s: qps 4.1 -> 107.6 (26x), mean 3.9s -> 149ms

Adds a many-partition regression test asserting exact row_id resolution.
@LuQQiu
LuQQiu force-pushed the lu/fts_resolve_row_ids_fix branch from 84b5cad to 7e2acc6 Compare July 21, 2026 22:46

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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/index.rs`:
- Around line 8523-8525: Update the expected_row_ids initialization in the
matching-partition loop to use Vec::with_capacity with the known upper bound of
NUM_MATCHING_PARTITIONS multiplied by three, while preserving the existing push
and iteration behavior.
🪄 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: 4294d2f0-8f7f-40cb-9d9e-9309bedb9b34

📥 Commits

Reviewing files that changed from the base of the PR and between 84b5cad and 7e2acc6.

📒 Files selected for processing (2)
  • rust/lance-index/src/scalar/inverted/index.rs
  • rust/lance-index/src/scalar/inverted/lazy_docset.rs

Comment thread rust/lance-index/src/scalar/inverted/index.rs Outdated
@LuQQiu
LuQQiu requested review from BubbleCal and Xuanwo July 21, 2026 23:01
…n test

Address review: the exact-row-ids test would also pass with the previous
scattered-read implementation. Wrap the index store with a docs-file
ROW_ID read counter and assert:
- one full-column ROW_ID read per resolving partition on the first query
- zero scattered single-row reads
- zero additional docs-file ROW_ID reads on a subsequent query
- the non-matching partition never reads its ROW_ID column

Also preallocate the expected_row_ids vector to its known bound.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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/index.rs`:
- Around line 8514-8518: Add #[cfg_attr(coverage, coverage(off))] to
DocsRowIdReadCounter and every associated test-only I/O wrapper struct and impl
block in this section, covering the ranges indicated by the comment. Do not
alter production code or the wrappers’ behavior.
🪄 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: 8c4fa95e-c41a-487b-91cd-0c6f7b8bdb54

📥 Commits

Reviewing files that changed from the base of the PR and between 7e2acc6 and 408350f.

📒 Files selected for processing (1)
  • rust/lance-index/src/scalar/inverted/index.rs

Comment thread rust/lance-index/src/scalar/inverted/index.rs
@codecov

codecov Bot commented Jul 21, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 90.73359% with 48 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
rust/lance-index/src/scalar/inverted/index.rs 92.59% 22 Missing and 8 partials ⚠️
...ust/lance-index/src/scalar/inverted/lazy_docset.rs 85.85% 6 Missing and 8 partials ⚠️
rust/lance-index/src/scalar/inverted/wand.rs 71.42% 4 Missing ⚠️

📢 Thoughts on this report? Let us know!

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

8551-8559: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win

Count only true full-column reads.

read_range increments full_column_reads for any projection containing ROW_ID, without checking range. A regression to partial ROW_ID reads could therefore still satisfy full_column_reads == 40, so the test does not fully verify the whole-column caching contract. Track partial reads separately and assert they remain zero; ideally also key counts by docs file to validate the non-matching partition guarantee.

Also applies to: 8773-8785

🤖 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 8551 - 8559,
Update the read-range instrumentation around read_range so full_column_reads
increments only when the requested range represents a true full-column read, not
merely when projection contains ROW_ID. Add separate tracking for partial reads
and assert they remain zero, preferably keyed by docs file, and apply the same
correction to the corresponding logic at the second referenced location.
🤖 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 8551-8559: Update the read-range instrumentation around read_range
so full_column_reads increments only when the requested range represents a true
full-column read, not merely when projection contains ROW_ID. Add separate
tracking for partial reads and assert they remain zero, preferably keyed by docs
file, and apply the same correction to the corresponding logic at the second
referenced location.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: QUIET

Plan: Pro Plus

Run ID: 276df044-e7aa-43ea-ba1c-7b3caa077762

📥 Commits

Reviewing files that changed from the base of the PR and between 408350f and 15f2e44.

📒 Files selected for processing (1)
  • rust/lance-index/src/scalar/inverted/index.rs

@Xuanwo

Xuanwo commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

@LuQQiu I think the repeated-resolution problem is real, and the whole-column result is convincing for the benchmarked serving shape. My concern is the ownership and accounting contract.

The scalar index is weighed when it is inserted, before any row-id column has been initialized. Moka keeps that insertion-time weight, so filling the OnceCell later does not add eviction pressure. The new ~8 bytes/doc also cannot be evicted per partition. For a 100M-doc index that is roughly 800 MB beyond index_cache_size_bytes; with multiple indexes, retained memory can exceed the configured capacity while cache stats still report it as within budget. I do not think we should leave this part as a follow-up, even though the existing num-tokens cache has the same lazy-accounting pattern.

Could we explore either of these directions?

  1. Resolve row IDs after the global top-k merge. Today each matching partition can resolve up to k candidates before the global heap keeps only k, so the sparse path may resolve up to partitions * k row IDs. Carrying (partition, doc_id, score) through the global heap and resolving only the final survivors may remove enough work without introducing whole-column residency.
  2. If whole-column caching is still needed, store each partition's row-id column as an independently weighed cache entry, loaded through WeakLanceCache::get_or_insert_with_key, without retaining a strong copy in the index's OnceCell. That would preserve warm in-memory lookup and single-flight loading while making the memory visible to capacity accounting and independently evictable. A chunked or admission-controlled variant could remain an option if cold rare-term queries over-read too much.

Would the benchmark harness allow a quick A/B of the late-resolution approach? If that does not recover enough throughput, does the partition-scoped cache entry look feasible for this PR?

LuQQiu added 2 commits July 22, 2026 11:02
…-id columns as weighed entries

- Carry (partition slot, doc_id) through the global top-k heap and
  resolve row ids only for the final survivors: at most limit lookups
  per query instead of up to partitions * limit, and only partitions
  holding survivors load their ROW_ID column.
- Store each partition's ROW_ID column as its own index-cache entry
  (DocRowIdsKey -> CachedDocRowIds) loaded through
  WeakLanceCache::get_or_insert_with_key: weighed at insert time so the
  cache accounts for the memory, independently evictable, and
  single-flight deduplicated on concurrent loads. The lazily populated
  OnceCell copy is removed.
The DocRowIdsKey index-cache entry is now the sole owner of a modern
partition's doc_id->row_id column. ensure_loaded (and therefore prewarm)
loads the column into the entry and keeps a resident DocSet of
num_tokens + inv only; the masked non-flat wand path borrows the entry
per query via a zero-copy view; top-k resolution always reads through
the entry so the LRU sees real usage. Evicting the entry now actually
frees the memory, and reloads are single-flight.

Flat-path shape selection shares wand's own should_flat_search
predicate so the two decisions cannot drift, and flat_search asserts it
got a reverse-lookup-capable DocSet. Legacy and frag-reuse partitions
keep their rewritten owned copies and do not populate the raw-column
entry. into_builder copies an owned DocSet out of the entry instead of
materializing one on the partition.
@LuQQiu

LuQQiu commented Jul 22, 2026

Copy link
Copy Markdown
Contributor Author

@Xuanwo Thanks for the careful review — we ended up doing both of your directions, plus one step further.

1. Late resolution (resolve after the global top-k merge) — adopted. The global heap now carries (partition, doc_id) for pending candidates and only the ≤k survivors resolve, grouped per partition. Lookups drop from partitions × k to k, and only survivor-holding partitions load their column.

2. Separate weighed cache entry — adopted, and we A/B'd the "without introducing whole-column residency" hypothesis on the 100M-doc bench: late resolution with scattered per-survivor reads (no column residency) caps at 19–24 qps vs ~108 blended / ~137 warm qps with the cached whole column — per-query file opens + random reads dominate regardless of how few rows are read. So the whole-column load stays, but as a DocRowIdsKey → CachedDocRowIds index-cache entry: weighed at insert, evictable, single-flight.

Follow-up commit (985d7d3): the entry is now the column's only home. ensure_loaded/prewarm no longer materialize a second owned copy inside the DocSet (that copy was the ~800MB invisible to accounting you flagged); the resident set keeps only num_tokens + inv, the masked wand path borrows the entry per query via a zero-copy view, and resolution always reads through the entry so the LRU sees real usage. Evicting the entry genuinely frees the memory; the next query reloads it single-flight. Tests cover: prewarm fills every partition's entry with zero further docs-file reads on warm queries, and a no-retention cache (always-evicted worst case) stays correct.

@Xuanwo Xuanwo left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for working on this!

@Xuanwo
Xuanwo merged commit 8d9c35f into lance-format:main Jul 23, 2026
41 of 48 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

A-index Vector index, linalg, tokenizer performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants