lance_core/utils/
fragment_reuse.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Readers for a single fragment rewrite. Graph traversal belongs to the FRI reader.
5
6use std::sync::Arc;
7
8use async_trait::async_trait;
9
10use super::row_addr_remap::RowAddrRemap;
11use crate::deepsize::{Context, DeepSizeOf};
12use crate::{Error, Result};
13
14/// The read contract for one mapping, independent of index types and manifests.
15///
16/// Row IDs are physical addresses. Results preserve input positions and `None`
17/// denotes a deleted source row. Addresses outside the source fragments, including
18/// out-of-range offsets, return [`Error::InvalidInput`]. Invalid persisted mapping
19/// contents return [`Error::CorruptFile`]; storage errors propagate to the caller.
20#[async_trait]
21pub trait MappingReader: Send + Sync + std::fmt::Debug + DeepSizeOf {
22    /// Remap one physical row ID, or return `None` if the row was deleted.
23    async fn remap_row_id(&self, row_id: u64) -> Result<Option<u64>>;
24
25    /// Remap a batch, preserving input order and duplicates, with one result per input.
26    async fn remap_row_ids(&self, row_ids: &[u64]) -> Result<Vec<Option<u64>>> {
27        let mut mapped = Vec::with_capacity(row_ids.len());
28        for &row_id in row_ids {
29            mapped.push(self.remap_row_id(row_id).await?);
30        }
31        Ok(mapped)
32    }
33}
34
35/// Ordered-compaction adapter for one rewrite group's existing bitmap/rank remap.
36/// This adapter does not replace the legacy FRI reader or writer.
37#[derive(Debug)]
38pub struct OrderedCompactionMapping {
39    remap: Arc<RowAddrRemap>,
40}
41
42impl OrderedCompactionMapping {
43    /// Use a remap whose fragment layout was validated when it was constructed.
44    pub fn new(remap: Arc<RowAddrRemap>) -> Self {
45        Self { remap }
46    }
47
48    fn map_row_id(&self, row_id: u64) -> Result<Option<u64>> {
49        self.remap.get(row_id).ok_or_else(|| {
50            Error::invalid_input(format!("row ID {row_id} is outside compaction mapping"))
51        })
52    }
53}
54
55impl DeepSizeOf for OrderedCompactionMapping {
56    fn deep_size_of_children(&self, context: &mut Context) -> usize {
57        self.remap.deep_size_of_children(context)
58    }
59}
60
61#[async_trait]
62impl MappingReader for OrderedCompactionMapping {
63    async fn remap_row_id(&self, row_id: u64) -> Result<Option<u64>> {
64        self.map_row_id(row_id)
65    }
66
67    async fn remap_row_ids(&self, row_ids: &[u64]) -> Result<Vec<Option<u64>>> {
68        row_ids
69            .iter()
70            .map(|&row_id| self.map_row_id(row_id))
71            .collect()
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78    use crate::utils::address::RowAddress;
79    use crate::utils::row_addr_remap::GroupInputWithLayout;
80    use roaring::RoaringTreemap;
81
82    #[derive(Debug, DeepSizeOf)]
83    struct SingleRowReader;
84
85    #[async_trait]
86    impl MappingReader for SingleRowReader {
87        async fn remap_row_id(&self, row_id: u64) -> Result<Option<u64>> {
88            if row_id == u64::MAX {
89                return Err(Error::invalid_input("invalid test row ID"));
90            }
91            Ok((row_id != 0).then_some(row_id))
92        }
93    }
94
95    #[tokio::test]
96    async fn default_batch_preserves_positions_and_propagates_errors() {
97        let reader = SingleRowReader;
98        assert_eq!(
99            reader.remap_row_ids(&[2, 0, 1, 2]).await.unwrap(),
100            vec![Some(2), None, Some(1), Some(2)],
101        );
102        assert!(reader.remap_row_ids(&[]).await.unwrap().is_empty());
103        let error = reader.remap_row_ids(&[1, u64::MAX]).await.unwrap_err();
104        assert!(matches!(error, Error::InvalidInput { .. }));
105        assert!(error.to_string().contains("invalid test row ID"));
106    }
107
108    #[tokio::test]
109    async fn compaction_reader_preserves_legacy_mapping_semantics() {
110        let address = |fragment, offset| u64::from(RowAddress::new_from_parts(fragment, offset));
111        let remap = Arc::new(
112            RowAddrRemap::compact_with_layout([GroupInputWithLayout {
113                rewritten_old_row_addrs: RoaringTreemap::from_iter([
114                    address(1, 0),
115                    address(1, 2),
116                    address(2, 0),
117                ]),
118                old_frags: vec![(1, 3), (2, 1)],
119                new_frags: vec![(3, 2), (4, 1)],
120            }])
121            .unwrap(),
122        );
123        let reader = OrderedCompactionMapping::new(remap.clone());
124        let input = [address(2, 0), address(1, 1), address(1, 0), address(1, 0)];
125        let expected: Vec<_> = input.iter().map(|&a| remap.get(a).unwrap()).collect();
126        assert_eq!(reader.remap_row_ids(&input).await.unwrap(), expected);
127        for (&row_id, &mapped) in input.iter().zip(&expected) {
128            assert_eq!(reader.remap_row_id(row_id).await.unwrap(), mapped);
129        }
130        assert!(reader.remap_row_ids(&[]).await.unwrap().is_empty());
131        for invalid in [address(99, 0), address(1, 3)] {
132            for error in [
133                reader.remap_row_id(invalid).await.unwrap_err(),
134                reader
135                    .remap_row_ids(&[input[0], invalid])
136                    .await
137                    .unwrap_err(),
138            ] {
139                assert!(matches!(error, Error::InvalidInput { .. }));
140                assert!(error.to_string().contains("outside compaction mapping"));
141            }
142        }
143    }
144}
lance_core/utils/fragment_reuse.rs

lance_core/utils/
fragment_reuse.rs
