diff --git a/rust/lance/src/io/exec/rowids.rs b/rust/lance/src/io/exec/rowids.rs index dc9a54cc182..a0add27d3f7 100644 --- a/rust/lance/src/io/exec/rowids.rs +++ b/rust/lance/src/io/exec/rowids.rs @@ -438,15 +438,19 @@ impl AddRowOffsetExec { frag_id_to_offset: &HashMap, ) -> Result { let row_addr_values = row_addr.as_primitive::().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::>(); + 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::().values(), + &[1_098, 1_049] + ); + } }