perf(vector): stop rebuilding the quantizer and copying masks per query by dshepelev15 · Pull Request #10829 · lance-format/lance

dshepelev15 · 2026-11-07T11:54:01Z

Summary

Allocation profiling of an IVF_PQ search with a with_row_id_prefilter mask (a counting global allocator plus perf stat page-fault, dTLB and L1 counters) found work that every query repeated although nothing about it changes between queries:

The quantizer was rebuilt from the cached state on every open. open_vector_index hits the IvfIndexState cache and calls reconstruct, which ran Q::from_metadata again: for a PQ codebook that widens the f16 codebook to f32 and builds the 48 prepared L2 targets, about 1.5 MB per query. IvfIndexState now carries the quantizer as a runtime-only Option<Quantizer>, excluded from the cache codec like rq_search_cache. It is set when the state is built (taken from the storage on export, built from the metadata on decode), so the cache charges for it on admission and a reconstruction cannot grow the entry; only storages that retain a quantizer (PQ) carry it. IvfQuantizationStorage::from_cached_with_quantizer takes the cached quantizer.
The external mask was deep-copied twice per query. MaskAndLoader::load returned the mask by value, so Arc::unwrap_or_clone copied it while the plan still held the Arc, and DatasetPreFilter::wait_for_ready cloned both the loaded mask and the deletion mask before intersecting. FilterLoader::load now returns Arc<RowAddrMask> (the trait is documented as internal; MaskAndLoader hands out its shared Arc, other loaders wrap a fresh mask), RowAddrMask::intersect_ref (with RowAddrTreeMap::intersection / difference) builds the combined mask from references, and a lone mask is shared with its cache entry.

Per query on the same 4M-row dataset as #10828 (fsl<f16,768>, IVF_PQ 128 × 48, k=10, nprobes=20, one client, measured with the scratch pool also cached; the quantizer is the larger part):

	before	after
plain: allocations / bytes	7 849 / 11.7 MB	6 489 / 4.3 MB
plain: p50	5.3 ms	4.4 ms
mask (5% allowlist): allocations / bytes	19 717 / 15.7 MB	17 828 / 9.0 MB
mask: p50	6.4 ms	6.0 ms

The remaining mask-side allocations are the per-run roaring intersections inside selected_indices; walking the bitmap's range instead allocates nothing but measured slower, so they stay.

Tests

test_state_entry_carries_quantizer_before_admission: an exported and a decoded PQ state entry report the same deep_size_of after a reconstruction, so the cache's admission charge covers everything the entry retains.
test_intersect_ref_matches_intersect: the by-reference intersection equals the consuming one for allow/allow (including an empty fragment and a whole fragment), allow/block and block/block masks.
Existing MaskAndLoader, DatasetPreFilter and prefilter loader tests cover the shared return.

cargo fmt --all and cargo clippy --all --tests --benches -- -D warnings are clean.

Opening a vector index from its cached IvfIndexState reconstructed the index object for every query: Q::from_metadata ran again (widening an f16 PQ codebook and preparing its L2 targets, about 1.5 MB per query) and a fresh QueryScratchPool was created whose scratches then grew on first use. The state now holds the quantizer and the scratch pool in runtime-only cells, excluded from the cache codec like rq_search_cache, so later reconstructions share them. With an external row-id mask the mask was deep-copied twice per query: MaskAndLoader::load returned it by value while the plan still held the Arc, and DatasetPreFilter cloned both the loaded mask and the deletion mask before intersecting. FilterLoader::load now returns the Arc, RowAddrMask::intersect_ref combines masks from references, and a lone mask is shared with its cache entry. Per query (4M rows, fsl<f16,768>, IVF_PQ 128 x 48, k=10, nprobes=20): plain 7 849 -> 6 489 allocations and 11.7 -> 4.3 MB, p50 5.3 -> 4.4 ms; 5% allowlist mask 19 717 -> 17 828 allocations and 15.7 -> 9.0 MB, p50 6.4 -> 6.0 ms.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

