diff --git a/rust/lance-table/src/rowids.rs b/rust/lance-table/src/rowids.rs index 4cadaf67b76..b6cf4c6dc79 100644 --- a/rust/lance-table/src/rowids.rs +++ b/rust/lance-table/src/rowids.rs @@ -18,6 +18,7 @@ use std::ops::{Range, RangeInclusive}; mod bitmap; mod encoded_array; mod index; +mod runs; pub mod segment; mod serde; pub mod version; @@ -317,6 +318,25 @@ impl RowIdSequence { self.0.extend(other.0); } + /// Re-encode every segment that is smaller as a run of `Range` segments + /// than as it is stored now (see [`U64Segment::as_ranges`]). Returns + /// whether anything changed. + /// + /// The result is written as plain `Range` segments, which every reader + /// understands, but a reader without the compact in-memory form handles + /// thousands of segments per fragment slowly, so this is only called for + /// tables that opted in. + pub fn use_range_segments(&mut self) -> bool { + let mut changed = false; + for segment in &mut self.0 { + if let Some(ranges) = segment.as_ranges() { + *segment = ranges; + changed = true; + } + } + changed + } + /// Remove a set of row ids from the sequence. pub fn delete(&mut self, row_ids: impl IntoIterator) { // Order the row ids by position in which they appear in the sequence. @@ -668,6 +688,24 @@ impl RowIdSequence { offset_start + position_in_range - holes_passed }))); } + U64Segment::Ranges { range, runs } => { + let offset_start = offset; + offset += runs.present_len() as u64; + let mut ids = RowAddrTreeMap::new(); + for present in runs.present_ranges() { + ids.insert_range( + (range.start + present.start as u64) + ..(range.start + present.end as u64), + ); + } + ids.mask(mask); + ranges.extend(GroupingIterator::new(ids.into_addr_iter().map(|addr| { + let position = runs + .position((addr - range.start) as u32) + .expect("addresses were inserted from the present ranges"); + offset_start + position as u64 + }))); + } U64Segment::SortedArray(array) | U64Segment::Array(array) => { // TODO: Could probably optimize the sorted array case to be O(N) instead of O(N log N) ranges.extend(GroupingIterator::new(array.iter().enumerate().filter_map( @@ -749,6 +787,14 @@ impl From<&RowIdSequence> for RowAddrTreeMap { seg.remove(hole); } } + U64Segment::Ranges { range, runs } => { + for present in runs.present_ranges() { + seg.insert_range( + (range.start + present.start as u64) + ..(range.start + present.end as u64), + ); + } + } U64Segment::SortedArray(array) | U64Segment::Array(array) => { for val in array.iter() { seg.insert(val); @@ -1829,4 +1875,63 @@ mod test { assert_eq!(*r.start(), 0); assert_eq!(*r.end(), 104); } + + #[test] + fn test_range_segments_sequence_matches_bitmap_sequence() { + let live: Vec = (0..600u64) + .filter(|v| !(50..250).contains(v) && !(300..310).contains(v)) + .collect(); + let bitmap_sequence = RowIdSequence::from(live.as_slice()); + assert!(matches!( + bitmap_sequence.0.as_slice(), + [U64Segment::RangeWithBitmap { .. }] + )); + let mut runs_sequence = bitmap_sequence.clone(); + assert!(runs_sequence.use_range_segments()); + assert!(matches!( + runs_sequence.0.as_slice(), + [U64Segment::Ranges { .. }] + )); + // Idempotent: a second pass has nothing left to convert. + let again = runs_sequence.clone(); + assert!(!runs_sequence.use_range_segments()); + assert_eq!(runs_sequence, again); + + assert_eq!(runs_sequence.iter().collect::>(), live); + assert_eq!(runs_sequence.len(), bitmap_sequence.len()); + let mut cursor = runs_sequence.cursor(); + let mut chunked = Vec::new(); + for start in (0..live.len()).step_by(37) { + let end = (start + 37).min(live.len()); + chunked.extend(runs_sequence.select_range_with_cursor(&mut cursor, start..end)); + } + assert_eq!(chunked, live); + let picks = [0usize, 5, 49, 50, 300]; + assert_eq!( + runs_sequence + .select(picks.iter().copied()) + .collect::>(), + bitmap_sequence + .select(picks.iter().copied()) + .collect::>() + ); + + for mask in [ + RowAddrMask::from_allowed(RowAddrTreeMap::from_iter(&[0, 49, 50, 100, 250, 251, 599])), + RowAddrMask::from_block(RowAddrTreeMap::from_iter(&[0, 250, 305, 599])), + ] { + assert_eq!( + runs_sequence.mask_to_offset_ranges(&mask), + bitmap_sequence.mask_to_offset_ranges(&mask) + ); + } + assert_eq!( + RowAddrTreeMap::from(&runs_sequence), + RowAddrTreeMap::from(&bitmap_sequence) + ); + assert_eq!( + read_row_ids(write_row_ids(&runs_sequence).as_slice()).unwrap(), + runs_sequence + ); + } } diff --git a/rust/lance-table/src/rowids/index.rs b/rust/lance-table/src/rowids/index.rs index eb5dcf05730..c4d7b8f3771 100644 --- a/rust/lance-table/src/rowids/index.rs +++ b/rust/lance-table/src/rowids/index.rs @@ -1609,4 +1609,33 @@ mod tests { prop_assert!(error_message.contains(&expected_message)); } } + + #[test] + fn test_index_resolves_range_segments() { + let live: Vec = (1000..3000u64) + .filter(|v| !(1200..1900).contains(v)) + .collect(); + let mut runs_sequence = RowIdSequence::from(live.as_slice()); + assert!(runs_sequence.use_range_segments()); + let index = RowIdIndex::new(&[fragment(7, runs_sequence)]).unwrap(); + for (position, &row_id) in live.iter().enumerate() { + assert_eq!( + index.get(row_id).unwrap(), + Some(RowAddress::new_from_parts(7, position as u32)), + "row id {row_id}" + ); + } + assert_eq!(index.get(1500).unwrap(), None); + let addr = |position: u32| Some(RowAddress::new_from_parts(7, position)); + assert_eq!( + index.get_many(&[2999, 1000, 1199, 1900, 1500]).unwrap(), + vec![ + addr(live.len() as u32 - 1), + addr(0), + addr(199), + addr(200), + None + ] + ); + } } diff --git a/rust/lance-table/src/rowids/runs.rs b/rust/lance-table/src/rowids/runs.rs new file mode 100644 index 00000000000..7a8a99814f7 --- /dev/null +++ b/rust/lance-table/src/rowids/runs.rs @@ -0,0 +1,325 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Run-length encoded holes of a [`U64Segment::Ranges`](super::U64Segment). + +use std::ops::Range; + +use lance_core::Result; +use lance_core::deepsize::DeepSizeOf; + +use super::bitmap::Bitmap; +use super::serde::corrupt_row_id_metadata; + +/// The missing offsets of a range, as maximal runs. +/// +/// Offsets are relative to the range start and fit `u32`, which bounds the span +/// of a range this encoding can describe to `u32::MAX`; a fragment never holds +/// more rows than that, and writers fall back to another encoding for wider +/// spans. Compared to a bitmap this costs 8 bytes per run rather than one bit +/// per offset, so it pays off exactly when deletions cluster, as they do after +/// compaction of fragments with deleted rows. +#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)] +pub struct HoleRuns { + /// Length of the range the runs live in (offsets are below this). + span: u32, + /// First missing offset of each run, strictly increasing. + starts: Vec, + /// `present_before[k]` is the number of present offsets before run `k`; + /// the final entry (index `num_runs`) is the number present in total. Run + /// ends follow from it: the offsets missing before run `k` are + /// `starts[k] - present_before[k]`, and a run is as long as the missing + /// count grows across it. + present_before: Vec, +} + +impl HoleRuns { + /// Build from run bounds. Rejects runs that are empty, out of order, + /// overlapping, adjacent (a maximal run never touches the next) or beyond + /// `span`. + pub fn try_new(span: u32, starts: Vec, ends: Vec) -> Result { + if starts.len() != ends.len() { + return Err(corrupt_row_id_metadata(format!( + "Ranges has {} run starts but {} run ends", + starts.len(), + ends.len() + ))); + } + let mut present_before = Vec::with_capacity(starts.len() + 1); + let mut missing: u32 = 0; + let mut previous_end: Option = None; + for (i, (&start, &end)) in starts.iter().zip(&ends).enumerate() { + if start >= end { + return Err(corrupt_row_id_metadata(format!( + "Ranges run {i} is empty or reversed: {start}..{end}" + ))); + } + if end > span { + return Err(corrupt_row_id_metadata(format!( + "Ranges run {i} ({start}..{end}) ends beyond the span {span}" + ))); + } + if let Some(previous_end) = previous_end + && previous_end >= start + { + return Err(corrupt_row_id_metadata(format!( + "Ranges run {i} starts at {start}, but the previous run ends at \ + {previous_end}: runs must be sorted, disjoint and non-adjacent" + ))); + } + present_before.push(start - missing); + // Disjoint runs below `span` cannot miss more than `span` offsets. + missing += end - start; + previous_end = Some(end); + } + present_before.push(span - missing); + Ok(Self { + span, + starts, + present_before, + }) + } + + /// The runs of cleared bits of `bitmap`, whose length must fit `u32`. + /// + /// Walks bytes rather than bits: on a compacted table almost every byte is + /// all-present or all-missing, and the bit loop only runs at run edges. + pub fn from_bitmap(bitmap: &Bitmap) -> Self { + let span = bitmap.len(); + let mut starts = Vec::new(); + let mut ends = Vec::new(); + let mut open_run: Option = None; + let mut offset: usize = 0; + for &byte in bitmap.bytes() { + let valid_bits = (span - offset).min(8); + // Bits past `span` read as present so they never open a run. + let mut present = byte | (u8::MAX.checked_shl(valid_bits as u32).unwrap_or(0)); + if present == u8::MAX { + if let Some(start) = open_run.take() { + starts.push(start); + ends.push(offset as u32); + } + } else if present == 0 { + open_run.get_or_insert(offset as u32); + } else { + for bit in 0..valid_bits { + let is_present = present & 1 != 0; + present >>= 1; + let here = (offset + bit) as u32; + match (is_present, open_run) { + (false, None) => open_run = Some(here), + (true, Some(start)) => { + starts.push(start); + ends.push(here); + open_run = None; + } + _ => {} + } + } + } + offset += 8; + if offset >= span { + break; + } + } + if let Some(start) = open_run { + starts.push(start); + ends.push(span as u32); + } + Self::try_new(span as u32, starts, ends).expect("runs derived from a bitmap are valid") + } + + pub fn num_runs(&self) -> usize { + self.starts.len() + } + + /// Offsets present over the whole span. + pub fn present_len(&self) -> u32 { + self.present_before[self.starts.len()] + } + + /// Start of run `k`, or the span end once the runs are exhausted. + fn start_or_span(&self, k: usize) -> u32 { + self.starts.get(k).copied().unwrap_or(self.span) + } + + /// Offsets missing in runs `0..k`. + fn missing_before(&self, k: usize) -> u32 { + self.start_or_span(k) - self.present_before[k] + } + + /// First offset after run `k`. + fn end(&self, k: usize) -> u32 { + self.starts[k] + (self.missing_before(k + 1) - self.missing_before(k)) + } + + /// Position of the present `offset`, or `None` when it is missing or out of range. + pub fn position(&self, offset: u32) -> Option { + if offset >= self.span { + return None; + } + let k = self.starts.partition_point(|&start| start <= offset); + if k > 0 && offset < self.end(k - 1) { + return None; + } + Some(offset - self.missing_before(k)) + } + + /// Runs entirely before the present value at `position`. + fn runs_before_position(&self, position: u32) -> usize { + self.present_before[..self.starts.len()].partition_point(|&present| present <= position) + } + + /// Offset of the present value at `position`. + pub fn offset_at(&self, position: u32) -> Option { + if position >= self.present_len() { + return None; + } + let k = self.runs_before_position(position); + Some(position + self.missing_before(k)) + } + + /// The present offsets, as the maximal ranges between runs. + pub fn present_ranges(&self) -> impl DoubleEndedIterator> + '_ { + (0..=self.starts.len()) + .map(move |k| { + let start = if k == 0 { 0 } else { self.end(k - 1) }; + start..self.start_or_span(k) + }) + .filter(|range| !range.is_empty()) + } + + /// Append `base + offset` for the present offsets at `positions`. + /// + /// Streaming readers materialize `_rowid` a batch at a time through this; + /// walking the present ranges keeps that linear instead of a binary search + /// per row. + pub fn extend_values(&self, base: u64, positions: Range, values: &mut Vec) { + let end = positions.end.min(self.present_len()); + let mut position = positions.start; + let mut k = self.runs_before_position(position); + while position < end && k <= self.starts.len() { + // Present range between run k - 1 and run k, in offsets and positions. + let gap_start = if k == 0 { 0 } else { self.end(k - 1) }; + let gap_end = self.start_or_span(k); + let first_position = gap_start - self.missing_before(k); + let gap_positions = first_position..first_position + (gap_end - gap_start); + let take = position.max(gap_positions.start)..end.min(gap_positions.end); + if !take.is_empty() { + let offset = gap_start + (take.start - gap_positions.start); + values.extend((base + offset as u64)..(base + offset as u64 + take.len() as u64)); + position = take.end; + } + k += 1; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn runs() -> HoleRuns { + // span 20: present 0..3, missing 3..7, present 7..8, missing 8..15, present 15..20 + HoleRuns::try_new(20, vec![3, 8], vec![7, 15]).unwrap() + } + + #[test] + fn test_counts_and_ranges() { + let runs = runs(); + assert_eq!(runs.span - runs.present_len(), 11); + assert_eq!(runs.present_len(), 9); + assert_eq!((0..2).map(|k| runs.end(k)).collect::>(), vec![7, 15]); + assert_eq!( + runs.present_ranges().collect::>(), + vec![0..3, 7..8, 15..20] + ); + assert_eq!( + runs.present_ranges().rev().collect::>(), + vec![15..20, 7..8, 0..3] + ); + } + + #[test] + fn test_position_and_offset_round_trip() { + let runs = runs(); + let present: Vec = runs.present_ranges().flatten().collect(); + assert_eq!(present, vec![0, 1, 2, 7, 15, 16, 17, 18, 19]); + for (position, &offset) in present.iter().enumerate() { + assert_eq!( + runs.position(offset), + Some(position as u32), + "offset {offset}" + ); + assert_eq!( + runs.offset_at(position as u32), + Some(offset), + "position {position}" + ); + } + for missing in [3, 4, 6, 8, 14, 20, 100] { + assert_eq!(runs.position(missing), None, "offset {missing}"); + } + assert_eq!(runs.offset_at(9), None); + } + + #[test] + fn test_extend_values_crosses_runs() { + let runs = runs(); + let mut values = Vec::new(); + runs.extend_values(100, 1..7, &mut values); + assert_eq!(values, vec![101, 102, 107, 115, 116, 117]); + values.clear(); + runs.extend_values(100, 4..40, &mut values); + assert_eq!(values, vec![115, 116, 117, 118, 119]); + values.clear(); + runs.extend_values(100, 9..12, &mut values); + assert!(values.is_empty()); + } + + #[test] + fn test_from_bitmap_matches_cleared_bits() { + for (len, cleared) in [ + (20usize, vec![3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14]), + (1, vec![0]), + (9, vec![]), + (17, vec![0, 16]), + (64, (8..56).collect()), + (70, vec![7, 8, 9, 15, 16, 63, 64, 65, 66, 67, 68, 69]), + ] { + let mut bitmap = Bitmap::new_full(len); + for &offset in &cleared { + bitmap.clear(offset); + } + let runs = HoleRuns::from_bitmap(&bitmap); + let missing: Vec = (0..len as u32) + .filter(|offset| runs.position(*offset).is_none()) + .collect(); + assert_eq!( + missing, + cleared.iter().map(|&c| c as u32).collect::>(), + "len {len}" + ); + assert_eq!((runs.span - runs.present_len()) as usize, cleared.len()); + // Runs are maximal: each ends strictly before the next starts. + for k in 1..runs.starts.len() { + assert!(runs.end(k - 1) < runs.starts[k]); + } + } + } + + #[test] + fn test_try_new_rejects_malformed_runs() { + let cases: [(&str, Vec, Vec); 5] = [ + ("length mismatch", vec![1], vec![]), + ("empty run", vec![3], vec![3]), + ("beyond span", vec![3], vec![21]), + ("overlapping", vec![3, 5], vec![7, 9]), + ("adjacent", vec![3, 7], vec![7, 9]), + ]; + for (name, starts, ends) in cases { + let error = HoleRuns::try_new(20, starts, ends).unwrap_err(); + assert!(error.to_string().contains("Ranges"), "{name}: {error}"); + } + } +} diff --git a/rust/lance-table/src/rowids/segment.rs b/rust/lance-table/src/rowids/segment.rs index b634812f164..604a708328d 100644 --- a/rust/lance-table/src/rowids/segment.rs +++ b/rust/lance-table/src/rowids/segment.rs @@ -3,7 +3,7 @@ use std::ops::{Range, RangeInclusive}; -use super::{bitmap::Bitmap, encoded_array::EncodedU64Array}; +use super::{bitmap::Bitmap, encoded_array::EncodedU64Array, runs::HoleRuns}; use lance_core::deepsize::DeepSizeOf; /// Convert an estimated serialized byte cost from `u128` to `usize`, saturating @@ -61,6 +61,18 @@ pub enum U64Segment { /// Total size: 24 bytes + ceil((max - min) / 8) bytes /// Use when: max - min > 16 * len RangeWithBitmap { range: Range, bitmap: Bitmap }, + /// Two or more sorted, disjoint ranges of row ids, held compactly. + /// + /// This is how a run of consecutive `Range` segments is kept in memory: + /// on the wire it is exactly those `Range` segments, so it needs no format + /// support (see `read_row_ids` in `serde`). In memory it is two + /// `u32` per range plus prefix sums, and lookups are a binary search, + /// where separate segments would cost a heap enum each and linear scans. + /// + /// Total size on the wire: ~16 bytes per range. + /// Use when: fewer than one range per 128 slots, and only for tables that + /// opted in (see [`Self::as_ranges`]); `from_slice` never picks it. + Ranges { range: Range, runs: HoleRuns }, /// A sorted array of row ids, that is sparse. /// /// Total size: 24 bytes + 2 * n_values bytes @@ -75,6 +87,7 @@ impl DeepSizeOf for U64Segment { Self::Range(_) => 0, Self::RangeWithHoles { holes, .. } => holes.deep_size_of_children(context), Self::RangeWithBitmap { bitmap, .. } => bitmap.deep_size_of_children(context), + Self::Ranges { runs, .. } => runs.deep_size_of_children(context), Self::SortedArray(array) => array.deep_size_of_children(context), Self::Array(array) => array.deep_size_of_children(context), } @@ -228,6 +241,53 @@ impl U64Segment { pub fn from_slice(slice: &[u64]) -> Self { Self::from_iter(slice.iter().copied()) } + + /// Two or more sorted, disjoint, non-adjacent ranges as one [`Self::Ranges`] + /// segment; `None` for fewer ranges or a span that does not fit `u32`. + pub fn from_sorted_ranges(ranges: &[Range]) -> Option { + let [first, .., last] = ranges else { + return None; + }; + let base = first.start; + let span = u32::try_from(last.end - base).ok()?; + let hole_starts = ranges[..ranges.len() - 1] + .iter() + .map(|range| (range.end - base) as u32) + .collect(); + let hole_ends = ranges[1..] + .iter() + .map(|range| (range.start - base) as u32) + .collect(); + let runs = HoleRuns::try_new(span, hole_starts, hole_ends).ok()?; + Some(Self::Ranges { + range: base..last.end, + runs, + }) + } + + /// The same values as a run of ranges, when writing them as `Range` + /// segments is smaller than the current encoding; `None` when it is not + /// (or the span does not fit the compact form's 32-bit offsets). + /// + /// A `Range` segment costs about 16 bytes on the wire, so ranges beat a + /// bitmap when there is less than one range per 128 slots. Hole arrays + /// are small by construction and are left alone. + pub fn as_ranges(&self) -> Option { + const RANGE_BYTES: usize = 16; + match self { + Self::RangeWithBitmap { range, bitmap } => { + if bitmap.len() > u32::MAX as usize { + return None; + } + let runs = HoleRuns::from_bitmap(bitmap); + ((runs.num_runs() + 1) * RANGE_BYTES < bitmap.bytes().len()).then(|| Self::Ranges { + range: range.clone(), + runs, + }) + } + _ => None, + } + } } impl FromIterator for U64Segment { @@ -256,6 +316,12 @@ impl U64Segment { bitmap.get(offset) })) } + Self::Ranges { range, runs } => { + let start = range.start; + Box::new(runs.present_ranges().flat_map(move |offsets| { + (start + offsets.start as u64)..(start + offsets.end as u64) + })) + } Self::SortedArray(array) => Box::new(array.iter()), Self::Array(array) => Box::new(array.iter()), } @@ -271,6 +337,7 @@ impl U64Segment { let holes = bitmap.count_zeros(); (range.end - range.start) as usize - holes } + Self::Ranges { runs, .. } => runs.present_len() as usize, Self::SortedArray(array) => array.len(), Self::Array(array) => array.len(), } @@ -298,6 +365,7 @@ impl U64Segment { match self { Self::Range(range) | Self::RangeWithBitmap { range, .. } + | Self::Ranges { range, .. } | Self::RangeWithHoles { range, .. } => { (!range.is_empty()).then(|| range.start..=(range.end - 1)) } @@ -341,6 +409,13 @@ impl U64Segment { } } } + Self::Ranges { range, runs } => { + if !range.contains(&val) { + return None; + } + runs.position((val - range.start) as u32) + .map(|position| position as usize) + } Self::RangeWithBitmap { range, bitmap } => { if range.contains(&val) && bitmap.get((val - range.start) as usize) { let offset = (val - range.start) as usize; @@ -385,6 +460,10 @@ impl U64Segment { Some(range.start + i as u64 + lo as u64) } Self::RangeWithBitmap { .. } => self.cursor().get(i), + Self::Ranges { range, runs } => u32::try_from(i) + .ok() + .and_then(|position| runs.offset_at(position)) + .map(|offset| range.start + offset as u64), Self::SortedArray(array) => array.get(i), Self::Array(array) => array.get(i), } @@ -417,6 +496,9 @@ impl U64Segment { let idx = (val - range.start) as usize; bitmap.get(idx) } + Self::Ranges { range, runs } => { + range.contains(&val) && runs.position((val - range.start) as u32).is_some() + } Self::SortedArray(array) => array.binary_search(val).is_ok(), Self::Array(array) => array.iter().any(|v| v == val), } @@ -496,6 +578,14 @@ impl U64Segment { bitmap: Bitmap::from(new_bitmap.as_slice()), } } + Self::Ranges { range, runs } => { + // Rare on the append path; rebuild from the values and let + // `from_slice` pick an encoding (the commit re-encodes runs). + let current = Self::Ranges { range, runs }; + let mut values: Vec = current.iter().collect(); + values.push(val); + Self::from_slice(&values) + } Self::SortedArray(array) => match array { EncodedU64Array::U64(mut vec) => { vec.push(val); @@ -614,6 +704,7 @@ impl U64Segment { Self::Range(_) => true, Self::RangeWithHoles { .. } => true, Self::RangeWithBitmap { .. } => true, + Self::Ranges { .. } => true, Self::SortedArray(_) => true, Self::Array(_) => false, }; @@ -696,6 +787,15 @@ impl SegmentCursorState { ); } } + U64Segment::Ranges { range, runs } => { + // Positions are below the span, which fits u32. + let clamp = |position: usize| position.min(u32::MAX as usize) as u32; + runs.extend_values( + range.start, + clamp(selection.start)..clamp(selection.end), + values, + ); + } _ => values.extend(selection.filter_map(|index| segment.get(index))), } return; @@ -1428,4 +1528,97 @@ mod test { "Empty segment should not contain anything" ); } + + /// Row ids 1000..2000 minus three clusters of holes: dense enough that + /// `from_slice` picks a bitmap, clustered enough that runs are smaller. + fn clustered_values() -> Vec { + (1000..2000u64) + .filter(|v| !(1100..1300).contains(v) && !(1500..1650).contains(v) && *v != 1999) + .collect() + } + + #[test] + fn test_ranges_segment_matches_bitmap_semantics() { + let values = clustered_values(); + let bitmap = U64Segment::from_slice(&values); + assert!(matches!(bitmap, U64Segment::RangeWithBitmap { .. })); + let runs = bitmap + .as_ranges() + .expect("three runs are smaller than a 125 byte bitmap"); + let U64Segment::Ranges { runs: holes, .. } = &runs else { + panic!("expected a run-length segment"); + }; + assert_eq!(holes.num_runs(), 2); + + assert_eq!(runs.len(), bitmap.len()); + assert_eq!(runs.range(), bitmap.range()); + assert_eq!(runs.iter().collect::>(), values); + assert_eq!( + runs.iter().rev().collect::>(), + values.iter().rev().copied().collect::>() + ); + for (position, &value) in values.iter().enumerate() { + assert_eq!(runs.position(value), Some(position), "value {value}"); + assert_eq!(runs.get(position), Some(value), "position {position}"); + assert!(runs.contains(value)); + } + for missing in [999, 1100, 1250, 1299, 1500, 1649, 1999, 2000, 5000] { + assert_eq!(runs.position(missing), None, "missing {missing}"); + assert!(!runs.contains(missing)); + } + assert_eq!(runs.get(values.len()), None); + assert_eq!( + runs.slice(95, 20).iter().collect::>(), + values[95..115].to_vec() + ); + + let deleted = [1000, 1050, 1400]; + assert_eq!( + runs.delete(&deleted).iter().collect::>(), + bitmap.delete(&deleted).iter().collect::>() + ); + let (mut masked_runs, mut masked_bitmap) = (runs.clone(), bitmap); + masked_runs.mask(&[0, 1, 7]); + masked_bitmap.mask(&[0, 1, 7]); + assert_eq!( + masked_runs.iter().collect::>(), + masked_bitmap.iter().collect::>() + ); + + // The cursor path streaming readers use. + let mut cursor = SegmentCursorState::default(); + let mut got = Vec::new(); + cursor.extend_range(&runs, 90..110, &mut got); + assert_eq!(got, values[90..110].to_vec()); + got.clear(); + cursor.extend_range(&runs, 300..values.len() + 10, &mut got); + assert_eq!(got, values[300..].to_vec()); + assert_eq!(runs.cursor().get(400), Some(values[400])); + + let higher = runs.clone().with_new_high(2500).unwrap(); + assert_eq!(higher.iter().last(), Some(2500)); + assert_eq!(higher.len(), values.len() + 1); + } + + #[test] + fn test_as_ranges_only_when_smaller() { + // Alternating holes: one run per two slots, far more than one per 64. + let alternating: Vec = (0..1024u64).filter(|v| v % 2 == 0).collect(); + let bitmap = U64Segment::from_slice(&alternating); + assert!(matches!(bitmap, U64Segment::RangeWithBitmap { .. })); + assert_eq!(bitmap.as_ranges(), None); + + assert_eq!(U64Segment::Range(0..10).as_ranges(), None); + assert_eq!( + U64Segment::SortedArray(vec![1, 1000, 1_000_000].into()).as_ranges(), + None + ); + + // Hole arrays are already small; only bitmaps convert. + let holes = U64Segment::RangeWithHoles { + range: 0..1000, + holes: vec![3, 77, 500].into(), + }; + assert_eq!(holes.as_ranges(), None); + } } diff --git a/rust/lance-table/src/rowids/serde.rs b/rust/lance-table/src/rowids/serde.rs index 02bc6c673cf..94e437cc722 100644 --- a/rust/lance-table/src/rowids/serde.rs +++ b/rust/lance-table/src/rowids/serde.rs @@ -5,11 +5,14 @@ use crate::{format::pb, rowids::bitmap::Bitmap}; use lance_core::{Error, Result}; use super::{RowIdSequence, U64Segment, encoded_array::EncodedU64Array}; +use std::ops::Range; + use prost::Message; +use prost::encoding::{WireType, decode_key, decode_varint}; const ROW_ID_METADATA: &str = "row ID metadata"; -fn corrupt_row_id_metadata(message: impl Into) -> Error { +pub(super) fn corrupt_row_id_metadata(message: impl Into) -> Error { Error::corrupt_file_named(ROW_ID_METADATA, message) } @@ -54,29 +57,41 @@ fn first_non_increasing_pair(array: &EncodedU64Array) -> Option<(usize, u64, u64 .find_map(|(index, (previous, next))| (previous >= next).then_some((index, previous, next))) } +/// Add a non-empty range to the run of consecutive, sorted, disjoint `Range` +/// segments being folded into one [`U64Segment::Ranges`], merging an adjacent +/// one and flushing the run first when `range` does not continue it. +/// +/// A table whose deletions cluster is written as many small `Range` segments +/// (see `RowIdSequence::use_range_segments`). Kept as separate segments they +/// would cost a heap-allocated enum each and linear scans in every lookup; +/// folded, they are two `u32` per range and a binary search. +fn push_range(pending: &mut Vec>, out: &mut Vec, range: Range) { + match pending.last_mut() { + Some(last) if last.end == range.start => last.end = range.end, + Some(last) if last.end > range.start => { + flush_ranges(pending, out); + pending.push(range); + } + _ => pending.push(range), + } +} + +fn flush_ranges(pending: &mut Vec>, out: &mut Vec) { + match U64Segment::from_sorted_ranges(pending.as_slice()) { + Some(segment) => out.push(segment), + // A single range, or a run wider than the compact form can index. + None => out.extend(pending.iter().cloned().map(U64Segment::Range)), + } + pending.clear(); +} + impl TryFrom for RowIdSequence { type Error = Error; + /// Goes through [`read_row_ids`] so the proto and wire paths fold `Range` + /// runs the same way. fn try_from(pb: pb::RowIdSequence) -> Result { - let segments = pb - .segments - .into_iter() - .map(U64Segment::try_from) - .collect::>>()?; - // Each segment length fits a usize on its own, but the total need not fit a u64. - // Reject that here so `RowIdSequence::len()` stays total for anything decoded. - segments - .iter() - .try_fold(0_u64, |total, segment| { - total.checked_add(segment.len() as u64) - }) - .ok_or_else(|| { - corrupt_row_id_metadata(format!( - "row ID sequence of {} segments has a total length exceeding u64::MAX", - segments.len() - )) - })?; - Ok(Self(segments)) + read_row_ids(&pb.encode_to_vec()) } } @@ -215,9 +230,26 @@ impl TryFrom for EncodedU64Array { impl From for pb::RowIdSequence { fn from(sequence: RowIdSequence) -> Self { - Self { - segments: sequence.0.into_iter().map(pb::U64Segment::from).collect(), - } + let segments = sequence + .0 + .into_iter() + .flat_map(|segment| match segment { + // On the wire a run of ranges is just its `Range` segments, which + // every reader understands; `read_row_ids` folds them back + // together. + U64Segment::Ranges { range, runs } => runs + .present_ranges() + .map(|present| { + U64Segment::Range( + range.start + present.start as u64..range.start + present.end as u64, + ) + }) + .map(pb::U64Segment::from) + .collect::>(), + other => vec![pb::U64Segment::from(other)], + }) + .collect(); + Self { segments } } } @@ -248,6 +280,13 @@ impl From for pb::U64Segment { }, )), }, + // A run of ranges has no message of its own: `pb::RowIdSequence` + // writes it as its `Range` segments. Nothing else produces this + // variant (`from_slice`, `mask` and `slice` never pick it), so a + // standalone conversion is a caller bug rather than an encoding. + U64Segment::Ranges { .. } => { + unreachable!("U64Segment::Ranges is only serialized as part of a RowIdSequence") + } U64Segment::SortedArray(array) => Self { segment: Some(pb::u64_segment::Segment::SortedArray(array.into())), }, @@ -304,11 +343,85 @@ pub fn write_row_ids(sequence: &RowIdSequence) -> Vec { } /// Deserialize a rowid sequence from some bytes. -pub fn read_row_ids(reader: &[u8]) -> Result { - let pb_sequence = pb::RowIdSequence::decode(reader).map_err(|error| { +/// +/// Walks the `RowIdSequence` wire format directly: a sequence written as +/// `Range` segments can hold tens of millions of them, and decoding through +/// `pb::RowIdSequence` would materialize a proto message per segment before +/// folding. Plain `Range` segments are read straight into the fold, everything +/// else goes through prost. +pub fn read_row_ids(mut buf: &[u8]) -> Result { + let corrupt = |error: prost::DecodeError| { corrupt_row_id_metadata(format!("failed to decode row ID sequence: {error}")) - })?; - RowIdSequence::try_from(pb_sequence) + }; + let mut segments = Vec::new(); + let mut pending: Vec> = Vec::new(); + while !buf.is_empty() { + let (field, wire_type) = decode_key(&mut buf).map_err(corrupt)?; + if field != 1 || wire_type != WireType::LengthDelimited { + return Err(corrupt_row_id_metadata(format!( + "unexpected field {field} with wire type {wire_type:?} in row ID sequence" + ))); + } + let len = decode_varint(&mut buf).map_err(corrupt)? as usize; + if len > buf.len() { + return Err(corrupt_row_id_metadata(format!( + "row ID segment of {len} bytes exceeds the {} remaining", + buf.len() + ))); + } + let (segment, rest) = buf.split_at(len); + buf = rest; + match plain_range(segment) { + Some(range) => push_range(&mut pending, &mut segments, range), + None => { + flush_ranges(&mut pending, &mut segments); + let segment = pb::U64Segment::decode(segment).map_err(corrupt)?; + segments.push(U64Segment::try_from(segment)?); + } + } + } + flush_ranges(&mut pending, &mut segments); + // Each segment length fits a usize on its own, but the total need not fit a u64. + // Reject that here so `RowIdSequence::len()` stays total for anything decoded. + segments + .iter() + .try_fold(0_u64, |total, segment| { + total.checked_add(segment.len() as u64) + }) + .ok_or_else(|| { + corrupt_row_id_metadata(format!( + "row ID sequence of {} segments has a total length exceeding u64::MAX", + segments.len() + )) + })?; + Ok(RowIdSequence(segments)) +} + +/// The non-empty range of a serialized `U64Segment` that holds nothing but a +/// `Range`; `None` for any other shape, which prost decodes instead. +fn plain_range(mut bytes: &[u8]) -> Option> { + let (field, wire_type) = decode_key(&mut bytes).ok()?; + if field != 1 || wire_type != WireType::LengthDelimited { + return None; + } + let len = decode_varint(&mut bytes).ok()? as usize; + if len != bytes.len() { + return None; + } + let (mut start, mut end) = (0_u64, 0_u64); + while !bytes.is_empty() { + let (field, wire_type) = decode_key(&mut bytes).ok()?; + if wire_type != WireType::Varint { + return None; + } + let value = decode_varint(&mut bytes).ok()?; + match field { + 1 => start = value, + 2 => end = value, + _ => return None, + } + } + (start < end).then_some(start..end) } #[cfg(test)] @@ -407,7 +520,14 @@ mod test { let sequence2 = read_row_ids(&serialized).unwrap(); - assert_eq!(sequence.0, sequence2.0); + // Decoding folds the two leading `Range` segments into one `Ranges` + // segment; everything else round-trips unchanged. + assert!(matches!(sequence2.0[0], U64Segment::Ranges { .. })); + assert_eq!(&sequence2.0[1..], &sequence.0[2..]); + assert_eq!( + sequence2.iter().collect::>(), + sequence.iter().collect::>() + ); } proptest! { @@ -697,4 +817,67 @@ mod test { "SortedArray values are not strictly increasing at indices 0 and 1", ); } + + #[test] + fn test_ranges_are_written_as_range_segments_and_folded_back() { + // 1000..1100, 1300..1500, 1650..2000 as one compact segment. + let ranges = vec![1000..1100, 1300..1500, 1650..2000]; + let sequence = RowIdSequence(vec![U64Segment::from_sorted_ranges(&ranges).unwrap()]); + let encoded = write_row_ids(&sequence); + + // On the wire: three plain `Range` segments, readable by any version. + let wire = pb::RowIdSequence::decode(encoded.as_slice()).unwrap(); + let wire_ranges: Vec> = wire + .segments + .iter() + .map(|segment| match &segment.segment { + Some(pb::u64_segment::Segment::Range(range)) => range.start..range.end, + other => panic!("expected a Range segment, got {other:?}"), + }) + .collect(); + assert_eq!(wire_ranges, ranges); + + // Decoding folds them back into the compact segment. + let decoded = read_row_ids(encoded.as_slice()).unwrap(); + assert_eq!(decoded, sequence); + assert_eq!( + decoded.iter().collect::>(), + sequence.iter().collect::>() + ); + + // A lone range, adjacent ranges and non-range neighbours: adjacent + // ranges merge, a single range stays a `Range`, others are untouched. + let mixed = pb::RowIdSequence { + segments: vec![ + U64Segment::Range(0..10), + U64Segment::Range(10..20), + U64Segment::SortedArray(vec![100, 200].into()), + U64Segment::Range(300..310), + U64Segment::Range(320..330), + U64Segment::Range(340..350), + ] + .into_iter() + .map(pb::U64Segment::from) + .collect(), + }; + let folded = RowIdSequence::try_from(mixed).unwrap(); + assert!(matches!( + folded.0.as_slice(), + [ + U64Segment::Range(_), + U64Segment::SortedArray(_), + U64Segment::Ranges { .. } + ] + )); + assert_eq!(folded.0[0], U64Segment::Range(0..20)); + assert_eq!( + folded.iter().collect::>(), + (0..20) + .chain([100, 200]) + .chain(300..310) + .chain(320..330) + .chain(340..350) + .collect::>() + ); + } } diff --git a/rust/lance-table/src/transaction.rs b/rust/lance-table/src/transaction.rs index 74fba1b33ba..97d63fa5caf 100644 --- a/rust/lance-table/src/transaction.rs +++ b/rust/lance-table/src/transaction.rs @@ -40,6 +40,7 @@ mod validate; pub(crate) mod test_support; pub use builder::{Transaction, TransactionBuilder}; +pub use manifest_build::RANGE_SEGMENTS_CONFIG_KEY; pub use operation::{ DataOverlayGroup, DataReplacementGroup, Operation, RewriteGroup, RewrittenIndex, UpdateMode, UpdatedFragmentOffsets, diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index de14803cafa..5551bd36e64 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -16,7 +16,7 @@ use crate::feature_flags::{ }; use crate::format::overlay::{OverlayCoverage, TOMBSTONE_FIELD_ID}; use crate::format::{ - DataFile, DataStorageFormat, Fragment, IndexMetadata, Manifest, ManifestBuildConfig, + DataFile, DataStorageFormat, Fragment, IndexMetadata, Manifest, ManifestBuildConfig, RowIdMeta, overlay::DataOverlayFile, }; use crate::io::{ @@ -24,6 +24,7 @@ use crate::io::{ manifest::{read_manifest, read_manifest_indexes}, }; use crate::rowids::version::build_version_meta; +use crate::rowids::{read_row_ids, write_row_ids}; use crate::system_index::frag_reuse::FRAG_REUSE_INDEX_NAME; use crate::system_index::is_system_index; use crate::system_index::mem_wal::{ @@ -52,6 +53,72 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Arc; use uuid::Uuid; +/// Table config key that opts a table into writing clustered deletions as +/// `Range` segments (`U64Segment::Ranges`). Set it to `true` with +/// `update_config`. +/// +/// Manifests written afterwards re-encode bitmap row id segments as runs of +/// `Range` segments where that is smaller. That is the +/// existing wire format, so any version reads the result, but a reader without +/// the compact in-memory form handles thousands of segments per fragment +/// slowly. That is why it is opt-in rather than the default. +/// +/// Setting the key back to `false` only stops further conversions. Fragments +/// already written as `Range` segments keep that encoding until something else +/// rewrites their row ids (compaction, for example), so coordinate reader +/// upgrades before enabling it. +pub const RANGE_SEGMENTS_CONFIG_KEY: &str = "lance.row_ids.range_segments"; + +fn range_segments_enabled(manifest: &Manifest) -> bool { + manifest + .config + .get(RANGE_SEGMENTS_CONFIG_KEY) + .is_some_and(|value| str_is_truthy(value)) +} + +/// Re-encode inline row id sequences as runs of `Range` segments when the +/// table has opted in. +/// +/// Fragments that the previous manifest already re-encoded are left alone +/// (their inline bytes are the very same allocation), so a commit pays only +/// for the fragments it changed; enabling the key re-encodes every fragment +/// once. +fn apply_range_segments( + manifest: &mut Manifest, + current_manifest: Option<&Manifest>, +) -> Result<()> { + if !range_segments_enabled(manifest) { + return Ok(()); + } + let unchanged: HashMap = current_manifest + .filter(|current| range_segments_enabled(current)) + .map(|current| { + current + .fragments + .iter() + .filter_map(|fragment| match &fragment.row_id_meta { + Some(RowIdMeta::Inline(data)) => Some((fragment.id, data.as_ptr())), + _ => None, + }) + .collect() + }) + .unwrap_or_default(); + let fragments = Arc::make_mut(&mut manifest.fragments); + for fragment in fragments.iter_mut() { + let Some(RowIdMeta::Inline(data)) = &fragment.row_id_meta else { + continue; + }; + if unchanged.get(&fragment.id) == Some(&data.as_ptr()) { + continue; + } + let mut sequence = read_row_ids(&data[..])?; + if sequence.use_range_segments() { + fragment.row_id_meta = Some(RowIdMeta::Inline(write_row_ids(&sequence).into())); + } + } + Ok(()) +} + impl Transaction { pub(super) fn fragments_with_ids<'a, T>( new_fragments: T, @@ -1539,6 +1606,8 @@ impl Transaction { _ => {} } + apply_range_segments(&mut manifest, current_manifest)?; + // Handle UpdateBases operation to update manifest base_paths if let Operation::UpdateBases { new_bases } = &self.operation { // Validate and add new base paths to the manifest diff --git a/rust/lance/src/dataset/delta.rs b/rust/lance/src/dataset/delta.rs index b9e12da2482..7092d9bce02 100644 --- a/rust/lance/src/dataset/delta.rs +++ b/rust/lance/src/dataset/delta.rs @@ -783,6 +783,19 @@ fn segment_ids<'a>( let (base, start) = (range.start, next_value.max(range.start)); Box::new((start..range.end).filter(move |&v| bitmap.get((v - base) as usize))) } + U64Segment::Ranges { range, runs } => { + let (base, start) = (range.start, next_value.max(range.start)); + // Walk the present ranges, clipping the one the resume point falls in. + Box::new( + runs.present_ranges() + .filter_map(move |present| { + let first = (base + present.start as u64).max(start); + let end = base + present.end as u64; + (first < end).then_some(first..end) + }) + .flatten(), + ) + } U64Segment::SortedArray(array) | U64Segment::Array(array) => { Box::new((consumed..array.len()).filter_map(move |i| array.get(i))) } diff --git a/rust/lance/src/dataset/rowids.rs b/rust/lance/src/dataset/rowids.rs index f6c98b7925d..b09e4933d5c 100644 --- a/rust/lance/src/dataset/rowids.rs +++ b/rust/lance/src/dataset/rowids.rs @@ -277,7 +277,8 @@ mod test { use std::ops::Range; use crate::dataset::{ - ReadParams, UpdateBuilder, WriteMode, WriteParams, builder::DatasetBuilder, + ProjectionRequest, ReadParams, UpdateBuilder, WriteMode, WriteParams, + builder::DatasetBuilder, }; use super::*; @@ -296,6 +297,7 @@ mod test { use lance_datagen::Dimension; use lance_index::{IndexType, scalar::ScalarIndexParams}; use lance_io::object_store::ObjectStoreParams; + use lance_table::rowids::segment::U64Segment; use std::collections::HashMap; use std::collections::HashSet; @@ -994,6 +996,117 @@ mod test { dataset.delete(expr).await.unwrap(); } + #[tokio::test] + async fn test_range_segments_are_opt_in() { + use lance_table::format::pb; + use lance_table::transaction::RANGE_SEGMENTS_CONFIG_KEY; + use prost::Message; + + /// Wire-level segment kinds of every fragment, and whether the decoded + /// form folded any of them into a compact `Ranges` segment. + fn segment_shapes(dataset: &Dataset) -> (Vec<&'static str>, bool) { + let mut kinds = Vec::new(); + let mut folded = false; + for fragment in dataset.manifest.fragments.iter() { + let Some(RowIdMeta::Inline(data)) = &fragment.row_id_meta else { + panic!("expected inline row ids, got {:?}", fragment.row_id_meta); + }; + let wire = pb::RowIdSequence::decode(&data[..]).unwrap(); + kinds.extend(wire.segments.iter().map(|segment| match segment.segment { + Some(pb::u64_segment::Segment::Range(_)) => "range", + Some(pb::u64_segment::Segment::RangeWithBitmap(_)) => "bitmap", + _ => "other", + })); + folded |= read_row_ids(&data[..]) + .unwrap() + .segments() + .iter() + .any(|segment| matches!(segment, U64Segment::Ranges { .. })); + } + (kinds, folded) + } + + // Deleting a contiguous block inside every fragment leaves each row id + // sequence as a range with one long run of holes, which `delete` + // encodes as a bitmap segment; compaction carries those segments over. + let mut dataset = lance_datagen::gen_batch() + .col("i", lance_datagen::array::step::()) + .into_ram_dataset_with_params( + FragmentCount::from(4), + FragmentRowCount::from(500), + Some(WriteParams { + max_rows_per_file: 500, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + delete(&mut dataset, "i % 500 >= 100 and i % 500 < 400").await; + compact(&mut dataset, 5000).await; + let map_before = scan_rowid_map(&dataset).await; + assert_eq!(map_before.len(), 800); + + // Without the opt-in the bitmaps stay. + let (kinds, folded) = segment_shapes(&dataset); + assert!(kinds.contains(&"bitmap") && !folded, "{kinds:?}"); + + // Opting in re-encodes every eligible fragment in the same commit: on + // the wire only plain `Range` segments, folded back when decoded. + dataset + .update_config([(RANGE_SEGMENTS_CONFIG_KEY, "true")]) + .await + .unwrap(); + let (kinds, folded) = segment_shapes(&dataset); + assert!(kinds.iter().all(|kind| *kind == "range"), "{kinds:?}"); + assert!( + folded, + "fragments with 300 contiguous holes each should fold into Ranges" + ); + assert_eq!( + dataset.manifest().reader_feature_flags, + lance_table::feature_flags::FLAG_STABLE_ROW_IDS, + "the wire format did not change, so no new reader flag" + ); + + // Reads see the same row ids; a take by row id resolves through the + // compact segments. + let map_after = scan_rowid_map(&dataset).await; + assert_eq!(map_before, map_after); + let mut sample: Vec = map_before.keys().copied().collect(); + sample.sort_unstable(); + let sample: Vec = sample.into_iter().step_by(97).collect(); + let taken = dataset + .take_rows( + &sample, + ProjectionRequest::from_columns(["i"], dataset.schema()), + ) + .await + .unwrap(); + assert_eq!(taken.num_rows(), sample.len()); + let index = get_row_id_index(&Arc::new(dataset.clone())) + .await + .unwrap() + .unwrap(); + for row_id in &sample { + assert!(index.get(*row_id).unwrap().is_some(), "row id {row_id}"); + } + + // Later commits keep re-encoding only what changed, and the manifest + // re-read from storage folds the same way. + delete(&mut dataset, "i = 10").await; + let reopened = dataset + .checkout_version(dataset.manifest().version) + .await + .unwrap(); + let (kinds, folded) = segment_shapes(&reopened); + assert!( + kinds.iter().all(|kind| *kind == "range") && folded, + "{kinds:?}" + ); + assert_eq!(scan_rowid_map(&reopened).await.len(), 799); + } + #[tokio::test] async fn test_stable_row_id_after_multiple_deletion_and_compaction() { async fn delete(dataset: &mut Dataset, expr: &str) {