feat(compaction): introduce RowAddrRemap structure to avoid remap OOM caused by HashMap (#7237)

Author: zhangyue19921010

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;

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

@@ -13,6 +13,7 @@ use datafusion::functions::string::contains::ContainsFunc;
 use datafusion::functions_nested::array_has;
 use datafusion::physical_plan::SendableRecordBatchStream;
 use datafusion_common::{Column, scalar::ScalarValue};
+use lance_core::utils::row_addr_remap::RowAddrRemap;
 use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
 use std::pin::Pin;
@@ -1108,7 +1109,7 @@ pub trait ScalarIndex: Send + Sync + std::fmt::Debug + Index + DeepSizeOf {
     /// Remap the row ids, creating a new remapped version of this index in `dest_store`
     async fn remap(
         &self,
-        mapping: &HashMap<u64, Option<u64>>,
+        mapping: &RowAddrRemap,
         dest_store: &dyn IndexStore,
     ) -> Result<CreatedIndex>;
 

rust/lance-index/src/scalar.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The Lance Authors
 
+use lance_core::utils::row_addr_remap::RowAddrRemap;
 use std::{
     any::Any,
     cmp::Reverse,
@@ -795,7 +796,7 @@ impl ScalarIndex for BitmapIndex {
     /// Remap the row ids, creating a new remapped version of this index in `dest_store`
     async fn remap(
         &self,
-        mapping: &HashMap<u64, Option<u64>>,
+        mapping: &RowAddrRemap,
         dest_store: &dyn IndexStore,
     ) -> Result<CreatedIndex> {
         let state = self.load_bitmap_index_state().await?;
@@ -1549,18 +1550,15 @@ impl BitmapIndexPlugin {
     /// Remaps every bitmap in a materialized bitmap-index state using row-id mappings.
     pub(crate) fn remap_bitmap_state(
         state: HashMap<ScalarValue, RowAddrTreeMap>,
-        mapping: &HashMap<u64, Option<u64>>,
+        mapping: &RowAddrRemap,
     ) -> HashMap<ScalarValue, RowAddrTreeMap> {
         state
             .into_iter()
             .map(|(key, bitmap)| {
                 let remapped_bitmap =
                     RowAddrTreeMap::from_iter(bitmap.row_addrs().unwrap().filter_map(|addr| {
                         let addr_as_u64 = u64::from(addr);
-                        mapping
-                            .get(&addr_as_u64)
-                            .copied()
-                            .unwrap_or(Some(addr_as_u64))
+                        mapping.get(addr_as_u64).unwrap_or(Some(addr_as_u64))
                     }));
                 (key, remapped_bitmap)
             })
@@ -1909,7 +1907,7 @@ mod tests {
     use lance_core::utils::{address::RowAddress, tempfile::TempObjDir};
     use lance_io::object_store::ObjectStore;
     use lance_select::RowSetOps;
-    use std::collections::HashMap;
+    use rstest::rstest;
 
     fn assert_state_roundtrips(state: &BitmapIndexState) {
         let mut buf = Vec::new();
@@ -2460,8 +2458,44 @@ mod tests {
         assert_eq!(red_rows_2, vec![0, 3, 6, 10, 11]);
     }
 
+    // frags 1 and 2 (3 rows each) are compacted into frag 3: the 6 rows are
+    // rewritten in order to frag 3 offsets 0..6.
+    fn bitmap_remap_compact() -> RowAddrRemap {
+        use lance_core::utils::row_addr_remap::GroupInput;
+        use roaring::RoaringTreemap;
+        RowAddrRemap::compact([GroupInput {
+            rewritten_old_row_addrs: RoaringTreemap::from_iter(
+                (0..3)
+                    .map(|o| u64::from(RowAddress::new_from_parts(1, o)))
+                    .chain((0..3).map(|o| u64::from(RowAddress::new_from_parts(2, o)))),
+            ),
+            old_frag_ids: vec![1, 2],
+            new_frags: vec![(3, 6)],
+        }])
+        .unwrap()
+    }
+
+    fn bitmap_remap_explicit() -> RowAddrRemap {
+        // The same mapping, listed out explicitly.
+        RowAddrRemap::direct(
+            (0..6u32)
+                .map(|i| {
+                    let (f, o) = if i < 3 { (1, i) } else { (2, i - 3) };
+                    (
+                        u64::from(RowAddress::new_from_parts(f, o)),
+                        Some(u64::from(RowAddress::new_from_parts(3, i))),
+                    )
+                })
+                .collect(),
+        )
+    }
+
+    // remap must behave identically whether the mapping is compact or explicit.
+    #[rstest]
+    #[case(bitmap_remap_compact())]
+    #[case(bitmap_remap_explicit())]
     #[tokio::test]
-    async fn test_remap_bitmap_with_null() {
+    async fn test_remap_bitmap_with_null(#[case] remap: RowAddrRemap) {
         use arrow_array::UInt32Array;
 
         // Create a temporary store.
@@ -2526,38 +2560,8 @@ mod tests {
         assert_eq!(index.index_map.len(), 2); // 2 non-null values (1 and 2)
         assert!(!index.null_map.is_empty()); // Should have null values
 
-        // Create a remap that simulates compaction of frags 1 and 2 into frag 3
-        let mut row_addr_map = HashMap::<u64, Option<u64>>::new();
-        row_addr_map.insert(
-            RowAddress::new_from_parts(1, 0).into(),
-            Some(RowAddress::new_from_parts(3, 0).into()),
-        );
-        row_addr_map.insert(
-            RowAddress::new_from_parts(1, 1).into(),
-            Some(RowAddress::new_from_parts(3, 1).into()),
-        );
-        row_addr_map.insert(
-            RowAddress::new_from_parts(1, 2).into(),
-            Some(RowAddress::new_from_parts(3, 2).into()),
-        );
-        row_addr_map.insert(
-            RowAddress::new_from_parts(2, 0).into(),
-            Some(RowAddress::new_from_parts(3, 3).into()),
-        );
-        row_addr_map.insert(
-            RowAddress::new_from_parts(2, 1).into(),
-            Some(RowAddress::new_from_parts(3, 4).into()),
-        );
-        row_addr_map.insert(
-            RowAddress::new_from_parts(2, 2).into(),
-            Some(RowAddress::new_from_parts(3, 5).into()),
-        );
-
         // Perform remap
-        index
-            .remap(&row_addr_map, test_store.as_ref())
-            .await
-            .unwrap();
+        index.remap(&remap, test_store.as_ref()).await.unwrap();
 
         // Reload and check
         let reloaded_idx = BitmapIndex::load(test_store, None, &LanceCache::no_cache())

rust/lance-index/src/scalar/bitmap.rs

@@ -20,12 +20,13 @@ use arrow_schema::{DataType, Field};
 use lance_arrow_stats::StatisticsAccumulator;
 use lance_core::utils::bloomfilter::as_bytes;
 use lance_core::utils::bloomfilter::sbbf::{Sbbf, SbbfBuilder};
+use lance_core::utils::row_addr_remap::RowAddrRemap;
 use serde::{Deserialize, Serialize};
 
 use std::sync::LazyLock;
 
 use datafusion::execution::SendableRecordBatchStream;
-use std::{collections::HashMap, sync::Arc};
+use std::sync::Arc;
 
 use crate::scalar::{
     AnyQuery, IndexStore, MetricsCollector, RowIdRemapper, ScalarIndex, SearchResult,
@@ -425,7 +426,7 @@ impl ScalarIndex for BloomFilterIndex {
 
     async fn remap(
         &self,
-        _mapping: &HashMap<u64, Option<u64>>,
+        _mapping: &RowAddrRemap,
         _dest_store: &dyn IndexStore,
     ) -> Result<CreatedIndex> {
         Err(Error::invalid_input_source(

rust/lance-index/src/scalar/bloomfilter.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The Lance Authors
 
+use lance_core::utils::row_addr_remap::RowAddrRemap;
 use std::{
     any::Any,
     cmp::Ordering,
@@ -2227,7 +2228,7 @@ impl ScalarIndex for BTreeIndex {
 
     async fn remap(
         &self,
-        mapping: &HashMap<u64, Option<u64>>,
+        mapping: &RowAddrRemap,
         dest_store: &dyn IndexStore,
     ) -> Result<CreatedIndex> {
         // (part_id, path)
@@ -2394,7 +2395,7 @@ pub trait BTreeSubIndex: Debug + Send + Sync + DeepSizeOf {
     async fn remap_subindex(
         &self,
         serialized: RecordBatch,
-        mapping: &HashMap<u64, Option<u64>>,
+        mapping: &RowAddrRemap,
     ) -> Result<RecordBatch>;
 }
 
@@ -3355,6 +3356,7 @@ impl ScalarIndexPlugin for BTreeIndexPlugin {
 
 #[cfg(test)]
 mod tests {
+    use lance_core::utils::row_addr_remap::RowAddrRemap;
     use std::sync::atomic::Ordering;
     use std::{collections::HashMap, sync::Arc};
 
@@ -3503,7 +3505,7 @@ mod tests {
 
         // Remap with a no-op mapping.  The remapped index should be identical to the original
         index
-            .remap(&HashMap::default(), remap_store.as_ref())
+            .remap(&RowAddrRemap::empty(), remap_store.as_ref())
             .await
             .unwrap();
 
@@ -5014,7 +5016,7 @@ mod tests {
 
         // Remap with a no-op mapping.  The remapped index should be identical to the original
         ranged_index
-            .remap(&HashMap::default(), remap_store.as_ref())
+            .remap(&RowAddrRemap::empty(), remap_store.as_ref())
             .await
             .unwrap();
 
@@ -5342,7 +5344,10 @@ mod tests {
         ));
 
         // Remap the index with our deletion mapping
-        index.remap(&mapping, remap_store.as_ref()).await.unwrap();
+        index
+            .remap(&RowAddrRemap::direct(mapping), remap_store.as_ref())
+            .await
+            .unwrap();
 
         let remapped_index = BTreeIndex::load(remap_store.clone(), None, &LanceCache::no_cache())
             .await

rust/lance-index/src/scalar/btree.rs

@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The Lance Authors
 
-use std::collections::{BTreeSet, HashMap};
+use lance_core::utils::row_addr_remap::RowAddrRemap;
+use std::collections::BTreeSet;
 use std::{ops::Bound, sync::Arc};
 
 use arrow_array::Array;
@@ -133,19 +134,15 @@ impl FlatIndex {
         NullableRowAddrSet::new(self.all_addrs_map.clone(), Default::default())
     }
 
-    pub fn remap_batch(
-        batch: RecordBatch,
-        mapping: &HashMap<u64, Option<u64>>,
-    ) -> Result<RecordBatch> {
+    pub fn remap_batch(batch: RecordBatch, mapping: &RowAddrRemap) -> Result<RecordBatch> {
         let row_ids = batch.column(IDS_COL_IDX).as_primitive::<UInt64Type>();
         let val_idx_and_new_id = row_ids
             .values()
             .iter()
             .enumerate()
             .filter_map(|(idx, old_id)| {
                 mapping
-                    .get(old_id)
-                    .copied()
+                    .get(*old_id)
                     .unwrap_or(Some(*old_id))
                     .map(|new_id| (idx, new_id))
             })
@@ -346,7 +343,11 @@ mod tests {
     use super::*;
     use arrow_array::{record_batch, types::Int32Type};
     use datafusion_common::ScalarValue;
+    use lance_core::utils::row_addr_remap::GroupInput;
     use lance_datagen::{RowCount, array, gen_batch};
+    use roaring::RoaringTreemap;
+    use rstest::rstest;
+    use std::collections::HashMap;
 
     fn example_index() -> FlatIndex {
         let batch = gen_batch()
@@ -490,9 +491,10 @@ mod tests {
         // 3 -> delete
         // Keep remaining as is
         let mapping = HashMap::<u64, Option<u64>>::from_iter(vec![(0, Some(2000)), (3, None)]);
-        let remapped =
-            FlatIndex::try_new(FlatIndex::remap_batch((*index.data).clone(), &mapping).unwrap())
-                .unwrap();
+        let remapped = FlatIndex::try_new(
+            FlatIndex::remap_batch((*index.data).clone(), &RowAddrRemap::direct(mapping)).unwrap(),
+        )
+        .unwrap();
 
         let expected = FlatIndex::try_new(
             gen_batch()
@@ -505,18 +507,22 @@ mod tests {
         assert_eq!(remapped.data, expected.data);
     }
 
-    // It's possible, during compaction, that an entire page of values is deleted.  We just serialize
-    // it as an empty record batch.
-    #[tokio::test]
-    async fn test_remap_to_nothing() {
+    // An entire page (frag 0) is deleted during compaction. remap_batch must
+    // drop every row regardless of which RowAddrRemap mode expresses it.
+    // example_index holds row ids 5, 0, 3, 100, all in frag 0.
+    #[rstest]
+    #[case::compact(RowAddrRemap::compact([GroupInput {
+        rewritten_old_row_addrs: RoaringTreemap::new(),
+        old_frag_ids: vec![0],
+        new_frags: vec![],
+    }])
+    .unwrap())]
+    #[case::explicit(RowAddrRemap::direct(
+        [5u64, 0, 3, 100].into_iter().map(|id| (id, None)).collect(),
+    ))]
+    fn test_remap_to_nothing(#[case] remap: RowAddrRemap) {
         let index = example_index();
-        let mapping = HashMap::<u64, Option<u64>>::from_iter(vec![
-            (5, None),
-            (0, None),
-            (3, None),
-            (100, None),
-        ]);
-        let remapped = FlatIndex::remap_batch((*index.data).clone(), &mapping).unwrap();
+        let remapped = FlatIndex::remap_batch((*index.data).clone(), &remap).unwrap();
         assert_eq!(remapped.num_rows(), 0);
     }
 

rust/lance-index/src/scalar/btree/flat.rs

@@ -32,6 +32,7 @@ use datafusion::execution::SendableRecordBatchStream;
 use futures::{StreamExt, TryStreamExt};
 use lance_core::cache::LanceCache;
 use lance_core::deepsize::DeepSizeOf;
+use lance_core::utils::row_addr_remap::RowAddrRemap;
 use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu};
 use lance_core::{Error, ROW_ADDR, Result};
 use roaring::RoaringBitmap;
@@ -1828,11 +1829,7 @@ impl ScalarIndex for FMIndexScalarIndex {
     fn can_remap(&self) -> bool {
         false
     }
-    async fn remap(
-        &self,
-        _: &HashMap<u64, Option<u64>>,
-        _: &dyn IndexStore,
-    ) -> Result<CreatedIndex> {
+    async fn remap(&self, _: &RowAddrRemap, _: &dyn IndexStore) -> Result<CreatedIndex> {
         Err(Error::not_supported("Fm does not support remap"))
     }
     async fn update(