perf(frag-reuse): build the index-open path on RowAddrRemap (release-9, for redap 0.16)

rust/lance-core/src/utils/row_addr_remap.rs

@@ -283,6 +283,64 @@ impl CompactRowAddrRemap {
     }
 }
 
+// --- rerun fork additions, not present upstream ---
+//
+// `FragReuseIndex` stores these and is `Debug` + `DeepSizeOf`; upstream does not need
+// either yet because it still holds `Vec<HashMap<u64, Option<u64>>>` 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::<usize>()
+                            + group.new_frag_row_ranges.len() * 16
+                    })
+                    .sum::<usize>()
+                    + compact.frag_to_group.len() * 12
+            }
+            Self::Direct(map) => map.deep_size_of_children(cx),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

rust/lance-index/src/vector/bq/storage.rs

@@ -2396,7 +2396,7 @@ fn build_frag_reuse_mapping(
     row_ids: &UInt64Array,
 ) -> Option<HashMap<u64, Option<u64>>> {
     let fri = fri?;
-    if fri.row_id_maps.is_empty() {
+    if fri.row_addr_maps.is_empty() {
         return None;
     }
     let mut mapping: HashMap<u64, Option<u64>> = HashMap::new();

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,60 @@ 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<HashMap<u64, Option<u64>>>,
+    /// 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) rather than the O(#rows) a materialized map would. On one production
+    /// payload the map form was 676M entries and 40 GB resident, which OOMs the pod.
+    pub row_addr_maps: Vec<RowAddrRemap>,
     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`].
     pub fn new(
         uuid: Uuid,
         row_id_maps: Vec<HashMap<u64, Option<u64>>>,
         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<RowAddrRemap>,
+        details: FragReuseIndexDetails,
     ) -> Self {
         Self {
             uuid,
-            row_id_maps,
+            row_addr_maps,
             details,
         }
     }
 
     pub fn remap_row_id(&self, row_id: u64) -> Option<u64> {
         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);
             }
         }

rust/lance/src/dataset/optimize.rs

@@ -1780,7 +1780,13 @@ async fn rewrite_files(
             let captured_ids = row_ids_rx
                 .try_recv()
                 .map_err(|err| Error::internal(format!("Failed to receive row ids: {}", err)))?;
-            let row_addrs = captured_ids.row_addrs(None).into_owned();
+            let mut row_addrs = captured_ids.row_addrs(None).into_owned();
+            // Compaction rewrites long contiguous spans of addresses, which roaring
+            // stores as dense bitmap containers unless asked otherwise: one bit per row,
+            // so a large table costs tens of megabytes here. `optimize` converts those to
+            // run containers, a few bytes per span. This payload is read on every index
+            // open, so it is worth shrinking at rest.
+            row_addrs.optimize();
             let mut serialized = Vec::with_capacity(row_addrs.serialized_size());
             row_addrs.serialize_into(&mut serialized)?;
             Ok(Some(serialized))

rust/lance/src/dataset/optimize/remapping.rs

@@ -213,7 +213,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(());
     }
 
@@ -306,10 +306,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 is what mattered.
     let composed_row_id_map: HashMap<u64, Option<u64>> = 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();