diff --git a/rust/lance-index/src/vector/v3/shuffler.rs b/rust/lance-index/src/vector/v3/shuffler.rs index 7b08260d0f9..b0c67493ed2 100644 --- a/rust/lance-index/src/vector/v3/shuffler.rs +++ b/rust/lance-index/src/vector/v3/shuffler.rs @@ -596,9 +596,11 @@ 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); + + // Partition ids are bounded by `num_partitions`, so counting sort orders the + // rows in O(n + num_partitions) passes instead of sorting n key tuples. + let mut partition_counts = vec![0u64; num_partitions]; for (batch_idx, col) in part_id_columns.iter().enumerate() { - let batch_idx = batch_idx as u32; for row_idx in 0..col.len() { if col.is_null(row_idx) { return Err(Error::invalid_input(format!( @@ -606,24 +608,35 @@ fn sort_to_interleave_indices( batch_idx, row_idx ))); } - keys.push((col.value(row_idx), batch_idx, row_idx as u32)); + let part_id = col.value(row_idx) as usize; + if part_id >= num_partitions { + return Err(Error::invalid_input(format!( + "partition ID {} is out of range [0, {})", + part_id, num_partitions + ))); + } + partition_counts[part_id] += 1; } } - keys.sort_unstable_by_key(|k| k.0); - let mut partition_counts = vec![0u64; num_partitions]; - 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 - ))); + // Prefix sums turn the counts into the first slot of each partition's run, + // then the scatter pass below consumes them as a write cursor per partition. + let mut cursors = Vec::with_capacity(num_partitions); + let mut running = 0usize; + for &count in partition_counts.iter() { + cursors.push(running); + running += count as usize; + } + + let mut interleave_indices = vec![(0usize, 0usize); total_rows]; + for (batch_idx, col) in part_id_columns.iter().enumerate() { + for (row_idx, &part_id) in col.values().iter().enumerate() { + let cursor = &mut cursors[part_id as usize]; + interleave_indices[*cursor] = (batch_idx, row_idx); + *cursor += 1; } - partition_counts[pid] += 1; - interleave_indices.push((*batch_idx as usize, *row_idx as usize)); } + Ok((interleave_indices, partition_counts)) } @@ -2851,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() {