Skip to content

fix(index): window IVF flat merge reads #7761

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jackye1995 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(index): window IVF flat merge reads #7761

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
895753a
fix(index): window IVF flat merge reads
jackye1995 Jul 13, 2026
ef28c42
fix(index): size IVF flat merge windows
jackye1995 Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
286 changes: 271 additions & 15 deletions rust/lance-index/src/vector/distributed/index_merger.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ const PARTITION_WINDOW_SIZE_ENV: &str = "LANCE_IVF_PQ_MERGE_PARTITION_WINDOW_SIZ
const DEFAULT_PARTITION_PREFETCH_WINDOW_COUNT: usize = 2;
const PARTITION_PREFETCH_WINDOW_COUNT_ENV: &str =
"LANCE_IVF_PQ_MERGE_PARTITION_PREFETCH_WINDOW_COUNT";
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";
Comment on lines +58 to +59

@coderabbitai coderabbitai Bot Jul 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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

static PARTITION_WINDOW_SIZE: LazyLock<usize> = LazyLock::new(|| {
std::env::var(PARTITION_WINDOW_SIZE_ENV)
.ok()
Expand All @@ -67,6 +69,67 @@ static PARTITION_PREFETCH_WINDOW_COUNT: LazyLock<usize> = LazyLock::new(|| {
.and_then(|v| v.parse::<usize>().ok())
.unwrap_or(DEFAULT_PARTITION_PREFETCH_WINDOW_COUNT)
});
static NON_TRANSFORMED_PARTITION_WINDOW_BYTES: LazyLock<usize> = LazyLock::new(|| {
std::env::var(NON_TRANSFORMED_PARTITION_WINDOW_BYTES_ENV)
.ok()
.and_then(|v| v.parse::<usize>().ok())
.filter(|v| *v > 0)
.unwrap_or(DEFAULT_NON_TRANSFORMED_PARTITION_WINDOW_BYTES)
});

fn fixed_width_bytes(data_type: &DataType) -> Option<usize> {
match data_type {
DataType::Boolean | DataType::Int8 | DataType::UInt8 => Some(1),
DataType::Int16 | DataType::UInt16 | DataType::Float16 => Some(2),
DataType::Int32
| DataType::UInt32
| DataType::Float32
| DataType::Date32
| DataType::Time32(_) => Some(4),
DataType::Int64
| DataType::UInt64
| DataType::Float64
| DataType::Date64
| DataType::Time64(_)
| DataType::Timestamp(_, _)
| DataType::Duration(_) => Some(8),
DataType::FixedSizeBinary(width) => Some(*width as usize),
DataType::FixedSizeList(field, len) => {
fixed_width_bytes(field.data_type()).map(|width| width * *len as usize)
}
_ => None,
}
}

fn estimate_schema_row_width_bytes(schema: &ArrowSchema) -> usize {
schema
.fields()
.iter()
.map(|field| fixed_width_bytes(field.data_type()).unwrap_or(64))
.sum::<usize>()
.max(1)
}

fn estimate_non_transformed_partition_window_size(
accumulated_lengths: &[u32],
row_width_bytes: usize,
configured_window_size: usize,
target_window_bytes: usize,
) -> usize {
let max_window_size = configured_window_size
.max(1)
.min(accumulated_lengths.len().max(1));
let Some(max_partition_rows) = accumulated_lengths.iter().copied().max() else {
return max_window_size;
};
if max_partition_rows == 0 {
return max_window_size;
}

let target_rows = (target_window_bytes as u128 / row_width_bytes.max(1) as u128).max(1);
let partitions = target_rows / max_partition_rows as u128;
partitions.clamp(1, max_window_size as u128) as usize
}

/// Strict bitwise equality check for FixedSizeListArray values.
/// Returns true only if length, value_length and all underlying primitive values are equal.
Expand Down Expand Up @@ -363,6 +426,13 @@ pub async fn write_partition_rows(
Ok(())
}

async fn write_partition_batches(w: &mut FileWriter, batches: Vec<RecordBatch>) -> Result<()> {
for batch in batches {
w.write_batch(&batch).await?;
}
Ok(())
}

/// Transpose the PQ code column for a batch and write it to the unified writer.
///
/// This helper assumes `batch` contains a contiguous range of rows for a single
Expand Down Expand Up @@ -750,6 +820,7 @@ pub async fn merge_partial_vector_auxiliary_files(
// Track per-shard readers, IVF lengths, and precomputed partition offsets.
// This avoids reopening each shard file for every partition during merge.
let mut shard_infos: Vec<ShardInfo> = Vec::new();
let mut non_transformed_row_width_bytes: Option<usize> = None;

progress
.stage_start(
Expand Down Expand Up @@ -877,6 +948,21 @@ pub async fn merge_partial_vector_auxiliary_files(
// Handle logic based on detected index type
let idx_type = detected_index_type
.ok_or_else(|| Error::index("Unable to detect index type".to_string()))?;
if matches!(
idx_type,
SupportedIvfIndexType::IvfFlat
| SupportedIvfIndexType::IvfSq
| SupportedIvfIndexType::IvfHnswFlat
| SupportedIvfIndexType::IvfHnswSq
) {
let schema_arrow: ArrowSchema = reader.schema().as_ref().into();
let row_width = estimate_schema_row_width_bytes(&schema_arrow);
non_transformed_row_width_bytes = Some(
non_transformed_row_width_bytes
.map(|existing| existing.max(row_width))
.unwrap_or(row_width),
);
}

// Compute format version once; defaults to V2_0 if no shards processed yet
let fv = format_version.unwrap_or(LanceFileVersion::V2_0);
Expand Down Expand Up @@ -1471,23 +1557,45 @@ pub async fn merge_partial_vector_auxiliary_files(
}
}
_ => {
for (pid, total_part_len) in accumulated_lengths.iter().copied().enumerate().take(nlist)
{
for shard in shard_infos.iter() {
let part_len = shard.lengths[pid] as usize;
if part_len == 0 {
continue;
}
let offset = shard.partition_offsets[pid];
if let Some(w) = v2w_opt.as_mut() {
write_partition_rows(shard.reader.as_ref(), w, offset..offset + part_len)
.await?;
}
}
if total_part_len == 0 {
// 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 row_width_bytes = non_transformed_row_width_bytes.unwrap_or(1);
let partition_window_size = estimate_non_transformed_partition_window_size(
&accumulated_lengths,
row_width_bytes,
*PARTITION_WINDOW_SIZE,
*NON_TRANSFORMED_PARTITION_WINDOW_BYTES,
);
log::info!(
"Using non-transformed IVF merge partition window size {} (estimated row width: {} bytes, target window bytes: {})",
partition_window_size,
row_width_bytes,
*NON_TRANSFORMED_PARTITION_WINDOW_BYTES
);
let prefetch_window_count = *PARTITION_PREFETCH_WINDOW_COUNT;
let mut shard_merge_reader = ShardMergeReader::new(
shard_infos,
nlist,
partition_window_size,
prefetch_window_count,
);

while let Some((pid, batches)) = shard_merge_reader.next_partition().await? {
if accumulated_lengths[pid] == 0 {
continue;
}
merged_rows = merged_rows.saturating_add(total_part_len as u64);
if batches.is_empty() {
return Err(Error::index(format!(
"No merged batches found for non-empty partition {}",
pid
)));
}

if let Some(w) = v2w_opt.as_mut() {
write_partition_batches(w, batches).await?;
}
merged_rows = merged_rows.saturating_add(accumulated_lengths[pid] as u64);
progress
.stage_progress("merge_partitions", merged_rows)
.await?;
Expand Down Expand Up @@ -1689,6 +1797,154 @@ mod tests {
Ok(total_rows)
}

async fn flat_shard_info(store: &ObjectStore, aux_path: &Path) -> Result<ShardInfo> {
let sched = ScanScheduler::new(
Arc::new(store.clone()),
SchedulerConfig::max_bandwidth(store),
);
let fh = sched
.open_file(aux_path, &CachedFileSize::unknown())
.await?;
let reader = V2Reader::try_open(
fh,
None,
Arc::default(),
&lance_core::cache::LanceCache::no_cache(),
V2ReaderOptions::default(),
)
.await?;

let ivf_idx: u32 = reader
.metadata()
.file_schema
.metadata
.get(IVF_METADATA_KEY)
.unwrap()
.parse()
.unwrap();
let bytes = reader.read_global_buffer(ivf_idx).await?;
let pb_ivf: pb::Ivf = prost::Message::decode(bytes)?;

let mut partition_offsets = Vec::with_capacity(pb_ivf.lengths.len());
let mut running_offset = 0usize;
for length in &pb_ivf.lengths {
partition_offsets.push(running_offset);
running_offset += *length as usize;
}

Ok(ShardInfo {
reader: Arc::new(reader),
lengths: pb_ivf.lengths,
partition_offsets,
total_rows: running_offset,
})
}

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()
}
Comment on lines +1843 to +1854

@coderabbitai coderabbitai Bot Jul 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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.rs

Repository: 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


#[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);
}
Comment on lines +1856 to +1888

@coderabbitai coderabbitai Bot Jul 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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


#[test]
fn test_estimate_schema_row_width_bytes_for_fixed_size_vector() {
let schema = ArrowSchema::new(vec![
(*ROW_ID_FIELD).clone(),
Field::new(
crate::vector::flat::storage::FLAT_COLUMN,
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 768),
true,
),
]);

assert_eq!(estimate_schema_row_width_bytes(&schema), 8 + 768 * 4);
}

#[tokio::test]
async fn test_shard_merge_reader_reads_flat_partitions_in_windows() {
let object_store = ObjectStore::memory();
let index_dir = Path::from("index/windowed_flat");

let partial0 = index_dir.clone().join("partial_0");
let partial1 = index_dir.clone().join("partial_1");
let aux0 = partial0.clone().join(INDEX_AUXILIARY_FILE_NAME);
let aux1 = partial1.clone().join(INDEX_AUXILIARY_FILE_NAME);

let lengths0 = vec![2_u32, 0_u32, 1_u32, 2_u32, 0_u32];
let lengths1 = vec![1_u32, 1_u32, 0_u32, 1_u32, 2_u32];
let dim = 2_i32;

write_flat_partial_aux(&object_store, &aux0, dim, &lengths0, 0, DistanceType::L2)
.await
.unwrap();
write_flat_partial_aux(&object_store, &aux1, dim, &lengths1, 100, DistanceType::L2)
.await
.unwrap();

let shard_infos = vec![
flat_shard_info(&object_store, &aux0).await.unwrap(),
flat_shard_info(&object_store, &aux1).await.unwrap(),
];
let mut reader = ShardMergeReader::new(shard_infos, lengths0.len(), 2, 1);

let mut observed = Vec::new();
while let Some((partition_id, batches)) = reader.next_partition().await.unwrap() {
observed.push((partition_id, row_ids_from_batches(&batches)));
}

assert_eq!(
observed,
vec![
(0, vec![0, 1, 100]),
(1, vec![101]),
(2, vec![2]),
(3, vec![3, 4, 102]),
(4, vec![103, 104]),
]
);
}

#[tokio::test]
async fn test_merge_ivf_flat_success_basic() {
let object_store = ObjectStore::memory();
Expand Down
Loading