diff --git a/Cargo.lock b/Cargo.lock index e6441ee06e9..4f22e37c93c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -132,7 +132,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -143,7 +143,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -1455,7 +1455,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2764,7 +2764,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2916,7 +2916,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3460,9 +3460,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.15" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27" dependencies = [ "atomic-waker", "bytes", @@ -4145,7 +4145,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5681,7 +5681,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -7529,9 +7529,9 @@ dependencies = [ [[package]] name = "rkyv" -version = "0.8.16" +version = "0.8.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73389e0c99e664f919275ab5b5b0471391fe9a8de61e1dff9b1eaf56a90f16e3" +checksum = "d9776093b7ca170454ab1406954f7b7d97a57c51dc6c0642957fb2ef25c2d399" dependencies = [ "bytecheck", "bytes", @@ -7548,13 +7548,13 @@ dependencies = [ [[package]] name = "rkyv_derive" -version = "0.8.16" +version = "0.8.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d2ed0b54125315fb36bd021e82d314d1c126548f871634b483f46b31d13cac6" +checksum = "1c25ef604ac7dd839d44d64648952ea23c97866f124ff671b0ed2cf3ad9bb06e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.118", + "syn 3.0.3", ] [[package]] @@ -7685,7 +7685,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -7744,7 +7744,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -8283,7 +8283,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -8553,6 +8553,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "sync_wrapper" version = "1.0.2" @@ -8627,10 +8638,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -9654,7 +9665,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/rust/lance-core/src/utils.rs b/rust/lance-core/src/utils.rs index c202329838c..51d72c3dae3 100644 --- a/rust/lance-core/src/utils.rs +++ b/rust/lance-core/src/utils.rs @@ -15,6 +15,7 @@ pub mod hash; pub mod io_stats; pub mod parse; pub mod path; +pub mod row_addr_remap; pub mod tempfile; pub mod testing; pub mod tokio; diff --git a/rust/lance-core/src/utils/row_addr_remap.rs b/rust/lance-core/src/utils/row_addr_remap.rs new file mode 100644 index 00000000000..1a723077cd1 --- /dev/null +++ b/rust/lance-core/src/utils/row_addr_remap.rs @@ -0,0 +1,471 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Compact row-address remapping for compaction. +//! +//! Compaction rewrites rows into new fragments, so indices that store physical +//! row addresses need an old-address to new-address mapping without building an +//! O(total rows) `HashMap>`. +//! +//! Layout: +//! +//! * Old rows: `old_fragment_id -> (old_offsets, old_rows_before)` +//! * `old_offsets`: rewritten old row offsets in this old fragment. +//! * `old_rows_before`: rewritten row count before this old fragment. +//! * New rows: ordered new-fragment ranges +//! `(fragment_id, new_rows_before, physical_rows)` +//! * `new_rows_before`: rewritten row count before this new fragment. +//! +//! Lookup: +//! +//! * An address whose fragment was not rewritten returns `None`. +//! * For an address whose fragment was rewritten: +//! * Read `(old_offsets, old_rows_before)` from the old-row layout. +//! * If `offset` is not in `old_offsets`, return `Some(None)` because the +//! row was deleted. +//! * Otherwise, `old_offsets.rank(offset) - 1` is this row's 0-based +//! position among rewritten old rows in this old fragment. Add +//! `old_rows_before` to get `k`, the row's 0-based position among all +//! rewritten old rows. +//! * In the new-row layout, find the range +//! `(fragment_id, new_rows_before, physical_rows)` where +//! `new_rows_before <= k < new_rows_before + physical_rows`. +//! * The new address is `(fragment_id, k - new_rows_before)`. +//! +//! Ordering: +//! +//! Compact remap does not store each old-to-new row mapping. It computes `k` +//! from the old-row layout, then maps it to the k-th row written to the new +//! fragments. This requires the reader-to-writer pipeline to preserve row order. +//! +//! * `old_frag_ids` must match the order old fragments are read. Within each +//! old fragment, rewritten rows are interpreted by ascending old row offset. +//! * `new_frags` must match the order new rows are written. +//! * Current compaction satisfies this because it scans selected fragments in +//! order and writes the resulting stream without reordering rows. + +use crate::utils::address::RowAddress; +use crate::{Error, Result}; +use roaring::{RoaringBitmap, RoaringTreemap}; +use std::collections::HashMap; + +/// A queryable row-address remapping with the exact semantics of +/// `HashMap>::get(&addr).copied()`: +/// +/// * `None` — the address is not affected by this remap (keep it unchanged) +/// * `Some(None)` — the row was deleted +/// * `Some(Some(addr))` — the row moved to `addr` +#[derive(Clone)] +pub enum RowAddrRemap { + /// Compact, `O(#fragments)` remap built from per-group rewritten-row + /// bitmaps and new-fragment layouts. + Compact(CompactRowAddrRemap), + /// Full materialized old-to-new address map. Uses `O(#rows)` memory. + Direct(HashMap>), +} + +impl RowAddrRemap { + pub fn compact(groups: impl IntoIterator) -> Result { + Ok(Self::Compact(CompactRowAddrRemap::new(groups)?)) + } + + /// Build a remap from a fully materialized old-to-new address map. + pub fn direct(map: HashMap>) -> Self { + Self::Direct(map) + } + + /// An empty remap that leaves every address unchanged. + pub fn empty() -> Self { + Self::Direct(HashMap::new()) + } + + /// Look up `addr`. See [`RowAddrRemap`] for the tri-state return semantics. + #[inline] + pub fn get(&self, addr: u64) -> Option> { + match self { + Self::Compact(c) => c.get(addr), + Self::Direct(m) => m.get(&addr).copied(), + } + } + + pub fn is_empty(&self) -> bool { + match self { + Self::Compact(c) => c.is_empty(), + Self::Direct(m) => m.is_empty(), + } + } + + pub fn affected_fragments(&self) -> RoaringBitmap { + match self { + Self::Compact(c) => RoaringBitmap::from_iter(c.frag_to_group.keys().copied()), + Self::Direct(m) => RoaringBitmap::from_iter(m.keys().map(|addr| (addr >> 32) as u32)), + } + } + + pub fn fully_deleted_fragments(&self) -> Option { + match self { + Self::Compact(c) => c.fully_deleted_fragments(), + Self::Direct(m) => { + if m.values().all(|v| v.is_none()) { + Some(RoaringBitmap::from_iter( + m.keys().map(|addr| (addr >> 32) as u32), + )) + } else { + None + } + } + } + } +} + +/// Input describing one rewrite group: the old row addresses that were +/// rewritten plus the fragment layout before/after the rewrite. +pub struct GroupInput { + /// Old row addresses that were read and re-written into the new fragments. + pub rewritten_old_row_addrs: RoaringTreemap, + /// Old fragment ids covered by this group. + pub old_frag_ids: Vec, + /// New fragments produced by this group, as `(fragment_id, physical_rows)`, + pub new_frags: Vec<(u32, u32)>, +} + +#[derive(Clone)] +struct GroupRemap { + /// Old fragment id -> (rewritten old row offsets in that fragment, + /// rewritten row count before this fragment in the group). + frags: HashMap, + /// New fragment ranges as `(fragment_id, rewritten_rows_before, physical_rows)`, + /// used to map a rewritten row's group-local index to its new address via binary search. + new_frag_row_ranges: Vec<(u32, u64, u32)>, +} + +impl GroupRemap { + fn new(input: GroupInput) -> Result { + // `compute_new_addr` maps a rewritten row's group-local index to a new + // address by accumulating `physical_rows` in `new_frags` order, so that + // order must be the order rows were written. New fragment ids are + // reserved monotonically in write order (see `reserve_fragment_ids` in + // compaction), so ascending id is a proxy for write order; reject any + // input that violates it before it can silently misplace addresses. + let mut new_frag_row_ranges = Vec::with_capacity(input.new_frags.len()); + let mut rewritten_rows_before = 0u64; + let mut prev_frag_id: Option = None; + for (frag_id, physical_rows) in input.new_frags { + if physical_rows == 0 { + continue; + } + if let Some(prev) = prev_frag_id + && frag_id <= prev + { + return Err(Error::invalid_input(format!( + "compaction new fragments must be in ascending id (write) order, but fragment {frag_id} follows {prev}", + ))); + } + prev_frag_id = Some(frag_id); + new_frag_row_ranges.push((frag_id, rewritten_rows_before, physical_rows)); + rewritten_rows_before += physical_rows as u64; + } + let total_new_rows = rewritten_rows_before; + + let mut per_frag: HashMap = input + .rewritten_old_row_addrs + .bitmaps() + .map(|(frag_id, bitmap)| (frag_id, bitmap.clone())) + .collect(); + let mut frags = HashMap::new(); + let mut rewritten_rows_before = 0u64; + for &frag_id in &input.old_frag_ids { + // A fragment with no rewritten rows (fully deleted) contributes + // nothing to the rewritten row sequence. + if let Some(bitmap) = per_frag.remove(&frag_id) { + let num_rewritten_rows = bitmap.len(); + frags.insert(frag_id, (bitmap, rewritten_rows_before)); + rewritten_rows_before += num_rewritten_rows; + } + } + // Rewritten old row addresses must reference only fragments listed in `old_frag_ids`. + if !per_frag.is_empty() { + return Err(Error::invalid_input(format!( + "compaction rewritten old row addresses reference fragments {:?} not in the rewrite group's old fragments {:?}", + per_frag.keys().collect::>(), + input.old_frag_ids, + ))); + } + + // Rewritten old rows are mapped positionally onto the new rows, so the + // two counts must match exactly + let total_rewritten_old_rows = input.rewritten_old_row_addrs.len(); + if total_new_rows != total_rewritten_old_rows { + return Err(Error::invalid_input(format!( + "compaction rewrote {total_rewritten_old_rows} old rows from fragments {:?} but the new fragments hold {total_new_rows} rows", + input.old_frag_ids, + ))); + } + + Ok(Self { + frags, + new_frag_row_ranges, + }) + } + + fn compute_new_addr(&self, rewritten_row_index: u64) -> u64 { + let idx = + match self + .new_frag_row_ranges + .binary_search_by(|(_, rewritten_rows_before, _)| { + rewritten_rows_before.cmp(&rewritten_row_index) + }) { + Ok(i) => i, + Err(i) => i - 1, + }; + let (frag_id, rewritten_rows_before, _rows) = self.new_frag_row_ranges[idx]; + let offset = (rewritten_row_index - rewritten_rows_before) as u32; + u64::from(RowAddress::new_from_parts(frag_id, offset)) + } + + /// Compute the new address for an old row in this group. + /// Returns `None` if the old row was not rewritten. + #[inline] + fn get(&self, frag: u32, offset: u32) -> Option { + match self.frags.get(&frag) { + Some((bitmap, rewritten_rows_before)) if bitmap.contains(offset) => { + let rewritten_row_index = rewritten_rows_before + bitmap.rank(offset) - 1; + Some(self.compute_new_addr(rewritten_row_index)) + } + _ => None, + } + } +} + +/// Compact remap backed by per-group rewritten row bitmaps + new-fragment layouts. +#[derive(Clone)] +pub struct CompactRowAddrRemap { + groups: Vec, + /// Old fragment id -> index into `groups`. Size is O(#fragments), not rows. + frag_to_group: HashMap, +} + +impl CompactRowAddrRemap { + fn new(groups: impl IntoIterator) -> Result { + let mut frag_to_group = HashMap::new(); + let mut group_remaps = Vec::new(); + for input in groups { + let gi = group_remaps.len(); + for &frag_id in &input.old_frag_ids { + frag_to_group.insert(frag_id, gi); + } + group_remaps.push(GroupRemap::new(input)?); + } + Ok(Self { + groups: group_remaps, + frag_to_group, + }) + } + + #[inline] + pub fn get(&self, addr: u64) -> Option> { + let frag = (addr >> 32) as u32; + // Not in any rewrite group -> unaffected by this remap. + let gi = *self.frag_to_group.get(&frag)?; + Some(self.groups[gi].get(frag, addr as u32)) + } + + pub fn is_empty(&self) -> bool { + self.groups.is_empty() + } + + fn fully_deleted_fragments(&self) -> Option { + // A group with any rewritten row moved at least one row. + if self.groups.iter().any(|g| !g.frags.is_empty()) { + return None; + } + Some(RoaringBitmap::from_iter(self.frag_to_group.keys().copied())) + } +} + +// --- rerun fork additions, not present in lance-format/lance release-9.0.0 --- +// +// `FragReuseIndex` stores these and is `Debug` + `DeepSizeOf`; upstream does not need +// either yet because it still holds `Vec>>` there. + +impl std::fmt::Debug for RowAddrRemap { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Compact(compact) => f + .debug_struct("RowAddrRemap::Compact") + .field("groups", &compact.num_groups()) + .field("fragments", &compact.num_fragments()) + .finish(), + Self::Direct(map) => f + .debug_struct("RowAddrRemap::Direct") + .field("entries", &map.len()) + .finish(), + } + } +} + +impl CompactRowAddrRemap { + /// Number of rewrite groups held. + pub fn num_groups(&self) -> usize { + self.groups.len() + } + + /// Number of old fragments covered. This is what the structure scales with. + pub fn num_fragments(&self) -> usize { + self.frag_to_group.len() + } +} + +impl crate::deepsize::DeepSizeOf for RowAddrRemap { + fn deep_size_of_children(&self, cx: &mut crate::deepsize::Context) -> usize { + match self { + // Bitmaps dominate; roaring is not `DeepSizeOf`, so approximate from its + // own serialized size rather than under-report it as zero. + Self::Compact(compact) => { + compact + .groups + .iter() + .map(|group| { + group + .frags + .values() + .map(|(bitmap, _)| bitmap.serialized_size() + 16) + .sum::() + + group.new_frag_row_ranges.len() * 16 + }) + .sum::() + + compact.frag_to_group.len() * 12 + } + Self::Direct(map) => map.deep_size_of_children(cx), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn addr(frag: u32, offset: u32) -> u64 { + u64::from(RowAddress::new_from_parts(frag, offset)) + } + + #[test] + fn test_compact_lookup() { + // Group A: out-of-order old frags [4, 3], split new frags (11 empty), + // some deletions. frag 4 (5 rows) keeps 0,2,4; frag 3 keeps 0,1, so the + // rewritten rows (4,0)(4,2)(4,4)(3,0)(3,1) go to new frags 10(2), 12(3). + // Group B is a fully-deleted fragment. + let group_a = GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter([ + addr(4, 0), + addr(4, 2), + addr(4, 4), + addr(3, 0), + addr(3, 1), + ]), + old_frag_ids: vec![4, 3], + new_frags: vec![(10, 2), (11, 0), (12, 3)], + }; + let group_b = GroupInput { + rewritten_old_row_addrs: RoaringTreemap::new(), + old_frag_ids: vec![7], + new_frags: vec![], + }; + let remap = RowAddrRemap::compact([group_a, group_b]).unwrap(); + + // Moves, in rewrite order; frag 4 comes first despite the larger id. + assert_eq!(remap.get(addr(4, 0)), Some(Some(addr(10, 0)))); + assert_eq!(remap.get(addr(4, 2)), Some(Some(addr(10, 1)))); + // Rank 2 skips the zero-row new fragment 11 and lands in fragment 12. + assert_eq!(remap.get(addr(4, 4)), Some(Some(addr(12, 0)))); + assert_eq!(remap.get(addr(3, 0)), Some(Some(addr(12, 1)))); + assert_eq!(remap.get(addr(3, 1)), Some(Some(addr(12, 2)))); + // Deleted offsets inside a rewritten fragment. + assert_eq!(remap.get(addr(4, 1)), Some(None)); + assert_eq!(remap.get(addr(4, 3)), Some(None)); + // Covered but fully-deleted fragment -> Some(None), not None. + assert_eq!(remap.get(addr(7, 0)), Some(None)); + // Fragment in no group -> unaffected. + assert_eq!(remap.get(addr(9, 0)), None); + assert!(!remap.is_empty()); + } + + #[test] + fn test_fragment_sets() { + // No rewritten rows at all: every covered fragment is fully deleted. + let dead = RowAddrRemap::compact([GroupInput { + rewritten_old_row_addrs: RoaringTreemap::new(), + old_frag_ids: vec![3, 7], + new_frags: vec![], + }]) + .unwrap(); + assert_eq!( + dead.fully_deleted_fragments(), + Some(RoaringBitmap::from_iter([3u32, 7u32])) + ); + assert_eq!( + dead.affected_fragments(), + RoaringBitmap::from_iter([3u32, 7u32]) + ); + + // At least one rewritten row -> not fully deleted, but both covered + // fragments (including the fully-deleted frag 1) are still affected. + let alive = RowAddrRemap::compact([GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter([addr(0, 0)]), + old_frag_ids: vec![0, 1], + new_frags: vec![(10, 1)], + }]) + .unwrap(); + assert!(alive.fully_deleted_fragments().is_none()); + assert_eq!( + alive.affected_fragments(), + RoaringBitmap::from_iter([0u32, 1u32]) + ); + } + + #[test] + fn test_compact_rejects_rewritten_addrs_outside_old_frags() { + // Rewritten addresses reference frag 5, not in old_frag_ids. The count + // still matches (2 == 2), so only the per-fragment split catches it. + let input = GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter([addr(0, 0), addr(5, 0)]), + old_frag_ids: vec![0], + new_frags: vec![(10, 2)], + }; + assert!(RowAddrRemap::compact([input]).is_err()); + } + + #[test] + fn test_compact_rejects_new_frags_out_of_write_order() { + // New fragments out of ascending id (write) order would make + // `compute_new_addr` accumulate rows in the wrong order, silently + // misplacing addresses. A zero-row fragment between them is ignored. + let input = GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter([addr(0, 0), addr(0, 1)]), + old_frag_ids: vec![0], + new_frags: vec![(12, 1), (11, 1)], + }; + assert!(RowAddrRemap::compact([input]).is_err()); + } + + #[test] + fn test_direct_and_empty() { + // Direct covers arbitrary maps the compact form can't express. + let mut map = HashMap::new(); + map.insert(addr(2, 0), Some(addr(9, 9))); + map.insert(addr(5, 1), None); + let remap = RowAddrRemap::direct(map); + assert_eq!(remap.get(addr(2, 0)), Some(Some(addr(9, 9)))); + assert_eq!(remap.get(addr(5, 1)), Some(None)); + assert_eq!(remap.get(addr(2, 1)), None); + // affected_fragments over an explicit map: the fragment of every key. + assert_eq!( + remap.affected_fragments(), + RoaringBitmap::from_iter([2u32, 5u32]) + ); + + let empty = RowAddrRemap::empty(); + assert!(empty.is_empty()); + assert_eq!(empty.get(addr(0, 0)), None); + } +} diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index 2f4fe69792a..6ecfd6e01a0 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -2395,7 +2395,7 @@ fn build_frag_reuse_mapping( row_ids: &UInt64Array, ) -> Option>> { let fri = fri?; - if fri.row_id_maps.is_empty() { + if fri.row_addr_maps.is_empty() { return None; } let mut mapping: HashMap> = HashMap::new(); diff --git a/rust/lance-table/src/system_index/frag_reuse.rs b/rust/lance-table/src/system_index/frag_reuse.rs index 40bbc4f58b6..4c6166e3e94 100644 --- a/rust/lance-table/src/system_index/frag_reuse.rs +++ b/rust/lance-table/src/system_index/frag_reuse.rs @@ -7,6 +7,7 @@ use arrow_array::cast::AsArray; use arrow_array::types::UInt64Type; use arrow_array::{Array, ArrayRef, PrimitiveArray, RecordBatch, UInt64Array}; use lance_core::deepsize::{Context, DeepSizeOf}; +use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_core::{Error, Result}; use lance_select::RowAddrTreeMap; use roaring::{RoaringBitmap, RoaringTreemap}; @@ -199,39 +200,61 @@ impl FragReuseIndexDetails { /// An index that stores row ID maps. /// A row ID map describes the mapping from old row address to new address after compactions. /// Each version contains the mapping for one round of compaction. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone)] pub struct FragReuseIndex { pub uuid: Uuid, - pub row_id_maps: Vec>>, + /// One remap per reuse version, oldest first. Order is load-bearing: each version is + /// applied to the previous version's output. + /// + /// Built as [`RowAddrRemap::Compact`] when the index is opened, which costs + /// O(#fragments) instead of the O(#rows) a materialized map would. A large compaction + /// here produced 676M row entries, 40 GB resident, which OOMs the pod that opens it. + pub row_addr_maps: Vec, pub details: FragReuseIndexDetails, } impl DeepSizeOf for FragReuseIndex { fn deep_size_of_children(&self, cx: &mut Context) -> usize { - self.row_id_maps.deep_size_of_children(cx) + self.details.deep_size_of_children(cx) + self.row_addr_maps.deep_size_of_children(cx) + self.details.deep_size_of_children(cx) } } impl FragReuseIndex { + /// Build from already-materialized maps, one per version. + /// + /// Kept for callers that hold maps already; it stores them as + /// [`RowAddrRemap::Direct`] and so costs O(#rows). Prefer + /// [`Self::new_from_remaps`] with [`RowAddrRemap::compact`]. pub fn new( uuid: Uuid, row_id_maps: Vec>>, details: FragReuseIndexDetails, + ) -> Self { + Self::new_from_remaps( + uuid, + row_id_maps.into_iter().map(RowAddrRemap::direct).collect(), + details, + ) + } + + pub fn new_from_remaps( + uuid: Uuid, + row_addr_maps: Vec, + details: FragReuseIndexDetails, ) -> Self { Self { uuid, - row_id_maps, + row_addr_maps, details, } } pub fn remap_row_id(&self, row_id: u64) -> Option { let mut mapped_value = Some(row_id); - for row_id_map in self.row_id_maps.iter() { + for row_addr_map in self.row_addr_maps.iter() { if mapped_value.is_some() { - mapped_value = row_id_map - .get(&mapped_value.unwrap()) - .copied() + mapped_value = row_addr_map + .get(mapped_value.unwrap()) .unwrap_or(mapped_value); } } diff --git a/rust/lance/src/dataset/optimize/remapping.rs b/rust/lance/src/dataset/optimize/remapping.rs index 266ac977a69..eb64149292a 100644 --- a/rust/lance/src/dataset/optimize/remapping.rs +++ b/rust/lance/src/dataset/optimize/remapping.rs @@ -216,7 +216,7 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { .await .unwrap(); - if frag_reuse_index.row_id_maps.is_empty() { + if frag_reuse_index.row_addr_maps.is_empty() { return Ok(()); } @@ -309,10 +309,23 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { // stale (an empty map makes `index::remap_index` return `Keep`). The map is // bounded by the rows the reuse index touched; addresses this index does not // store are simply never looked up. + // + // The compact remap deliberately cannot enumerate its keys: it stores per-fragment + // bitmaps, not rows, and treats any unlisted offset in a rewritten fragment as + // deleted, so "every key" is not a finite set it knows. Rebuild the per-row maps from + // the details here to get the key set. That is O(rows) again, but only on this path, + // which no production caller reaches; the cached open path above is what mattered. let composed_row_id_map: HashMap> = frag_reuse_index - .row_id_maps + .details + .versions .iter() - .flat_map(|row_id_map| row_id_map.keys().copied()) + .flat_map(|version| version.groups.iter()) + .flat_map(|group| { + let changed = + RoaringTreemap::deserialize_from(std::io::Cursor::new(&group.changed_row_addrs)) + .expect("fragment reuse index details were already parsed"); + transpose_row_ids_from_digest(changed, &group.old_frags, &group.new_frags).into_keys() + }) .map(|old_addr| (old_addr, frag_reuse_index.remap_row_id(old_addr))) .collect(); diff --git a/rust/lance/src/index/frag_reuse.rs b/rust/lance/src/index/frag_reuse.rs index 23a8fec5145..3ff49e7d01a 100644 --- a/rust/lance/src/index/frag_reuse.rs +++ b/rust/lance/src/index/frag_reuse.rs @@ -2,9 +2,9 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use crate::Dataset; -use crate::dataset::optimize::remapping::transpose_row_ids_from_digest; use crate::index::DatasetIndexExt; use lance_core::Error; +use lance_core::utils::row_addr_remap::{GroupInput, RowAddrRemap}; use lance_index::frag_reuse::{ FRAG_REUSE_DETAILS_FILE_NAME, FRAG_REUSE_INDEX_NAME, FragReuseGroup, FragReuseIndex, FragReuseIndexDetails, FragReuseVersion, @@ -14,7 +14,6 @@ use lance_table::format::pb::fragment_reuse_index_details::{Content, InlineConte use lance_table::format::pb::{ExternalFile, FragmentReuseIndexDetails}; use prost::Message; use roaring::{RoaringBitmap, RoaringTreemap}; -use std::collections::HashMap; use std::io::Cursor; use std::sync::Arc; use tokio::io::AsyncWriteExt; @@ -72,24 +71,33 @@ pub(crate) async fn open_frag_reuse_index( uuid: Uuid, details: &FragReuseIndexDetails, ) -> lance_core::Result { - let mut row_id_maps: Vec>> = - Vec::with_capacity(details.versions.len()); + // Build the compact form rather than a materialized per-row map. This runs on every + // index open and the result is cached, so the map's O(#rows) cost is paid by readers: + // one production payload here is 88 MB on disk and 40 GB once expanded. + let mut row_addr_maps: Vec = Vec::with_capacity(details.versions.len()); for version in &details.versions { - let mut row_id_map = HashMap::>::new(); + let mut groups = Vec::with_capacity(version.groups.len()); for group in version.groups.iter() { let cursor = Cursor::new(&group.changed_row_addrs); - let changed_row_addrs = RoaringTreemap::deserialize_from(cursor).unwrap(); - let group_row_id_map = transpose_row_ids_from_digest( - changed_row_addrs, - &group.old_frags, - &group.new_frags, - ); - row_id_map.extend(group_row_id_map); + let rewritten_old_row_addrs = RoaringTreemap::deserialize_from(cursor)?; + groups.push(GroupInput { + rewritten_old_row_addrs, + old_frag_ids: group.old_frags.iter().map(|frag| frag.id as u32).collect(), + new_frags: group + .new_frags + .iter() + .map(|frag| (frag.id as u32, frag.physical_rows as u32)) + .collect(), + }); } - row_id_maps.push(row_id_map); + row_addr_maps.push(RowAddrRemap::compact(groups)?); } - Ok(FragReuseIndex::new(uuid, row_id_maps, details.clone())) + Ok(FragReuseIndex::new_from_remaps( + uuid, + row_addr_maps, + details.clone(), + )) } pub(crate) async fn build_new_frag_reuse_index(