FilterLoader is public, so its load signature stays; a provided load_shared hands out the mask behind an Arc and MaskAndLoader overrides it to share the Arc it already holds. The cached quantizer is only built for storages that keep one, so an RQ index does not carry its rotation matrices in the cache for nothing. The intersect_ref test now pins an empty intersection and a Full/Partial pair.

The quantizer is now built when a state entry is exported or decoded, so the cache charges for it on admission and a reconstruction can no longer grow the entry. The scratch pool is no longer kept in the state: its buffers grow with use after admission, which the cache would not see.

dshepelev15 · 2026-11-07T12:58:34Z

Addressed the accounting note: the quantizer cell is now filled when the state entry is exported or decoded, before the entry enters the cache, so its size is part of the admission charge, and the scratch pool is no longer kept in the state since its buffers grow after admission. test_state_entry_carries_quantizer_before_admission checks that both an exported and a decoded entry keep the same deep_size_of across a reconstruction. The FilterLoader::load signature is unchanged (load_shared is a provided method), so no breaking-change label should be needed.

wjones127

Nice profiling-driven PR: caching the PQ quantizer in the index cache entry and intersecting masks by reference are both clear wins, and the cache-accounting issue from the first revision is fixed.

Requesting changes for one thing: the new RowAddrTreeMap::intersection/difference sit on every prefiltered query, and their whole-fragment branches (including the deleted-fragments path) aren't covered by the parity test. The rest is non-blocking: making the quantizer field a plain Option, tightening the state-entry test, reusing the storage's quantizer on export, and collapsing load/load_shared into one method.

wjones127 · 2026-11-11T20:11:49Z

+        disjoint.insert_fragment(1);
+        disjoint.insert(2 << 32 | 1);
+        let disjoint = RowAddrMask::from_allowed(disjoint);
+        for (a, b) in [


issue(blocking): None of these pairs reach the whole-fragment branches of the new difference/intersection: an allow list minus a block list that holds a whole fragment ((_, Some(Full)) => return None), a whole allowed fragment minus a partial block, a whole allowed fragment with no block entry, or Full ∩ Full. The first one is the common case. DatasetPreFilter builds deleted_fragments with insert_fragment, so any prefiltered query on a dataset with a deleted fragment goes through it. Could you add pairs that cover those four branches? Done when each match arm in intersection and difference is hit by at least one pair compared against intersect.

Done in da9cdb7. test_intersect_ref_matches_intersect now builds the two sides fragment by fragment so that every arm is hit: whole/whole, whole/partial, partial/whole, partial/partial with and without rows left, and a fragment absent from the other side, for both intersection and difference (including allow minus a block list that holds a whole fragment, a whole allowed fragment minus a partial block, and a whole allowed fragment with no block entry). Each pair is compared against intersect, and the two allow-list results are also spelled out so the implementations cannot agree on a wrong answer.

wjones127 · 2026-11-11T20:11:49Z

+    /// reconstruction of this entry. It is filled before the entry enters the
+    /// cache (on export and on decode) so the cache charges for it, and is
+    /// excluded from the CacheCodec wire format.
+    pub(crate) quantizer: Arc<OnceLock<Option<Quantizer>>>,


suggestion: Both constructors (export and decode) fill this cell eagerly, so the get_or_init fallback in reconstruct_typed is only reachable from the partition_serde test fixture. What if this were just quantizer: Option<Quantizer>? Then "filled before admission" holds by construction, rather than depending on every future constructor remembering to fill a lazy cell. A lazy fill after admission is exactly the accounting bug the earlier revision had. deep_size_of becomes .as_ref(), and reconstruct_typed becomes state.quantizer.clone().

Done in da9cdb7: the field is quantizer: Option<Quantizer>, set when the state is built on export and on decode. reconstruct_typed just clones it and the get_or_init fallback is gone.

wjones127 · 2026-11-11T20:11:49Z

+            )
+            .await
+            .unwrap();
+        assert_eq!(entry.deep_size_of(), charged);


