fix(index): window IVF flat merge reads#7761
Conversation
github-actions
Bot
added
A-index
📝 WalkthroughWalkthroughThe non-transformed IVF merge path now estimates row widths, derives configurable partition windows, reads partitions through ChangesWindowed non-transformed merge
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Merge as merge_partial_vector_auxiliary_files
participant Reader as ShardMergeReader
participant Writer as write_partition_batches
participant File as FileWriter
Merge->>Reader: next_partition()
Reader-->>Merge: windowed partition RecordBatches
Merge->>Writer: write partition batches
Writer->>File: write each RecordBatch sequentially
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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: 1
🤖 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 1752-1763: Update the test helper row_ids_from_batches to access
the row-ID column with indexing syntax, replacing
column_by_name(ROW_ID_FIELD.name()) while preserving the existing UInt64
extraction and collection behavior.
🪄 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: ASSERTIVE
Plan: Pro Plus
Run ID: 80dfeef1-20ce-46b6-a125-3232ec6073da
📒 Files selected for processing (1)
rust/lance-index/src/vector/distributed/index_merger.rs
| fn row_ids_from_batches(batches: &[RecordBatch]) -> Vec<u64> { | ||
| batches | ||
| .iter() | ||
| .flat_map(|batch| { | ||
| let row_ids = batch | ||
| .column_by_name(ROW_ID_FIELD.name()) | ||
| .unwrap() | ||
| .as_primitive::<arrow_array::types::UInt64Type>(); | ||
| row_ids.values().iter().copied().collect::<Vec<_>>() | ||
| }) | ||
| .collect() | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n 'use arrow_array' rust/lance-index/src/vector/distributed/index_merger.rs
rg -n 'AsArray' rust/lance-index/src/vector/distributed/index_merger.rsRepository: lance-format/lance
Length of output: 354
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file='rust/lance-index/src/vector/distributed/index_merger.rs'
echo '--- outline ---'
ast-grep outline "$file" --view expanded | sed -n '1,220p'
echo '--- surrounding lines 1700-1785 ---'
sed -n '1700,1785p' "$file"
echo '--- test module import area around 1535-1565 ---'
sed -n '1535,1565p' "$file"Repository: lance-format/lance
Length of output: 7117
Use indexing syntax for this test helper Replace column_by_name() with batch[ROW_ID_FIELD.name()] in row_ids_from_batches to match the test-code guideline.
🤖 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 1752 -
1763, Update the test helper row_ids_from_batches to access the row-ID column
with indexing syntax, replacing column_by_name(ROW_ID_FIELD.name()) while
preserving the existing UInt64 extraction and collection behavior.
Source: Coding guidelines
| // Non-transformed storage (FLAT, SQ, and HNSW variants) can be | ||
| // copied directly, but still needs the windowed reader to avoid | ||
| // one read stream per (partition, shard). | ||
| let partition_window_size = *PARTITION_WINDOW_SIZE; |
There was a problem hiding this comment.
NIT: this window size was chosen for PQ, it might be too much for FLAT, say 4k rows per partition for 768d, it would be 6GiB.
There was a problem hiding this comment.
Addressed in ef28c42. Non-transformed IVF merge now estimates row width from the auxiliary schema and caps the partition window by a byte budget (LANCE_IVF_MERGE_PARTITION_WINDOW_BYTES, default 128MiB), while preserving the existing configured partition-window size as the upper bound. With 4k rows/partition, a 16-byte hash-like payload keeps the 512 window, 768d f32 becomes 10 partitions, and 1024d f32 becomes 8 partitions.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
rust/lance-index/src/vector/distributed/index_merger.rs (2)
1560-1598: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the silent
unwrap_or(1)fallback with an explicit invariant check.
non_transformed_row_width_bytesis only populated whenidx_typeis FLAT/SQ/HNSW_FLAT/HNSW_SQ during the shard-metadata loop (Lines 951-965), and this_ =>arm is only reached for that exact same set of types — sounwrap_or(1)is guarding an effectively unreachable case. If that invariant is ever broken by a future refactor, silently falling back torow_width_bytes = 1would bypass the byte-budget cap entirely (sincetarget_rowswould become enormous) without any signal.🛡️ Suggested fix
+ debug_assert!( + non_transformed_row_width_bytes.is_some(), + "row width estimate must be populated for non-transformed (FLAT/SQ/HNSW) merges" + ); let row_width_bytes = non_transformed_row_width_bytes.unwrap_or(1);As per coding guidelines, "Do not silently guard against impossible conditions; use
debug_assert!, return an explicit error, or remove the check."🤖 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 1560 - 1598, Replace the unwrap_or(1) fallback in the non-transformed merge branch with an explicit invariant check that verifies non_transformed_row_width_bytes is present. Preserve the existing row-width and partition-window calculations for valid FLAT/SQ/HNSW variants, and signal an error or assertion if the invariant is violated instead of silently using 1.Source: Coding guidelines
1560-1598: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDelete
write_partition_rows
rust/lance-index/src/vector/distributed/index_merger.rs:408has no call sites in the crate anymore; the new windowed merge path uses the PQ/RQ helpers directly.🤖 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 1560 - 1598, The unused write_partition_rows function should be deleted from index_merger.rs. Confirm no remaining references in the crate, then remove only that function and leave the windowed merge flow and PQ/RQ helpers unchanged.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/distributed/index_merger.rs`:
- Around line 58-59: The byte-budget constants and
estimate_non_transformed_partition_window_size transformation lack documentation
explaining their semantics and rationale. Add concise doc comments to
DEFAULT_NON_TRANSFORMED_PARTITION_WINDOW_BYTES, the 64-byte fallback width for
non-fixed-width fields, and estimate_non_transformed_partition_window_size,
including how the environment override relates to the default and how the
estimation is used.
- Around line 1856-1888: Refactor
test_estimate_non_transformed_window_size_accounts_for_row_width into an rstest
parameterized test with named #[case::...] entries for the three row-width and
expected-window combinations. Move the differing values into row_width and
expected_window parameters, while keeping the shared lengths, configured window
size, target bytes, and assertion unchanged.
---
Outside diff comments:
In `@rust/lance-index/src/vector/distributed/index_merger.rs`:
- Around line 1560-1598: Replace the unwrap_or(1) fallback in the
non-transformed merge branch with an explicit invariant check that verifies
non_transformed_row_width_bytes is present. Preserve the existing row-width and
partition-window calculations for valid FLAT/SQ/HNSW variants, and signal an
error or assertion if the invariant is violated instead of silently using 1.
- Around line 1560-1598: The unused write_partition_rows function should be
deleted from index_merger.rs. Confirm no remaining references in the crate, then
remove only that function and leave the windowed merge flow and PQ/RQ helpers
unchanged.
🪄 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: ASSERTIVE
Plan: Pro Plus
Run ID: 724f407d-2a25-4cc7-a9b3-c0f6ed005819
📒 Files selected for processing (1)
rust/lance-index/src/vector/distributed/index_merger.rs
| const DEFAULT_NON_TRANSFORMED_PARTITION_WINDOW_BYTES: usize = 128 * 1024 * 1024; | ||
| const NON_TRANSFORMED_PARTITION_WINDOW_BYTES_ENV: &str = "LANCE_IVF_MERGE_PARTITION_WINDOW_BYTES"; |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Document the magic byte budgets and the estimation algorithm.
The 128MiB default budget, the 64-byte fallback width for non-fixed-width fields, and the whole estimate_non_transformed_partition_window_size transformation lack doc comments explaining the rationale/semantics.
📝 Suggested doc comments
+/// Default byte budget for a single windowed partition read when merging
+/// non-transformed (FLAT/SQ/HNSW_FLAT/HNSW_SQ) auxiliary storage. Chosen to
+/// bound peak memory for a window regardless of vector dimensionality.
const DEFAULT_NON_TRANSFORMED_PARTITION_WINDOW_BYTES: usize = 128 * 1024 * 1024;
const NON_TRANSFORMED_PARTITION_WINDOW_BYTES_ENV: &str = "LANCE_IVF_MERGE_PARTITION_WINDOW_BYTES";+/// Estimates the per-row byte width of `schema` for window-sizing purposes.
+/// Falls back to 64 bytes per field for variable-width/unsupported types,
+/// a conservative approximation since exact sizes aren't known ahead of read.
fn estimate_schema_row_width_bytes(schema: &ArrowSchema) -> usize {+/// Computes how many partitions (rows of `accumulated_lengths`) should be
+/// read per window so that `max_partition_rows * row_width_bytes` windows
+/// stay within `target_window_bytes`, capped by `configured_window_size`.
fn estimate_non_transformed_partition_window_size(As per coding guidelines, "Add doc comments to magic constants, thresholds, and non-obvious transformation functions, explaining what the value represents and why it was chosen."
Also applies to: 104-111, 113-132
🤖 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 58 -
59, The byte-budget constants and estimate_non_transformed_partition_window_size
transformation lack documentation explaining their semantics and rationale. Add
concise doc comments to DEFAULT_NON_TRANSFORMED_PARTITION_WINDOW_BYTES, the
64-byte fallback width for non-fixed-width fields, and
estimate_non_transformed_partition_window_size, including how the environment
override relates to the default and how the estimation is used.
Source: Coding guidelines
| #[test] | ||
| fn test_estimate_non_transformed_window_size_accounts_for_row_width() { | ||
| let lengths = vec![4_000_u32; 1_024]; | ||
| let target_window_bytes = 128 * 1024 * 1024; | ||
| let configured_window_size = 512; | ||
|
|
||
| let hash_like_row_width = 16; | ||
| let f32_768_row_width = 8 + 768 * 4; | ||
| let f32_1024_row_width = 8 + 1024 * 4; | ||
|
|
||
| let hash_window = estimate_non_transformed_partition_window_size( | ||
| &lengths, | ||
| hash_like_row_width, | ||
| configured_window_size, | ||
| target_window_bytes, | ||
| ); | ||
| let f32_768_window = estimate_non_transformed_partition_window_size( | ||
| &lengths, | ||
| f32_768_row_width, | ||
| configured_window_size, | ||
| target_window_bytes, | ||
| ); | ||
| let f32_1024_window = estimate_non_transformed_partition_window_size( | ||
| &lengths, | ||
| f32_1024_row_width, | ||
| configured_window_size, | ||
| target_window_bytes, | ||
| ); | ||
|
|
||
| assert_eq!(hash_window, configured_window_size); | ||
| assert_eq!(f32_768_window, 10); | ||
| assert_eq!(f32_1024_window, 8); | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Consider rstest #[case::...] for the three window-size scenarios.
hash_window/f32_768_window/f32_1024_window differ only by row_width_bytes and expected result — a textbook rstest parametrization case.
♻️ Suggested refactor
#[rstest]
#[case::hash_like(16, 512)]
#[case::f32_768d(8 + 768 * 4, 10)]
#[case::f32_1024d(8 + 1024 * 4, 8)]
fn test_estimate_non_transformed_window_size_accounts_for_row_width(
#[case] row_width: usize,
#[case] expected_window: usize,
) {
let lengths = vec![4_000_u32; 1_024];
let window = estimate_non_transformed_partition_window_size(&lengths, row_width, 512, 128 * 1024 * 1024);
assert_eq!(window, expected_window);
}As per coding guidelines, "Use rstest for Rust tests... when cases differ only by inputs; use #[case::{name}(...)] for readable Rust case names."
🤖 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 1856 -
1888, Refactor test_estimate_non_transformed_window_size_accounts_for_row_width
into an rstest parameterized test with named #[case::...] entries for the three
row-width and expected-window combinations. Move the differing values into
row_width and expected_window parameters, while keeping the shared lengths,
configured window size, target bytes, and assertion unchanged.
Source: Coding guidelines
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
2 participants
Fix the IVF merge path for non-PQ/RQ auxiliary files (IVF_FLAT, SQ, and HNSW variants) to reuse the existing ShardMergeReader windowed path.
The previous branch copied these partitions with one read stream per (partition, shard), which scales poorly at high partition counts. This keeps partition-contiguous output while reducing read stream setup to windowed shard reads.
Validated with targeted lance-index merge tests, distributed-indexer cargo check, and clippy without new warnings in the changed path.
Summary by CodeRabbit
Bug Fixes
Tests