fix(index): narrow f16/f64 training chunks to f32 in streaming IVF trainers - #8996
Open
LuciferYang wants to merge 2 commits into
Open
fix(index): narrow f16/f64 training chunks to f32 in streaming IVF trainers#8996LuciferYang wants to merge 2 commits into
LuciferYang wants to merge 2 commits into
Conversation
…ainers The streaming refine and coreset trainers accumulate and re-dispatch in f32, but the chunk normalization called convert_to_floating_point, which only widens integer types and returns f16/f64 inputs unchanged. A float16 or float64 column with streaming training and more than 256 partitions then hit the unconditional Float32 downcast inside the trainers and panicked the build. Route the three streaming normalization sites through one helper that casts to Float32 explicitly. An f64 above f32::MAX saturates to an infinity rather than failing the cast, so the helper re-runs the samplers' finite filter after that narrowing. The num_partitions <= 256 streaming path dispatches on value type already and is untouched.
…ning The streaming IVF training tests await the whole training stack. Its future type nests past the default limit of 128 when the lib test target is built with the full feature set, which is what CI does; benches/streaming_ivf_training.rs already raises the limit for the same stack. Scoped to test builds so the library's own limit is unchanged.
Contributor
There was a problem hiding this comment.
✅ Gate recommendation: approve.
This restores the streaming trainer’s intended Float32 boundary at the three f32-only coreset/refinement consumers while keeping native normalization and the ≤256 typed path unchanged. Post-cast overflow filtering preserves the finite-input invariant, and the regression coverage exercises both sampler routes and raw refinement.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The three streaming IVF normalization sites called
convert_to_floating_point, which returns f16 and f64 unchanged and only widens integer types. Those columns reached the trainers' unconditional Float32 downcast and panicked the build.Fixes #8995.
What this changes
All three sites now go through one helper that casts the chunk to Float32. Two details worth calling out.
An f64 above
f32::MAXsaturates to an infinity rather than failing the cast. The samplers run their finite filter before the cast (ivf.rs:3193 and :3518), so without something else the trainers could receive infinities, andlance-index/src/vector/kmeans.rs:840has no finite guard: the centroids would come out inf or NaN with no error.drop_rows_that_saturatedremoves exactly the rows the narrowing broke and logs how many, so a null row still passes through the way it does for every other value type.The helper takes the array apart with
into_partsrather than rebuilding the field by hand, which keeps the inner field's metadata and drops an unreachable branch. Its accept list is f16/f32/f64/i8, matching what a vector column can actually hold (rust/lance/src/index/vector/utils.rs:244).The
num_partitions <= 256streaming path dispatches on value type already and is untouched.Note for whoever reviews #8610: that PR moves this same
convert_to_floating_pointcall into a new helper without changing it, so it carries the panic forward. Landing this first, or folding it in, avoids that.Test plan
Five unit tests on the helper: the dtype matrix compared value by value, a sliced chunk mapping to the rows it points at, nulls staying on their row, an f64 that overflows f32 being dropped while a null row beside it survives, and an unsupported type returning
InvalidInputwith the type named. Removing the saturation filter makes the overflow test report[[1.0, 2.0], [inf, 4.0], [5.0, 6.0]].Two end-to-end cases train a 257 partition model with streaming enabled, one f16 and one f64. They take different routes on purpose: a non-nullable column gets the fixed-range sampler, a nullable one the resampling path, and each has its own normalization site, so the two cases cover two of the three sites rather than differing only in dtype. Restoring the old normalization at the coreset site makes both panic in arrow's cast.
cargo test -p lance --lib3387 passed, 3 ignoredcargo test -p lance --lib index::vector::ivf::tests::51 passedcargo clippy --all --tests --benches -- -D warningscleancargo fmt --all --checkclean