ailake_query/
dv.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// Iceberg V3 Deletion Vector support — Phase B (read only).
3//
4// A Deletion Vector (DV) is a Roaring Bitmap stored inside a Puffin `.dvd` file.
5// The manifest entry carries `(path, offset, length)` that address the bitmap blob
6// bytes directly, so no full Puffin footer parse is needed for read support.
7//
8// Phase C will add write support (producing DVs on row delete operations).
9
10use ailake_catalog::provider::DeletionVector;
11use ailake_core::{AilakeError, AilakeResult};
12use ailake_store::Store;
13use arrow_array::RecordBatch;
14use roaring::RoaringBitmap;
15use std::sync::Arc;
16
17/// Fetch and deserialize a Deletion Vector bitmap from a Puffin `.dvd` file.
18///
19/// Uses a range GET (`offset..offset+length`) so only the bitmap bytes are
20/// transferred from S3 — no full file download required.
21///
22/// The returned bitmap contains the row positions (0-based within the data file)
23/// that have been deleted. Callers must filter HNSW results against this bitmap:
24/// ```ignore
25/// results.retain(|(row_id, _)| !bitmap.contains(row_id.as_u64() as u32));
26/// ```
27pub async fn load_deletion_vector(
28    store: &Arc<dyn Store>,
29    dv: &DeletionVector,
30) -> AilakeResult<RoaringBitmap> {
31    let bytes = store
32        .get_range(&dv.path, dv.offset..dv.offset + dv.length)
33        .await?;
34
35    RoaringBitmap::deserialize_from(bytes.as_ref()).map_err(|e| {
36        AilakeError::Io(std::io::Error::other(format!(
37            "ailake: failed to deserialize Deletion Vector bitmap from '{}' \
38             (offset={}, length={}): {e}",
39            dv.path, dv.offset, dv.length
40        )))
41    })
42}
43
44/// Removes DV-masked rows from a Parquet-read `(batch, parallel)` pair, where
45/// `parallel[i]` corresponds positionally to `batch` row `i` (embeddings, texts, ...).
46///
47/// Used by rewrite jobs (compaction, migration, backfill) that read a file's raw rows
48/// and write them into a brand-new physical file: since the new file has fresh row
49/// positions, a deleted row can't just be re-masked with the old bitmap after the fact —
50/// it must be dropped from the input before the merge, or it silently reappears (no DV
51/// on the new `DataFileEntry`, or a DV bitmap that no longer lines up with the new
52/// row order).
53pub fn filter_deleted_rows<T>(
54    batch: RecordBatch,
55    parallel: Vec<T>,
56    bitmap: &RoaringBitmap,
57) -> AilakeResult<(RecordBatch, Vec<T>)> {
58    if bitmap.is_empty() {
59        return Ok((batch, parallel));
60    }
61    let n = batch.num_rows();
62    let keep: Vec<bool> = (0..n).map(|i| !bitmap.contains(i as u32)).collect();
63    let filtered_parallel: Vec<T> = parallel
64        .into_iter()
65        .zip(keep.iter())
66        .filter_map(|(v, &k)| k.then_some(v))
67        .collect();
68    let mask = arrow_array::BooleanArray::from(keep);
69    let filtered_batch = arrow_select::filter::filter_record_batch(&batch, &mask)
70        .map_err(|e| AilakeError::Arrow(e.to_string()))?;
71    Ok((filtered_batch, filtered_parallel))
72}
73
74/// Returns true when any row in `row_ids` is deleted according to `bitmap`.
75/// Used for early-exit pruning: if zero deletions touch the candidate set, skip
76/// the per-row check.
77#[inline]
78pub fn has_deletions(bitmap: &RoaringBitmap, row_ids: &[u64]) -> bool {
79    row_ids.iter().any(|&id| bitmap.contains(id as u32))
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use ailake_store::LocalStore;
86    use bytes::Bytes;
87    use roaring::RoaringBitmap;
88
89    fn make_bitmap_bytes(deleted: &[u32]) -> Vec<u8> {
90        let mut bm = RoaringBitmap::new();
91        for &r in deleted {
92            bm.insert(r);
93        }
94        let mut buf = Vec::new();
95        bm.serialize_into(&mut buf).unwrap();
96        buf
97    }
98
99    #[tokio::test]
100    async fn load_dv_roundtrip() {
101        let dir = tempfile::tempdir().unwrap();
102        let bitmap_bytes = make_bitmap_bytes(&[0, 5, 42, 1000]);
103
104        // Write a minimal Puffin-like file: just the bitmap bytes at a known offset.
105        // In real Puffin files there is a header and footer; the DV manifest entry
106        // gives us the exact offset+length so we only fetch those bytes.
107        let offset: u64 = 16; // simulate a 16-byte Puffin header before the blob
108        let mut file_bytes = vec![0u8; offset as usize]; // fake puffin header
109        file_bytes.extend_from_slice(&bitmap_bytes);
110
111        let dvd_path = "data/dv-0001.dvd";
112        let store: Arc<dyn Store> = Arc::new(LocalStore::new(dir.path()));
113        store.put(dvd_path, Bytes::from(file_bytes)).await.unwrap();
114
115        let dv = DeletionVector {
116            path: dvd_path.to_string(),
117            offset,
118            length: bitmap_bytes.len() as u64,
119            cardinality: 4,
120        };
121
122        let bm = load_deletion_vector(&store, &dv).await.unwrap();
123        assert!(bm.contains(0));
124        assert!(bm.contains(5));
125        assert!(bm.contains(42));
126        assert!(bm.contains(1000));
127        assert!(!bm.contains(1)); // not deleted
128        assert_eq!(bm.len(), 4);
129    }
130
131    #[test]
132    fn has_deletions_detects_overlap() {
133        let mut bm = RoaringBitmap::new();
134        bm.insert(10);
135        bm.insert(20);
136
137        assert!(has_deletions(&bm, &[5, 10, 15])); // 10 is deleted
138        assert!(!has_deletions(&bm, &[1, 2, 3])); // none deleted
139    }
140}
ailake_query/dv.rs

ailake_query/
dv.rs
