ailake_query/
equality_delete.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Equality delete filter — Phase H.
3//!
4//! Loads Iceberg equality delete files from the object store and builds an in-memory
5//! predicate set. Applied to each `RecordBatch` during scan to mask logically deleted rows.
6//!
7//! Scope: single-column equality predicates (most common pattern: document_id, agent_id,
8//! session_id). Multi-column AND predicates are supported as long as each column is checked
9//! independently (conservative: a row is deleted if ALL delete-file columns match).
10
11use std::collections::{HashMap, HashSet};
12use std::sync::Arc;
13
14use ailake_catalog::{read_equality_delete_values, EqualityDeleteFile};
15use ailake_core::{AilakeError, AilakeResult};
16use ailake_store::Store;
17use arrow_array::{
18    Array, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch,
19    StringArray,
20};
21use arrow_schema::DataType;
22
23/// In-memory equality delete filter built from one or more delete files.
24///
25/// Each entry is `column_name → set of string-normalised values to delete`.
26/// A data row is deleted if, for every column in the filter, the row's value
27/// is a member of that column's set.
28pub struct EqualityDeleteFilter {
29    /// column_name → set of values (string-normalised) to delete
30    filters: HashMap<String, HashSet<String>>,
31}
32
33impl EqualityDeleteFilter {
34    /// Build filter from a list of equality delete file references.
35    ///
36    /// For each file, downloads the Avro payload from `store` and extracts
37    /// `(column_name, value)` pairs. All files are merged into one filter.
38    pub async fn from_files(
39        store: &Arc<dyn Store>,
40        files: &[EqualityDeleteFile],
41    ) -> AilakeResult<Self> {
42        let mut filters: HashMap<String, HashSet<String>> = HashMap::new();
43        for edf in files {
44            let bytes = store.get(&edf.path).await?;
45            let pairs = read_equality_delete_values(&bytes)
46                .map_err(|e| AilakeError::Catalog(e.to_string()))?;
47            for (col, val) in pairs {
48                filters.entry(col).or_default().insert(val);
49            }
50        }
51        Ok(Self { filters })
52    }
53
54    pub fn empty() -> Self {
55        Self {
56            filters: HashMap::new(),
57        }
58    }
59
60    pub fn is_empty(&self) -> bool {
61        self.filters.is_empty()
62    }
63
64    /// Check whether a single row (by its physical index in `batch`) matches the delete predicate.
65    ///
66    /// Returns `true` if the row should be logically deleted.
67    /// Iceberg equality delete semantics: a row is deleted only when ALL columns in the
68    /// delete predicate match (AND, not OR). Used in the per-row HNSW result loop.
69    pub fn should_delete_row(&self, batch: &RecordBatch, row_idx: usize) -> bool {
70        if self.filters.is_empty() {
71            return false;
72        }
73        let mut any_column_found = false;
74        for (col_name, delete_values) in &self.filters {
75            let col_idx = match batch.schema().index_of(col_name.as_str()) {
76                Ok(i) => i,
77                Err(_) => continue, // column absent — skip (schema evolution)
78            };
79            any_column_found = true;
80            let array = batch.column(col_idx);
81            if array.is_null(row_idx) {
82                return false; // null never matches — AND tuple fails
83            }
84            let val_str: Option<String> = match array.data_type() {
85                DataType::Utf8 => array
86                    .as_any()
87                    .downcast_ref::<StringArray>()
88                    .map(|a| a.value(row_idx).to_string()),
89                DataType::LargeUtf8 => array
90                    .as_any()
91                    .downcast_ref::<arrow_array::LargeStringArray>()
92                    .map(|a| a.value(row_idx).to_string()),
93                DataType::Int32 => array
94                    .as_any()
95                    .downcast_ref::<Int32Array>()
96                    .map(|a| a.value(row_idx).to_string()),
97                DataType::Int64 => array
98                    .as_any()
99                    .downcast_ref::<Int64Array>()
100                    .map(|a| a.value(row_idx).to_string()),
101                DataType::Float32 => array
102                    .as_any()
103                    .downcast_ref::<Float32Array>()
104                    .map(|a| a.value(row_idx).to_string()),
105                DataType::Float64 => array
106                    .as_any()
107                    .downcast_ref::<Float64Array>()
108                    .map(|a| a.value(row_idx).to_string()),
109                _ => None,
110            };
111            match val_str {
112                Some(s) if delete_values.contains(&s) => {} // column matches, continue AND check
113                Some(_) => return false,                    // column mismatch — AND tuple fails
114                None => {}                                  // unknown type — skip column
115            }
116        }
117        any_column_found // true only when all checked columns matched
118    }
119
120    /// Apply the filter to `batch`, returning a new batch with matching rows removed.
121    ///
122    /// Iceberg equality delete semantics: a row is removed only when ALL columns in the
123    /// delete predicate match (AND). Columns absent from the batch are ignored.
124    pub fn apply(&self, batch: RecordBatch) -> AilakeResult<RecordBatch> {
125        if self.filters.is_empty() {
126            return Ok(batch);
127        }
128        let n = batch.num_rows();
129        let keep: Vec<bool> = (0..n).map(|i| !self.should_delete_row(&batch, i)).collect();
130        let mask = BooleanArray::from(keep);
131        arrow_select::filter::filter_record_batch(&batch, &mask)
132            .map_err(|e| AilakeError::Arrow(e.to_string()))
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use std::sync::Arc;
139
140    use arrow_array::{Int32Array, RecordBatch, StringArray};
141    use arrow_schema::{DataType, Field, Schema};
142
143    use super::EqualityDeleteFilter;
144    use std::collections::{HashMap, HashSet};
145
146    fn make_batch() -> RecordBatch {
147        let schema = Arc::new(Schema::new(vec![
148            Field::new("doc_id", DataType::Utf8, true),
149            Field::new("score", DataType::Int32, true),
150        ]));
151        RecordBatch::try_new(
152            schema,
153            vec![
154                Arc::new(StringArray::from(vec!["doc-a", "doc-b", "doc-c", "doc-d"])),
155                Arc::new(Int32Array::from(vec![1, 2, 3, 4])),
156            ],
157        )
158        .unwrap()
159    }
160
161    fn filter_with(filters: HashMap<String, HashSet<String>>) -> EqualityDeleteFilter {
162        EqualityDeleteFilter { filters }
163    }
164
165    #[test]
166    fn empty_filter_is_no_op() {
167        let batch = make_batch();
168        let f = filter_with(HashMap::new());
169        let result = f.apply(batch.clone()).unwrap();
170        assert_eq!(result.num_rows(), 4);
171    }
172
173    #[test]
174    fn single_value_deleted() {
175        let mut filters = HashMap::new();
176        filters.insert("doc_id".into(), ["doc-b".to_string()].into());
177        let f = filter_with(filters);
178        let result = f.apply(make_batch()).unwrap();
179        assert_eq!(result.num_rows(), 3);
180        let ids = result
181            .column(0)
182            .as_any()
183            .downcast_ref::<StringArray>()
184            .unwrap();
185        assert_eq!(ids.value(0), "doc-a");
186        assert_eq!(ids.value(1), "doc-c");
187        assert_eq!(ids.value(2), "doc-d");
188    }
189
190    #[test]
191    fn multiple_values_deleted() {
192        let mut filters = HashMap::new();
193        filters.insert(
194            "doc_id".into(),
195            ["doc-a".to_string(), "doc-c".to_string()].into(),
196        );
197        let f = filter_with(filters);
198        let result = f.apply(make_batch()).unwrap();
199        assert_eq!(result.num_rows(), 2);
200        let ids = result
201            .column(0)
202            .as_any()
203            .downcast_ref::<StringArray>()
204            .unwrap();
205        assert_eq!(ids.value(0), "doc-b");
206        assert_eq!(ids.value(1), "doc-d");
207    }
208
209    #[test]
210    fn column_absent_from_batch_is_skipped() {
211        let mut filters = HashMap::new();
212        filters.insert("nonexistent_col".into(), ["x".to_string()].into());
213        let f = filter_with(filters);
214        let result = f.apply(make_batch()).unwrap();
215        assert_eq!(result.num_rows(), 4); // no rows deleted
216    }
217
218    #[test]
219    fn numeric_column_deletion() {
220        let mut filters = HashMap::new();
221        filters.insert("score".into(), ["2".to_string(), "4".to_string()].into());
222        let f = filter_with(filters);
223        let result = f.apply(make_batch()).unwrap();
224        assert_eq!(result.num_rows(), 2);
225        let ids = result
226            .column(0)
227            .as_any()
228            .downcast_ref::<StringArray>()
229            .unwrap();
230        assert_eq!(ids.value(0), "doc-a");
231        assert_eq!(ids.value(1), "doc-c");
232    }
233}
ailake_query/equality_delete.rs

ailake_query/
equality_delete.rs
