perf(index): bucket shuffle rows by partition id instead of sorting keys

rust/lance-index/src/vector/v3/shuffler.rs

@@ -596,34 +596,47 @@ fn sort_to_interleave_indices(
     num_partitions: usize,
 ) -> Result<InterleaveResult> {
     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!(
                     "null partition ID at batch {} row {}",
                     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() {