Skip to content

fix(index): validate the width and dtype of pre-computed IVF centroids - #9366

Open
LuciferYang wants to merge 3 commits into
lance-format:mainfrom
LuciferYang:fix/ivf-rq-centroids-width-validation
Open

LuciferYang wants to merge 3 commits into
lance-format:mainfrom
LuciferYang:fix/ivf-rq-centroids-width-validation

Conversation

@LuciferYang

Copy link
Copy Markdown
Contributor

Closes #9365

The pre-computed-centroids branch of build_ivf_model checked only the flattened length, so 64 centroids of width 8 passed as 32 of width 16 and broke the build somewhere well past this boundary, and f16 centroids over an f32 column passed into an assignment pair that does not exist for that combination. It now checks the per-row width and the value type where it already checks the length, and both errors name the centroids' value alongside the column's.

The dtype check is not an equality on the column's element type, which is the trap here. An Int8 column is trained as f32 — train_ivf_kmeans_step takes (DataType::Int8, L2 | Dot | Cosine) through convert_to_floating_point — and the only assignment pair the kernels implement for such a column is (DataType::Float32, DataType::Int8); there is no (Int8, Int8) arm. So f32 centroids are the expected ones for an Int8 column, and the check compares against that expectation rather than against the column type. Java always builds centroids as f32 from a float[], which is exactly this case.

How was this patch tested?

test_build_ivf_model_rejects_wrong_width_centroids covers both rejections against a width-16 f32 column: 64×8 centroids, whose flattened length matches 32×16 and which the length check alone lets through, and 32×16 f16 centroids. Each case asserts its own message, since a substring shared by both would let the width case pass on the dtype error.

test_build_ivf_model_accepts_f32_centroids_for_int8_column is the other side: f32 centroids over an Int8 column build a model rather than being rejected. Without the Int8 mapping above, a strict equality check fails it — that combination is the only one Int8 columns can use.

build_ivf_model checked only the flattened length of supplied centroids,
so wrongly-shaped arrays passed validation: a half-width array with
twice the rows broke the build far from this boundary, and same-width
wrong-dtype centroids (e.g. f16 over an f32 column) panicked in a blind
downcast inside kmeans. Reject both at the boundary with the offending
values.

Assisted-by: GLM-5.3
@github-actions github-actions Bot added the bug Something isn't working label Sep 18, 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 fresh-build boundary now rejects supplied IVF centroids whose row width or assignment-compatible value type disagrees with the vector column. The Int8-to-Float32 exception matches the existing assignment kernels, while persisted index readers and on-disk formats remain unchanged.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Sep 18, 2026
@LuciferYang

Copy link
Copy Markdown
Contributor Author

The one red check is unrelated: linux-arm fails in index::vector::ivf::v2::tests::test_legacy_ivf_pq_cosine_multivec_smoke on a recall assertion, recall: 0.48 against the 0.5 threshold at rust/lance/src/index/vector/ivf/v2.rs:6698. That test builds its index through the normal training path and passes no pre-computed centroids, so it never reaches either branch this PR adds. The same test has flaked the same way on unrelated PRs (#9302, #9307).

Happy to rerun the job if you would rather see a clean board.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

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.

bug: pre-computed IVF centroids are only length-checked, so wrong-shaped ones break the build later

1 participant