feat(index): inject globally-trained SQ bounds through the build params#7901
feat(index): inject globally-trained SQ bounds through the build params#7901xuanyu-z wants to merge 4 commits into
Conversation
Distributed (fragment-sharded) builds of SQ-quantized indexes (IVF_SQ,
IVF_HNSW_SQ) currently let each shard's `ScalarQuantizer::build`
self-train its own min/max bounds from its own sample, so per-shard
codes end up on different scales and the segment merge silently
produces incorrect quantized values (all merged codes are dequantized
against the first shard's bounds).
Mirror `PQBuildParams::with_codebook`: add an optional
`bounds: Option<Range<f64>>` to `SQBuildParams` plus a
`SQBuildParams::with_bounds(num_bits, bounds)` constructor. When set,
`ScalarQuantizer::build` constructs the quantizer directly from the
injected bounds and skips the per-shard `update_bounds` training scan,
so every shard quantizes on an identical scale. The field rides the
existing `quantizer_params -> load_or_build_quantizer -> Q::build`
channel; no builder plumbing changes.
Also harden the segment merger: both SQ paths (IVF_SQ and IVF_HNSW_SQ)
previously captured only the first shard's `ScalarQuantizationMetadata`
and validated only `dim`. They now return an error ("SQ bounds mismatch
across shards") when later shards carry different bounds, turning
silent corruption into a hard failure.
Backward compatible: `bounds` defaults to `None`, which keeps the
existing self-training behavior, and the on-disk format is unchanged
(bounds were already persisted in `ScalarQuantizationMetadata`; the
proto still carries only `num_bits`).
Relates to lance-format#6309.
github-actions
Bot
added
A-index
📝 WalkthroughWalkthroughScalar quantization now accepts injected bounds, uses them during quantizer construction, and validates matching SQ metadata when merging IVF-SQ and IVF-HNSW-SQ shards. Tests cover shared-bound merges and mismatched-bound failures. ChangesScalar quantization bounds
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant SQBuildParams
participant ScalarQuantizer
participant Shards
participant IndexMerger
SQBuildParams->>ScalarQuantizer: supply optional bounds
ScalarQuantizer->>Shards: build SQ shard indexes
Shards->>IndexMerger: submit shard metadata
IndexMerger->>IndexMerger: compare SQ metadata
IndexMerger-->>Shards: merge indexes or return mismatch error
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 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: 6
🤖 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/vector/distributed/index_merger.rs`:
- Around line 186-194: Update ensure_sq_bounds_match to validate
candidate.num_bits against reference.num_bits in addition to bounds. Return an
index error on either mismatch that includes both shards’ bounds and bit-width
values, with variable names and relevant types, while preserving the successful
Ok path when metadata matches.
In `@rust/lance-index/src/vector/sq.rs`:
- Around line 148-153: Validate cloned bounds in the params.bounds branch before
calling Self::with_bounds: require both endpoints to be finite and bounds.start
strictly less than bounds.end. For invalid values, return Error::invalid_input
with both bound values included in the message, and add rejection tests covering
empty, reversed, infinite, and NaN bounds.
In `@rust/lance-index/src/vector/sq/builder.rs`:
- Around line 14-19: Document the public SQBuildParams bounds API, including
SQBuildParams::with_bounds, with a compiling usage example that matches the
current method signature. Link the documentation to the relevant quantizer build
method and associated structs, and keep the example synchronized with the actual
API.
In `@rust/lance/src/index/vector/details.rs`:
- Around line 1314-1318: Update the SQBuildParams fixture in
test_vector_index_details_roundtrip to use explicit non-default bounds via
Some(...), then extend the restored SQ-stage assertions to verify the exact
bounds range is preserved through protobuf/details round-trip.
In `@rust/lance/src/index/vector/ivf/v2.rs`:
- Around line 4115-4120: Update the regression assertion around the SQ bounds
mismatch error to verify that the error is specifically the Error::index
variant, while retaining the existing “SQ bounds mismatch across shards” message
check. Ensure both the variant and message content must match before the test
succeeds.
- Around line 4063-4065: Update the HNSW graph construction assertion in the
surrounding test to validate recall against ground_truth, not merely that
ids_split contains K results. Compute or use the existing recall comparison and
require recall to be at least 0.5, while preserving the current result-count
handling where applicable.
🪄 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: 1fe2ffb1-3e99-48c1-9197-711679e003c9
📒 Files selected for processing (7)
java/lance-jni/src/utils.rsrust/lance-index/src/vector/distributed/index_merger.rsrust/lance-index/src/vector/sq.rsrust/lance-index/src/vector/sq/builder.rsrust/lance/src/index/vector.rsrust/lance/src/index/vector/details.rsrust/lance/src/index/vector/ivf/v2.rs
…e, harden tests - Reject non-finite (NaN/±inf) injected bounds and non-float element types on the injected-bounds fast path (previously deferred to quantize time); `start == end` stays valid (a globally constant column). Tests added. - Merge guard now compares `num_bits` as well as `bounds` (`ensure_sq_metadata_compatible`), so shards on different scales can't merge even with matching bounds (matters for future SQ4); error names both values. - IVF_HNSW_SQ merge test asserts a K/3 top-k overlap floor vs the single-index baseline (the file's existing non-flaky idiom) instead of only length. - Test cleanup: reuse the existing `build_segments_for_fragment_groups` helper in the merge test; extract an `f32_fsl` fixture for the SQ build tests; use 8 bits (not the unsupported 4) in the params test.
There was a problem hiding this comment.
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance/src/index/vector/ivf/v2.rs (1)
4001-4010: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winCover
num_bitsin the merge tests.The new merge contract validates both
boundsandnum_bits, but this test only asserts bounds and only creates a bounds mismatch. Assertsq_meta.num_bits == 8and add a differing-bit-width case, or directly unit-test the compatibility helper.As per coding guidelines, “Every bugfix and feature must have corresponding tests.”
🤖 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/src/index/vector/ivf/v2.rs` around lines 4001 - 4010, The merge test around ds_single and ds_split only validates SQ bounds; extend coverage to the full compatibility contract by asserting sq_meta.num_bits equals 8 and adding a merge case with differing bit widths, or add a focused unit test for the compatibility helper that verifies rejection. Keep the existing bounds-mismatch coverage intact.Source: Coding guidelines
♻️ Duplicate comments (3)
rust/lance-index/src/vector/sq.rs (1)
163-168: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReject reversed injected bounds.
The fast path accepts ranges such as
2.0..1.0and stores them unchanged. Equal endpoints are intentionally supported, but valid bounds must satisfystart <= end; otherwise invalid metadata can reach quantization and shard merging.Proposed fix
- if !bounds.start.is_finite() || !bounds.end.is_finite() { + if !bounds.start.is_finite() + || !bounds.end.is_finite() + || bounds.start > bounds.end + { return Err(Error::invalid_input(format!( - "SQ builder: injected bounds must be finite, got {}..{}", + "SQ builder: invalid injected bounds: start={}, end={}; expected finite start <= end", bounds.start, bounds.end ))); }As per coding guidelines, “Validate inputs at API boundaries and reject invalid values with descriptive errors; never silently clamp or adjust them.”
🤖 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/vector/sq.rs` around lines 163 - 168, Update the injected-bounds validation in the SQ builder to reject finite ranges where bounds.start is greater than bounds.end, returning a descriptive invalid-input error that includes both values. Preserve acceptance of equal endpoints and the existing non-finite validation behavior.Source: Coding guidelines
rust/lance/src/index/vector/ivf/v2.rs (2)
4060-4072: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winRequire at least 0.5 HNSW-SQ overlap.
With
K = 10,overlap >= K / 3accepts only3/10 = 0.3overlap, below the required 0.5 recall floor.Proposed fix
- overlap >= K / 3, + overlap as f32 / K as f32 >= 0.5,As per coding guidelines, “Include ... vector-index recall assertions of at least 0.5.”
🤖 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/src/index/vector/ivf/v2.rs` around lines 4060 - 4072, Update the overlap assertion in the single-vs-split HNSW-SQ comparison to enforce at least 0.5 recall, replacing the K / 3 threshold with a threshold that guarantees half of K while preserving the existing overlap calculation and diagnostic message.Source: Coding guidelines
4122-4127: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAssert the error variant as well.
This regression test checks only the message substring. Capture the error and assert the concrete index-error variant in addition to the
"SQ metadata mismatch across shards"message.As per coding guidelines, “Assert on both the error variant and the message content in tests; do not check only
is_err().”🤖 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/src/index/vector/ivf/v2.rs` around lines 4122 - 4127, Update the regression test around the existing error assertion to inspect the captured error’s concrete index-error variant in addition to checking that its message contains “SQ metadata mismatch across shards.” Preserve the current failure message while replacing the message-only validation with assertions covering both the variant and message content.Source: Coding guidelines
🟡 Other comments (1)
rust/lance-index/src/vector/sq.rs-421-451 (1)
421-451: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAssert the concrete error contract in these tests.
The rejection tests only call
is_err(), so they do not verify theInvalidInputvariant or message content. Capture each error and assert both the concrete variant and descriptive message, including for unsupported element types.As per coding guidelines, “Assert on both the error variant and the message content in tests; do not check only
is_err().”🤖 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/vector/sq.rs` around lines 421 - 451, The tests test_build_rejects_non_finite_injected_bounds and test_build_with_injected_bounds_rejects_unsupported_type should validate the concrete InvalidInput error contract instead of only calling is_err(). Capture each ScalarQuantizer::build error, assert it matches the InvalidInput variant, and verify the message contains the descriptive text for non-finite bounds and unsupported element types.Source: Coding guidelines
🧹 Nitpick comments (1)
rust/lance/src/index/vector/ivf/v2.rs (1)
3922-3937: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse an enum for the SQ index kind.
make_sq_index_paramsuses arbitrary strings as a variant discriminator. PassIndexType::IvfSq/IndexType::IvfHnswSqinstead and match exhaustively; update the rstest cases and comparisons accordingly.As per coding guidelines, “Use enums instead of magic numbers for format versions, variant types, and discriminators, and rely on exhaustive
match.”🤖 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/src/index/vector/ivf/v2.rs` around lines 3922 - 3937, Update make_sq_index_params to accept IndexType instead of &str, matching exhaustively on IndexType::IvfSq and IndexType::IvfHnswSq and removing the unexpected-string panic branch. Update all rstest cases, callers, and related comparisons to pass and compare the enum variants directly.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/src/index/vector/ivf/v2.rs`:
- Around line 4001-4010: The merge test around ds_single and ds_split only
validates SQ bounds; extend coverage to the full compatibility contract by
asserting sq_meta.num_bits equals 8 and adding a merge case with differing bit
widths, or add a focused unit test for the compatibility helper that verifies
rejection. Keep the existing bounds-mismatch coverage intact.
---
Other comments:
In `@rust/lance-index/src/vector/sq.rs`:
- Around line 421-451: The tests test_build_rejects_non_finite_injected_bounds
and test_build_with_injected_bounds_rejects_unsupported_type should validate the
concrete InvalidInput error contract instead of only calling is_err(). Capture
each ScalarQuantizer::build error, assert it matches the InvalidInput variant,
and verify the message contains the descriptive text for non-finite bounds and
unsupported element types.
---
Duplicate comments:
In `@rust/lance-index/src/vector/sq.rs`:
- Around line 163-168: Update the injected-bounds validation in the SQ builder
to reject finite ranges where bounds.start is greater than bounds.end, returning
a descriptive invalid-input error that includes both values. Preserve acceptance
of equal endpoints and the existing non-finite validation behavior.
In `@rust/lance/src/index/vector/ivf/v2.rs`:
- Around line 4060-4072: Update the overlap assertion in the single-vs-split
HNSW-SQ comparison to enforce at least 0.5 recall, replacing the K / 3 threshold
with a threshold that guarantees half of K while preserving the existing overlap
calculation and diagnostic message.
- Around line 4122-4127: Update the regression test around the existing error
assertion to inspect the captured error’s concrete index-error variant in
addition to checking that its message contains “SQ metadata mismatch across
shards.” Preserve the current failure message while replacing the message-only
validation with assertions covering both the variant and message content.
---
Nitpick comments:
In `@rust/lance/src/index/vector/ivf/v2.rs`:
- Around line 3922-3937: Update make_sq_index_params to accept IndexType instead
of &str, matching exhaustively on IndexType::IvfSq and IndexType::IvfHnswSq and
removing the unexpected-string panic branch. Update all rstest cases, callers,
and related comparisons to pass and compare the enum variants directly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 846465e2-7724-4c2d-a9e6-15ac92c97c4e
📒 Files selected for processing (3)
rust/lance-index/src/vector/distributed/index_merger.rsrust/lance-index/src/vector/sq.rsrust/lance/src/index/vector/ivf/v2.rs
…act tests, robust HNSW recall - Document that injected bounds are in the quantizer's POST-TRANSFORM space: cosine indexes normalize first, so cosine bounds cover the normalized components (-1..1), not the raw values -- wrong-space bounds clip codes and silently collapse the index. (Mechanism is space-agnostic; this is a caller contract.) Add a compiling doc example on `with_bounds`. - Reject reversed injected bounds (start > end) too, alongside NaN/inf; start == end stays valid. - Error-path tests now assert the exact variant (InvalidInput / Index) plus message, not only is_err()/substring. - IVF_HNSW_SQ merge test: aggregate top-k overlap across all queries and check the K/3 recall floor (the file's existing single-vs-split SQ idiom); the build is non-deterministic, so a per-query threshold flaked -- verified stable across repeated runs.
There was a problem hiding this comment.
Actionable comments posted: 1
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
🟡 Other comments (1)
rust/lance-index/src/vector/sq/builder.rs-19-24 (1)
19-24: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocument the endpoint validity contract.
The build path rejects non-finite or reversed bounds and permits equal endpoints. State those valid-value rules here so callers can validate inputs before build time.
As per coding guidelines, “Add doc comments to public API elements that convey semantic meaning, valid values, and effects.”
🤖 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/vector/sq/builder.rs` around lines 19 - 24, Update the public `bounds` field documentation to state that bounds must have finite endpoints, may use equal lower and upper endpoints, and must not be reversed. Keep the existing post-transform-space and clipping behavior documentation intact.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.
Inline comments:
In `@rust/lance-index/src/vector/sq/builder.rs`:
- Around line 14-24: The public SQBuildParams change adding bounds breaks
existing exhaustive struct literals. Preserve source compatibility for prior
SQBuildParams literals, or coordinate this as an intentional breaking API
change; do not rely on with_bounds, and follow the project’s
deprecation-and-replacement convention if an API transition is required.
---
Other comments:
In `@rust/lance-index/src/vector/sq/builder.rs`:
- Around line 19-24: Update the public `bounds` field documentation to state
that bounds must have finite endpoints, may use equal lower and upper endpoints,
and must not be reversed. Keep the existing post-transform-space and clipping
behavior documentation intact.
🪄 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: 9d538b88-b40d-415c-801e-7869d6cf8e5d
📒 Files selected for processing (3)
rust/lance-index/src/vector/sq.rsrust/lance-index/src/vector/sq/builder.rsrust/lance/src/index/vector/ivf/v2.rs
| /// User-provided, globally-trained quantization bounds. When set, per-build | ||
| /// training is skipped and these exact bounds are used, so codes are | ||
| /// comparable across independently built shards (mirrors | ||
| /// `PQBuildParams::codebook`). | ||
| /// | ||
| /// Must be the global min/max in the quantizer's **post-transform** space: | ||
| /// cosine indexes L2-normalize first, so cosine bounds cover the normalized | ||
| /// components (`-1.0..=1.0`), not the raw values; L2/dot use the raw | ||
| /// min/max. Bounds that don't cover that space clip every code to 0 and | ||
| /// silently collapse the index. | ||
| pub bounds: Option<std::ops::Range<f64>>, |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Preserve compatibility for existing SQBuildParams literals.
Adding bounds makes downstream exhaustive literals fail to compile; with_bounds does not preserve callers that only set the prior public fields. Keep this source-compatible or release it as a coordinated breaking API change.
As per coding guidelines, “Do not break public API signatures; deprecate old APIs with #[deprecated] or @deprecated and add a replacement.”
🤖 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/vector/sq/builder.rs` around lines 14 - 24, The public
SQBuildParams change adding bounds breaks existing exhaustive struct literals.
Preserve source compatibility for prior SQBuildParams literals, or coordinate
this as an intentional breaking API change; do not rely on with_bounds, and
follow the project’s deprecation-and-replacement convention if an API transition
is required.
Source: Coding guidelines
…s merge guard - Aggregate IVF_HNSW_SQ recall floor raised K/3 -> 0.5 (the repo minimum). The aggregate measures ~0.71-0.79 across runs with tight variance, so 0.5 is a stable, non-flaky floor (the earlier K/3 was over-anchored on a single noisy per-query sample). - Add a focused unit test for `ensure_sq_metadata_compatible` covering the num_bits mismatch path (identical bounds, different num_bits -> rejected), which the integration test alone didn't exercise; also assert num_bits == 8 in the successful merge test.
There was a problem hiding this comment.
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance/src/index/vector/ivf/v2.rs (1)
4128-4136: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winMatch
errorby reference.matches!(error, lance_core::Error::Index { .. })consumeserror, so the laterto_string()call on the same binding won’t compile. Usematches!(&error, ...).🤖 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/src/index/vector/ivf/v2.rs` around lines 4128 - 4136, Update the first assertion’s matches! call to match error by reference using &error, preserving ownership so the subsequent error.to_string() assertion remains compilable.
🟡 Other comments (1)
rust/lance-index/src/vector/distributed/index_merger.rs-1679-1692 (1)
1679-1692: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAssert the error variant and diagnostic content.
The
is_err()checks can pass for any error. Capture each failure and assertError::Indexplus the"SQ metadata mismatch across shards"message, including the relevantnum_bits/boundscontext.As per coding guidelines, “Assert on both the error variant and the message content in tests; do not check only
is_err().”🤖 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/vector/distributed/index_merger.rs` around lines 1679 - 1692, Update ensure_sq_metadata_compatible_checks_bounds_and_num_bits to capture each incompatible-metadata result and assert it is the Error::Index variant with the “SQ metadata mismatch across shards” diagnostic, including the relevant num_bits and bounds values for each failure case. Replace the generic is_err() assertions while preserving the existing compatible-metadata check.Source: Coding guidelines
🧹 Nitpick comments (1)
rust/lance/src/index/vector/ivf/v2.rs (1)
3922-3939: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
IndexTypefor the SQ test cases instead of strings. PassIndexType::IvfSq/IndexType::IvfHnswSqintomake_sq_index_params, and switch the laterindex_kind == "IVF_SQ"check to an enum match so the variants stay typed consistently.🤖 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/src/index/vector/ivf/v2.rs` around lines 3922 - 3939, Update make_sq_index_params to accept IndexType instead of &str and match on IndexType::IvfSq and IndexType::IvfHnswSq. Change all SQ test call sites to pass those enum variants, and replace the later index_kind == "IVF_SQ" comparison with an enum-based match while preserving the existing 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/src/index/vector/ivf/v2.rs`:
- Around line 4128-4136: Update the first assertion’s matches! call to match
error by reference using &error, preserving ownership so the subsequent
error.to_string() assertion remains compilable.
---
Other comments:
In `@rust/lance-index/src/vector/distributed/index_merger.rs`:
- Around line 1679-1692: Update
ensure_sq_metadata_compatible_checks_bounds_and_num_bits to capture each
incompatible-metadata result and assert it is the Error::Index variant with the
“SQ metadata mismatch across shards” diagnostic, including the relevant num_bits
and bounds values for each failure case. Replace the generic is_err() assertions
while preserving the existing compatible-metadata check.
---
Nitpick comments:
In `@rust/lance/src/index/vector/ivf/v2.rs`:
- Around line 3922-3939: Update make_sq_index_params to accept IndexType instead
of &str and match on IndexType::IvfSq and IndexType::IvfHnswSq. Change all SQ
test call sites to pass those enum variants, and replace the later index_kind ==
"IVF_SQ" comparison with an enum-based match while preserving the existing
behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: c1f5644c-eb1d-4515-9791-7ca72ab431dd
📒 Files selected for processing (2)
rust/lance-index/src/vector/distributed/index_merger.rsrust/lance/src/index/vector/ivf/v2.rs
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
Not sure if the left over coderabbit comment need to be addressed, need a maintainer here |
1 participant
Problem
A distributed (fragment-sharded) build of an SQ-quantized index (
IVF_SQ,IVF_HNSW_SQ) lets each shard'sScalarQuantizer::buildself-train its ownmin/maxbounds from its own sample. SQ codes are only meaningful relative to the bounds they were trained with, so per-shard codes end up on different scales. When the shards' segments are merged, the merger keeps only the first shard'sScalarQuantizationMetadataand dequantizes every shard's codes against it — silently producing wrong quantized values and wrong search results, with no error.Fix
Mirror
PQBuildParams::with_codebook: give SQ the same "inject a globally-trained quantizer" hook so every shard quantizes on an identical scale.bounds: Option<Range<f64>>toSQBuildParams, plusSQBuildParams::with_bounds(num_bits, bounds).ScalarQuantizer::buildconstructs the quantizer directly from the injected bounds (via the pre-existingScalarQuantizer::with_bounds) and skips the per-shard training scan. The field rides the existingquantizer_params → load_or_build_quantizer → Q::buildchannel — no builder plumbing changes.IVF_SQ,IVF_HNSW_SQ) previously validated onlydim. They now reject shards whose SQ metadata differs —num_bitsandbounds(ensure_sq_metadata_compatible) — turning silent corruption into a hard error.Why exact-equality (not the tolerance used for PQ codebooks / centroids)
ensure_fixed_size_list_compatiblecompares PQ codebooks and IVF centroids with a float tolerance, because those can drift within training. SQnum_bits/boundsare different: they're injected verbatim into every shard, so any difference is not benign drift — it means a shard used a different scale and its codes are incomparable. Exact inequality is exactly the check we want; a tolerance would let a divergent scale slip through.Compatibility
boundsdefaults toNone→ existing self-training behavior is byte-for-byte the same; bounds were already persisted inScalarQuantizationMetadata, so no format/proto change.SQBuildParams, mirroring exactly howcodebook: Option<ArrayRef>was added toPQBuildParams(also a plain public struct). Code constructingSQBuildParamsviaDefault/..Default::default()is unaffected; exhaustive struct literals (SQBuildParams { num_bits, sample_rate }) need the new field — the same source-level break thePQBuildParams::codebookaddition carried. Flagging so it can be labeled/documented per the repo's convention if needed.Out of scope (caller's responsibility)
This adds the mechanism. Computing the global bounds — a coordinator-side pass over all data to find the true global
min/maxbefore sharding — is the caller's job, exactly as computing a global PQ codebook is today. Injecting bounds that don't cover all data would clip out-of-range values (scale_to_u8already saturates rather than panicking). Part of #6309 (Distributed Indexes Search).Tests
test_merge_sq_segments_with_injected_bounds(IVF_SQ, IVF_HNSW_SQ): a sharded-then-merged index carries exactly the injected bounds; for the deterministic IVF_SQ case it returns byte-identical Top-K to a single-shot baseline built with the same bounds.test_merge_sq_segments_rejects_mismatched_bounds: merging shards with different bounds errors.test_build_with_injected_bounds/test_build_without_bounds_self_trains: the build short-circuit and theNonedefault.test_build_rejects_non_finite_injected_bounds(NaN, ±inf;start == endstays valid) andtest_build_with_injected_bounds_rejects_unsupported_type(Int32): input validation on the fast path.Follow-ups (out of scope here)
bounds; JNI setsbounds: None. So distributed SQ builds are Rust-only for now — a binding user's independently-trained segments will be rejected by the new merger guard (a hard error, not silent corruption). Exposingsq_boundsthrough both bindings is a follow-up.load_or_build_quantizerstill samples vectors beforebuildsees the injected bounds; only the final min/max pass is skipped. A pre-trained quantizer could take dim/type from the schema and skip sampling entirely — a perf follow-up for distributed builds.num_bitsvalidation at the build boundary (SQ only encodes u8 today; SQ4 is a TODO).