diff --git a/rust/lance-table/benches/row_id_index.rs b/rust/lance-table/benches/row_id_index.rs index a951622005c..dd20fccb639 100644 --- a/rust/lance-table/benches/row_id_index.rs +++ b/rust/lance-table/benches/row_id_index.rs @@ -27,7 +27,7 @@ use lance_core::utils::deletion::DeletionVector; use lance_io::ReadBatchParams; use lance_select::{RowAddrMask, RowAddrTreeMap}; use lance_table::format::pb; -use lance_table::rowids::FragmentRowIdIndex; +use lance_table::rowids::{FragmentRowIdIndex, segment::U64Segment}; use lance_table::{ rowids::{RowIdIndex, RowIdSequence, read_row_ids, write_row_ids}, utils::stream::{RowIdAndDeletesConfig, apply_row_id_and_deletes}, @@ -70,6 +70,36 @@ fn make_frag_sequences( .collect() } +fn make_clustered_range_frag_sequences(num_rows: u64) -> Vec { + let rows_per_fragment = num_rows / 100; + (0_u32..100) + .map(|fragment_id| { + let start = fragment_id as u64 * rows_per_fragment; + let end = start + rows_per_fragment; + let mut ranges = Vec::new(); + let mut row_id = start; + while row_id < end { + let present_end = (row_id + 250).min(end); + ranges.push(row_id..present_end); + row_id = present_end.saturating_add(250); + } + let wire_sequence = pb::RowIdSequence { + segments: ranges + .into_iter() + .map(|range| pb::U64Segment::from(U64Segment::Range(range))) + .collect(), + }; + // Decode the wire ranges into the compact form used after PR #9311. + let sequence = read_row_ids(&wire_sequence.encode_to_vec()).unwrap(); + FragmentRowIdIndex { + fragment_id, + row_id_sequence: Arc::new(sequence), + deletion_vector: Arc::new(DeletionVector::default()), + } + }) + .collect() +} + // For range of values // https://bheisler.github.io/criterion.rs/book/user_guide/benchmarking_with_inputs.html @@ -199,6 +229,11 @@ fn bench_creation(c: &mut Criterion) { ); } + let clustered_indices = make_clustered_range_frag_sequences(num_rows()); + group.bench_function("BuildIndexClusteredRanges", |b| { + b.iter(|| std::hint::black_box(RowIdIndex::new(&clustered_indices).unwrap())); + }); + group.finish(); } diff --git a/rust/lance-table/src/rowids/index.rs b/rust/lance-table/src/rowids/index.rs index c4d7b8f3771..b04ff744eab 100644 --- a/rust/lance-table/src/rowids/index.rs +++ b/rust/lance-table/src/rowids/index.rs @@ -557,6 +557,15 @@ fn decompose_segment_no_deletions(segment: &U64Segment, start_address: u64) -> O let coverage = range.start..=range.end - 1; Some((coverage, (row_id_segment, address_segment))) } + U64Segment::Ranges { .. } if !segment.is_empty() => { + // Without deletions, a live ID's position is also its physical + // offset within this segment. + let row_id_segment = segment.clone(); + let address_segment = + U64Segment::Range(start_address..start_address + row_id_segment.len() as u64); + let coverage = row_id_segment.range()?; + Some((coverage, (row_id_segment, address_segment))) + } _ if segment.is_empty() => None, _ => { // Non-Range segments: must iterate to build address mapping. @@ -801,6 +810,100 @@ mod tests { ); } + #[test] + fn test_merged_index_keeps_ranges_segments_compact() { + let row_id_ranges = [100..103, 108..110, 115..119]; + let sequence = RowIdSequence(vec![ + U64Segment::Range(50..53), + U64Segment::from_sorted_ranges(&row_id_ranges).unwrap(), + ]); + let source = fragment(7, sequence); + let index = RowIdIndex::new(std::slice::from_ref(&source)).unwrap(); + + let merged = index.merged.as_ref().expect("small input should be merged"); + assert!(matches!( + merged.get(&100).unwrap().0, + U64Segment::Ranges { .. } + )); + + for (row_id, row_offset) in [ + (50, 0), + (51, 1), + (52, 2), + (100, 3), + (101, 4), + (102, 5), + (108, 6), + (109, 7), + (115, 8), + (116, 9), + (117, 10), + (118, 11), + ] { + assert_eq!( + index.get(row_id).unwrap(), + Some(RowAddress::new_from_parts(7, row_offset)), + "row id {row_id}" + ); + } + for missing in [49, 53, 99, 103, 107, 110, 114, 119, 120] { + assert_eq!(index.get(missing).unwrap(), None, "row id {missing}"); + } + assert_eq!( + index.get_many(&[108, 53, 100, 108]).unwrap(), + vec![ + Some(RowAddress::new_from_parts(7, 6)), + None, + Some(RowAddress::new_from_parts(7, 3)), + Some(RowAddress::new_from_parts(7, 6)), + ] + ); + } + + #[test] + fn test_overlapping_ranges_chunks_keep_row_addresses() { + let left_ranges = [100..103, 108..110, 115..119]; + let right_ranges = [103..108, 110..115]; + let sources = [ + fragment( + 7, + RowIdSequence(vec![U64Segment::from_sorted_ranges(&left_ranges).unwrap()]), + ), + fragment( + 8, + RowIdSequence(vec![U64Segment::from_sorted_ranges(&right_ranges).unwrap()]), + ), + ]; + let index = RowIdIndex::new(&sources).unwrap(); + + for row_id in 100..119 { + let (fragment_id, row_offset) = if (100..103).contains(&row_id) + || (108..110).contains(&row_id) + || (115..119).contains(&row_id) + { + let row_offset = match row_id { + 100..=102 => row_id - 100, + 108..=109 => row_id - 105, + 115..=118 => row_id - 110, + _ => unreachable!(), + }; + (7, row_offset) + } else { + let row_offset = match row_id { + 103..=107 => row_id - 103, + 110..=114 => row_id - 105, + _ => unreachable!(), + }; + (8, row_offset) + }; + assert_eq!( + index.get(row_id).unwrap(), + Some(RowAddress::new_from_parts(fragment_id, row_offset as u32)), + "row id {row_id}" + ); + } + } + #[test] fn test_deep_overlap_merges_however_many_rows_it_reads() { // Just past the row budget in total, interleaved so every fragment