Skip to content

fix(index): reject RabitQ indices whose dimension is not a multiple of 8 - #8869

Open
LuciferYang wants to merge 1 commit into
lance-format:mainfrom
LuciferYang:fix/rq-dist-table-bounds
Open

fix(index): reject RabitQ indices whose dimension is not a multiple of 8#8869
LuciferYang wants to merge 1 commit into
lance-format:mainfrom
LuciferYang:fix/rq-dist-table-bounds

Conversation

@LuciferYang

Copy link
Copy Markdown
Contributor

What this changes

RabitQuantizationStorage::try_from_batch now rejects an index whose rotated_dim is not a multiple of 8, and sum_4bit_dist_table_uninit gains the debug_assert! its own # Safety section already promises.

Why

The FastScan LUT is allocated 4 * code_dim bytes, while the kernels index it as BATCH_SIZE * code_dim.div_ceil(8) with BATCH_SIZE = 32. Those agree only when code_dim % 8 == 0, and then exactly, with no slack. For code_dim = 8q + r with 1 <= r <= 7 the table is 32 - 4r bytes short: 28 at r = 1, 4 at r = 7.

RabitQuantizer::build has rejected a non-multiple since #6024, so no current writer can produce one. But IVF_RQ shipped in #4344 and that gate landed five months later, and try_from_batch validates num_bits and the code byte width without ever checking the dimension. binary_code_bytes() is itself div_ceil(8), so a code_dim = 100 index with 13-byte codes passes the existing width check.

Loading one is not a clean failure. sum_4bit_dist_table_uninit reads the LUT through _mm256_loadu_si256(dist_table.as_ptr().add(i)) on AVX2, the equivalent on AVX-512, and vld1q_u8 on NEON, none of which is length-checked, so all three read past the allocation and return distances computed from whatever follows. Only the scalar fallback panics. This is ApproxMode::Normal, the default query path.

What changes for callers

An index in that state stops loading and returns invalid_input naming the dimension. It used to load and give wrong distances on any host with SIMD, so failing is the better of the two, but it is a behaviour change for anyone holding such an index rather than a pure hardening.

The debug_assert!

sum_4bit_dist_table_uninit's # Safety section lists four obligations and the body checked three. The missing one is the table bound, which is the only one no slice re-checks. Its sibling sum_4bit_hacc_dist_table_uninit has had the equivalent all along.

Test plan

test_try_from_batch_rejects_dim_not_multiple_of_eight builds metadata at code_dim = 12 with a matching 2-byte code column, so it clears the existing width check, and asserts the load fails with a message naming the dimension.

Measured: replacing the new condition with if false takes that test to failed. With the fix, cargo test --profile ci -p lance-index --lib vector::bq is 223 passed and -p lance-linalg --lib simd::dist_table is 9 passed. cargo fmt --all -- --check and cargo clippy -p lance-linalg -p lance-index --all-targets -- -D warnings are clean.

Not covered: the over-read itself. Reproducing it needs a host with SIMD and an index that current code cannot write, so the test pins the rejection rather than the behaviour it prevents.

Provenance

Found while re-checking an old note of mine that claimed the shortfall was a fixed 16 bytes and that the fix belonged at index creation. Both were wrong: the shortfall varies with the dimension, and the creation-time gate already exists. The read path was the part still open.

@github-actions github-actions Bot added A-index Vector index, linalg, tokenizer bug Something isn't working labels Aug 29, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 29, 2026
@LuciferYang
LuciferYang force-pushed the fix/rq-dist-table-bounds branch from 13a991f to 4b62154 Compare August 29, 2026 16:12
@lance-gatekeeper lance-gatekeeper Bot removed the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 29, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 29, 2026
@LuciferYang
LuciferYang force-pushed the fix/rq-dist-table-bounds branch from 4b62154 to d4dfdaf Compare August 29, 2026 16:18
@lance-gatekeeper lance-gatekeeper Bot removed the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 29, 2026

@lance-gatekeeper lance-gatekeeper 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.

Gate recommendation: approve.

The loader rejects byte-misaligned dimensions before any FastScan path can violate its LUT-size safety invariant, while preserving the format emitted by current writers. The regression test also verifies the public InvalidInput failure.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 29, 2026
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 bug Something isn't working K-approved Latest Gatekeeper recommendation permits acceptance.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant