fix: compute row offsets for unsorted addresses

rust/lance/src/io/exec/rowids.rs

@@ -438,15 +438,19 @@ impl AddRowOffsetExec {
         frag_id_to_offset: &HashMap<u32, FragInfo>,
     ) -> Result<ArrayRef> {
         let row_addr_values = row_addr.as_primitive::<UInt64Type>().values();
-        let mut row_offsets = Vec::with_capacity(row_addr_values.len());
+        let mut row_offsets = vec![0; row_addr_values.len()];
+        // The deletion iterator only moves forward, so compute in address order and
+        // scatter the offsets back into input order.
+        let mut sorted_row_indices = (0..row_addr_values.len()).collect::<Vec<_>>();
+        sorted_row_indices.sort_unstable_by_key(|index| row_addr_values[*index]);
 
         let mut last_frag_id = u32::MAX;
         let mut last_frag_offset = 0;
         let mut last_frag_delete_count = 0;
         let mut dv_iter = None;
 
-        for addr in row_addr_values {
-            let addr = RowAddress::new_from_u64(*addr);
+        for row_index in sorted_row_indices {
+            let addr = RowAddress::new_from_u64(row_addr_values[row_index]);
             let frag_id = addr.fragment_id();
             if frag_id != last_frag_id {
                 last_frag_id = frag_id;
@@ -476,10 +480,10 @@ impl AddRowOffsetExec {
                         break;
                     }
                 }
-                row_offsets
-                    .push(last_frag_offset + row_offset as u64 - last_frag_delete_count as u64);
+                row_offsets[row_index] =
+                    last_frag_offset + row_offset as u64 - last_frag_delete_count as u64;
             } else {
-                row_offsets.push(last_frag_offset + row_offset as u64);
+                row_offsets[row_index] = last_frag_offset + row_offset as u64;
             }
         }
 
@@ -727,4 +731,27 @@ mod test {
             .fold(0, |acc, col| acc + col.get_array_memory_size());
         assert_eq!(stats.total_byte_size, Precision::Exact(actual_byte_size));
     }
+
+    #[test]
+    fn test_row_offsets_with_unsorted_addresses() {
+        let row_addrs: ArrayRef = Arc::new(UInt64Array::from(vec![
+            u64::from(RowAddress::new_from_parts(0, 100)),
+            u64::from(RowAddress::new_from_parts(0, 50)),
+        ]));
+        let frag_id_to_offset = HashMap::from([(
+            0,
+            FragInfo {
+                row_offset: 1_000,
+                deletion_vector: Some(Arc::new(DeletionVector::from_iter([10, 60]))),
+            },
+        )]);
+
+        let row_offsets =
+            AddRowOffsetExec::compute_row_offsets(&row_addrs, &frag_id_to_offset).unwrap();
+
+        assert_eq!(
+            row_offsets.as_primitive::<UInt64Type>().values(),
+            &[1_098, 1_049]
+        );
+    }
 }