issue(non-blocking): This passes trivially if the quantizer is never built. retained_quantizer swallows from_metadata errors with .ok(), so if that returned None for both entries, deep_size_of would be unchanged and the test would still pass. Could you also assert that the quantizer is actually there? For example, check that the cell is Some, or that charged exceeds the size of the same state without a quantizer.

Good catch. da9cdb7 asserts entry.0.has_quantizer() for both the exported and the decoded entry before the size check. has_quantizer is a #[cfg(test)] method on IvfStateEntry, since the test only holds the type-erased entry.

wjones127 · 2026-11-11T20:11:49Z

            index_file_size: self.reader.metadata().file_size(),
            aux_file_size: self.storage.reader().metadata().file_size(),
            rq_search_cache: rabit_search_cache_cell(self.rq_search_cache.clone()),
+            quantizer: Arc::new(OnceLock::from(retained_quantizer::<Q>(


question: For PQ, self.storage already retains a quantizer built from this same metadata. Is there a reason to rebuild it here with retained_quantizer instead of cloning the storage's copy? Reusing it saves one codebook widening plus the prepared targets on every cache export.

No reason, it was an oversight. da9cdb7 takes the storage's copy through self.storage.quantizer() when the storage retains one, so export no longer widens the codebook or prepares the targets again. retained_quantizer is now used only on decode, where there is no storage yet.

wjones127 · 2026-11-11T20:11:49Z

+
+    /// [`Self::load`] behind an `Arc`. A loader that already holds a shared
+    /// mask overrides this to hand it out without a copy.
+    async fn load_shared(self: Box<Self>) -> Result<Arc<RowAddrMask>> {


suggestion: Every production caller now uses load_shared, and load is only reached from a test and from load_shared's default body. Since FilterLoader is documented as internal with no stability guarantees, could we collapse these into one method, async fn load(self: Box<Self>) -> Result<Arc<RowAddrMask>>? Loaders that build a fresh mask wrap it in Arc::new, and MaskAndLoader hands out its shared one. That removes the duplicated combine logic here and leaves one path to implement.

Done in da9cdb7: FilterLoader::load returns Result<Arc<RowAddrMask>> and load_shared is removed. The loaders that build a fresh mask wrap it in Arc::new, MaskAndLoader hands out its shared one. I had kept load unchanged only to avoid touching a public trait signature; if this needs the breaking-change label, I do not have permission to apply it.

wjones127 · 2026-11-11T20:11:49Z

+            .filter_map(|(fragment, lhs)| {
+                let selection = match (lhs, other.inner.get(fragment)?) {
+                    (RowAddrSelection::Full, RowAddrSelection::Full) => RowAddrSelection::Full,
+                    (RowAddrSelection::Full, RowAddrSelection::Partial(set))


nitpick: Full ∩ Partial(set) keeps the fragment even when set is empty. BitAndAssign drops empty partials in its final retain, so the two can disagree under PartialEq, and is_empty() on the result reads false. This only matters if an empty Partial can reach here, but an if set.is_empty() { return None; } like the partial/partial arm has would keep them identical.

Fixed in da9cdb7. intersection now drops a fragment whose result is an empty partial in one place after the match, so it agrees with BitAndAssign whichever side the empty set came from. The new test has a whole fragment against an empty partial to cover it.

wjones127 · 2026-11-11T20:11:49Z

+                + deleted.is_some() as usize
+                + deleted_fragments.is_some() as usize
+                + self.overlay_block.is_some() as usize;
+            if count <= 1 {


suggestion: The count sum and the two expect("at least two masks") calls restate the same thing twice. Pulling two items from parts up front (let (Some(first), Some(second)) = (parts.next(), parts.next()) else { ...lone/default... }) and then folding the rest would drop both the counter and the expects. Only worth doing if the borrow checker cooperates.

Done in da9cdb7, the borrow checker cooperated once every mask is an Arc: overlay_block is stored as Option<Arc<RowAddrMask>> (wrapped once in with_overlay_block), the first two masks are pulled with let Some(..) = masks.next() else, and the rest are folded. The counter and both expect calls are gone, and a lone overlay mask is no longer copied.

- Cover every arm of `RowAddrTreeMap::intersection` and `difference` against the consuming `intersect`, and spell out the expected masks. - `intersection` drops a fragment left without rows whichever side the empty set came from, as `BitAndAssign` does. - `FilterLoader::load` returns `Arc<RowAddrMask>`; `load_shared` is gone. - `IvfIndexState::quantizer` is a plain `Option<Quantizer>`, taken from the storage on export instead of being rebuilt. - The state test asserts that the quantizer is present. - `DatasetPreFilter` folds its masks without a counter or `expect`.

lance-gatekeeper

✅ Gate recommendation: approve.

Retained PQ data is present before cache admission, and the updated tests cover the whole-fragment mask combinations. No blocking findings remain.

Please mark this PR with the breaking-change label: lance_index::prefilter::FilterLoader::load now returns Result<Arc<RowAddrMask>>, so Rust implementations and callers must use the shared return type.

wjones127

Thanks for working through the last round so thoroughly. The parity test now reaches every arm of intersection/difference, including the deleted-fragment block list, and spells out the expected masks. The quantizer is a plain field set at construction, export reuses the storage's copy, and load/load_shared are one method. Those were all my earlier points.

Blocking (1): the FilterLoader::load signature change doesn't compile against current main, because PrefilterUnionExec (#10001) landed after this branch was cut. Rebasing and updating that call site should clear the CI failures.

Encouraged: find a way to share the set-algebra logic between the in-place and by-reference mask operations without giving up performance, plus a couple of small simplifications inline.

Re the gatekeeper's breaking-change ask: FilterLoader is documented as internal, so no label is needed.

wjones127 · 2026-11-12T23:00:35Z

-    async fn load(self: Box<Self>) -> Result<RowAddrMask>;
+    /// The mask is shared so that a loader which already holds one behind an
+    /// `Arc` (a cached mask, say) hands it out without a copy.
+    async fn load(self: Box<Self>) -> Result<Arc<RowAddrMask>>;


issue(blocking): This signature change doesn't compile against current main. PrefilterUnionExec (added in #10001 after this branch was cut) still treats load() as returning an owned mask:

lance/rust/lance/src/io/exec/prefilter_union.rs

Lines 277 to 286 in 0599360

let missing_future = async {

match missing {

Some(stream) => {

Box::new(FilteredRowIdsToPrefilter::new(stream).with_metrics(&metrics, 0))

.load()

.await

}

None => Ok(RowAddrMask::allow_nothing()),

}

};

That's the E0308 at L284 and the follow-on E0599 at L297 behind the linux-build/clippy/rustdoc/MSRV/JNI/Python failures. Could you rebase and update that caller? The None arm can wrap allow_nothing() in Arc::new, and the scalar | missing below can take Arc::unwrap_or_clone(missing), which is free since FilteredRowIdsToPrefilter hands back a fresh Arc. Done when the branch is rebased on main and cargo check -p lance --tests passes.

wjones127 · 2026-11-12T23:00:35Z


+impl RowAddrTreeMap {
+    /// Rows held by both maps, allocating only the result.
+    pub fn intersection(&self, other: &Self) -> Self {


suggestion: RowAddrTreeMap now has two implementations of the same set algebra: these by-reference methods, and the in-place BitAndAssign<&Self> / SubAssign<&Self>. Only the parity test keeps them in sync, and the empty-partial fix last round shows how easily they drift.

I'd like to keep this path fast, but sharing the logic matters just as much. Is there a way to get both? For example, could both forms go through one per-fragment helper that decides what each pairing (full/partial/empty on each side) produces, with the in-place and by-reference versions differing only in whether they write into self or build a new map? Or, if one form can be written in terms of the other without losing the allocation savings, that would be even simpler. Same question for RowAddrMask::intersect vs intersect_ref.

wjones127 · 2026-11-12T23:00:35Z

+    /// `intersect_ref` builds the same mask as the consuming `intersect`, for
+    /// every pairing of whole, partial, empty and absent fragments.
+    #[test]
+    fn test_intersect_ref_matches_intersect() {


praise: Nice test. Laying out each fragment's pairing in a comment, then spelling out the expected masks as well as comparing against intersect, makes it easy to check that every arm is reached and that both sides aren't wrong together.

wjones127 · 2026-11-12T23:00:36Z

            rq_search_cache: rabit_search_cache_cell(self.rq_search_cache.clone()),
+            // The storage's own copy, when it retains one: building another
+            // would widen the codebook and prepare its targets a second time.
+            quantizer: Q::retain_quantizer_for_partition_storage()


suggestion: This repeats the retain_quantizer_for_partition_storage() gate that the storage already applies, and storage.quantizer() would fall back to from_metadata if the storage's copy were absent. What if IvfQuantizationStorage exposed fn retained_quantizer(&self) -> Option<&Quantizer> { self.quantizer.as_ref() }? Then this would be quantizer: self.storage.retained_quantizer().cloned(), with no gate and no chance of a rebuild on export.

wjones127 · 2026-11-12T23:00:36Z

+    /// the same metadata, so reconstructing an index does not rebuild it (and
+    /// its prepared targets) from the metadata each time.
+    #[doc(hidden)]
+    pub fn from_cached_with_quantizer(


suggestion: That's now three from_cached* constructors, each forwarding to the next with one more argument. Since from_cached_with_remapper is #[doc(hidden)], could it just take the quantizer: Option<Quantizer> parameter, with existing callers passing None? That would leave one hidden constructor instead of two.

wjones127 · 2026-11-12T23:00:36Z

+        let allow = RowAddrMask::from_allowed(lhs);
+        let other_allow = RowAddrMask::from_allowed(rhs.clone());
+        let block = RowAddrMask::from_block(rhs);
+        let other_block = RowAddrMask::from_block(other_block);


nitpick: other_block is first a tree (L1519) and then shadowed by the mask built from it. Naming the tree block_rows (or building the mask in place) would make the pairs loop easier to read.

wjones127 · 2026-11-12T23:00:36Z

-            }
+                Arc::new(RowAddrMask::from_block(block_list))
+            });
+            let mut masks = [


praise: This reads much better than the old accumulate-and-clone version. Generalizing the deletions-only shortcut to any single mask falls out naturally, and the count and the expects are gone.

claude Bot reviewed Nov 7, 2026

View reviewed changes

github-actions Bot added A-index Vector index, linalg, tokenizer performance labels Nov 7, 2026

This comment was marked as outdated.

Sign in to view

lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Nov 7, 2026

lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Nov 7, 2026

This comment was marked as outdated.

Sign in to view

lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Nov 7, 2026

lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Nov 7, 2026

This comment was marked as outdated.

Sign in to view

lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Nov 7, 2026

wjones127 requested changes Nov 11, 2026

View reviewed changes

lance-gatekeeper Bot removed the K-approved Latest Gatekeeper recommendation permits acceptance. label Nov 12, 2026

lance-gatekeeper Bot reviewed Nov 12, 2026

View reviewed changes

lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Nov 12, 2026

wjones127 requested changes Nov 12, 2026

View reviewed changes

	let missing_future = async {
	match missing {
	Some(stream) => {
	Box::new(FilteredRowIdsToPrefilter::new(stream).with_metrics(&metrics, 0))
	.load()
	.await
	}
	None => Ok(RowAddrMask::allow_nothing()),
	}
	};

Conversation

dshepelev15 commented Nov 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Tests

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

dshepelev15 commented Nov 7, 2026

Uh oh!

This comment was marked as outdated.

Uh oh!

wjones127 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lance-gatekeeper Bot left a comment

Choose a reason for hiding this comment

Uh oh!

wjones127 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

dshepelev15 commented Nov 7, 2026 •

edited

Loading