From da4f0957356323c18a537051562aece8e2aa1707 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 14 Sep 2026 08:00:57 +0800 Subject: [PATCH 1/2] perf(index): counting sort for shuffle interleave indices sort_to_interleave_indices materialized a 12-byte (partition_id, batch_idx, row_idx) tuple per row and ran an O(n log n) comparison sort, although partition ids are already bounded to [0, num_partitions). Replace with an O(n + num_partitions) counting sort that also produces the per-partition counts as a byproduct, removing the per-row key allocation entirely from the shuffle critical path. Assisted-by: GLM-5.3 --- rust/lance-index/src/vector/v3/shuffler.rs | 46 ++++++++++++++-------- 1 file changed, 30 insertions(+), 16 deletions(-) diff --git a/rust/lance-index/src/vector/v3/shuffler.rs b/rust/lance-index/src/vector/v3/shuffler.rs index ee156961d24..a8652026b33 100644 --- a/rust/lance-index/src/vector/v3/shuffler.rs +++ b/rust/lance-index/src/vector/v3/shuffler.rs @@ -584,28 +584,42 @@ fn sort_to_interleave_indices( num_partitions: usize, ) -> Result { let total_rows: usize = part_id_columns.iter().map(|a| a.len()).sum(); - let mut keys: Vec<(u32, u32, u32)> = Vec::with_capacity(total_rows); - for (batch_idx, col) in part_id_columns.iter().enumerate() { - let batch_idx = batch_idx as u32; - for (row_idx, &part_id) in col.values().iter().enumerate() { - keys.push((part_id, batch_idx, row_idx as u32)); + + // Counting sort: partition ids are bounded to [0, num_partitions), so + // bucket by id in O(n + num_partitions) instead of the previous + // O(n log n) comparison sort over 12-byte key tuples. + let mut partition_counts = vec![0u64; num_partitions]; + for col in part_id_columns { + for &part_id in col.values().iter() { + let pid = part_id as usize; + if pid >= num_partitions { + return Err(Error::invalid_input(format!( + "partition ID {} is out of range [0, {})", + pid, num_partitions + ))); + } + partition_counts[pid] += 1; } } - keys.sort_unstable_by_key(|k| k.0); - let mut partition_counts = vec![0u64; num_partitions]; + // Prefix sums give the starting slot for each partition bucket. + let mut offsets = vec![0u32; num_partitions]; + let mut running = 0u32; + for (offset, &count) in offsets.iter_mut().zip(partition_counts.iter()) { + *offset = running; + running += count as u32; + } + let mut interleave_indices = Vec::with_capacity(total_rows); - for (part_id, batch_idx, row_idx) in &keys { - let pid = *part_id as usize; - if pid >= num_partitions { - return Err(Error::invalid_input(format!( - "partition ID {} is out of range [0, {})", - pid, num_partitions - ))); + interleave_indices.resize(total_rows, (0usize, 0usize)); + for (batch_idx, col) in part_id_columns.iter().enumerate() { + for (row_idx, &part_id) in col.values().iter().enumerate() { + let slot = offsets[part_id as usize] as usize; + offsets[part_id as usize] += 1; + interleave_indices[slot] = (batch_idx, row_idx); } - partition_counts[pid] += 1; - interleave_indices.push((*batch_idx as usize, *row_idx as usize)); } + Ok((interleave_indices, partition_counts)) } From cbd1602405ece701e3ccc2f5519895b95f48608e Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 18 Sep 2026 19:52:37 +0800 Subject: [PATCH 2/2] perf(index): bucket shuffle rows by partition id instead of sorting keys --- rust/lance-index/src/vector/v3/shuffler.rs | 33 ++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/rust/lance-index/src/vector/v3/shuffler.rs b/rust/lance-index/src/vector/v3/shuffler.rs index 235ac5039fe..b0c67493ed2 100644 --- a/rust/lance-index/src/vector/v3/shuffler.rs +++ b/rust/lance-index/src/vector/v3/shuffler.rs @@ -2864,6 +2864,39 @@ mod tests { assert_eq!(p1.num_rows(), 130); } + /// Two non-empty batches in one flush group: the bucketing has to keep each + /// row with its own batch, so a partition's rows come out batch by batch and + /// in row order within a batch. A wrong bucket offset would move values into + /// the neighbouring partition while leaving every partition size intact. + #[tokio::test] + async fn test_two_file_shuffler_groups_two_batches_by_partition() { + let dir = TempStrDir::default(); + let output_dir = Path::from(dir.as_ref()); + + let batch1 = make_batch(&[1, 0, 2], &[10, 20, 30], None); + let batch2 = make_batch(&[2, 1, 0], &[40, 50, 60], None); + + let shuffler = TwoFileShuffler::new(output_dir, 3); + let reader = shuffler + .shuffle(batches_to_stream(vec![batch1, batch2])) + .await + .unwrap(); + + let expected = [vec![20, 60], vec![10, 50], vec![30, 40]]; + for (partition_id, expected_values) in expected.iter().enumerate() { + assert_eq!(reader.partition_size(partition_id).unwrap(), 2); + let partition = collect_partition(reader.as_ref(), partition_id) + .await + .unwrap(); + let values: &Int32Array = partition["val"].as_primitive(); + assert_eq!( + values.values(), + expected_values, + "partition {partition_id} holds the wrong rows" + ); + } + } + /// Nullable `__ivf_part_id` must not be treated as partition 0 via `values()`. #[tokio::test] async fn test_two_file_shuffler_rejects_null_partition_ids() {