laminar_connectors/lakehouse/
iceberg_lookup.rs

1//! Iceberg on-demand lookup source for cache-miss fallback.
2//!
3//! Implements `LookupSource` via the native Iceberg scan with a batched
4//! `pk IN (...)` filter pushed down (`with_filter`), so all missed keys of a
5//! probe fold into one manifest-pruned scan. [`KeyAligner`] handles key decode
6//! and result realignment.
7
8use std::sync::Arc;
9
10use arrow_array::{Array, ArrayRef, RecordBatch};
11use arrow_row::SortField;
12use arrow_schema::SchemaRef;
13use iceberg::expr::{Predicate as IcebergPredicate, Reference};
14use iceberg::spec::Datum;
15use iceberg::Catalog;
16
17use laminar_core::lookup::predicate::Predicate;
18use laminar_core::lookup::source::{
19    projection_names, ColumnId, LookupError, LookupSource, LookupSourceCapabilities,
20};
21use laminar_core::lookup::KeyAligner;
22
23use crate::lakehouse::iceberg_config::IcebergCatalogConfig;
24
25/// Configuration for [`IcebergLookupSource`].
26#[derive(Debug, Clone)]
27pub struct IcebergLookupSourceConfig {
28    /// Shared catalog connection settings (also carries namespace + table).
29    pub catalog: IcebergCatalogConfig,
30    /// Primary key column names.
31    pub primary_key_columns: Vec<String>,
32}
33
34/// Iceberg lookup source for on-demand/partial cache mode.
35pub struct IcebergLookupSource {
36    catalog: Arc<dyn Catalog>,
37    namespace: String,
38    table_name: String,
39    schema: SchemaRef,
40    aligner: KeyAligner,
41}
42
43impl IcebergLookupSource {
44    /// Opens the catalog, loads the table, and derives the Arrow schema.
45    ///
46    /// # Errors
47    ///
48    /// Returns `LookupError` if the catalog/table cannot be opened or a primary
49    /// key column is missing from the table schema.
50    pub async fn open(config: IcebergLookupSourceConfig) -> Result<Self, LookupError> {
51        let catalog = crate::lakehouse::iceberg_io::build_catalog(&config.catalog)
52            .await
53            .map_err(|e| LookupError::Connection(format!("iceberg catalog: {e}")))?;
54        let table = crate::lakehouse::iceberg_io::load_table(
55            catalog.as_ref(),
56            &config.catalog.namespace,
57            &config.catalog.table_name,
58        )
59        .await
60        .map_err(|e| LookupError::Connection(format!("load iceberg table: {e}")))?;
61
62        let iceberg_schema = table.current_schema_ref();
63        let schema: SchemaRef = Arc::new(
64            iceberg::arrow::schema_to_arrow_schema(&iceberg_schema)
65                .map_err(|e| LookupError::Internal(format!("iceberg schema to arrow: {e}")))?,
66        );
67
68        let pk_sort_fields = config
69            .primary_key_columns
70            .iter()
71            .map(|name| {
72                let idx = schema
73                    .index_of(name)
74                    .map_err(|_| LookupError::Internal(format!("pk column not found: {name}")))?;
75                Ok(SortField::new(schema.field(idx).data_type().clone()))
76            })
77            .collect::<Result<Vec<_>, LookupError>>()?;
78        let aligner = KeyAligner::new(pk_sort_fields, config.primary_key_columns)?;
79
80        Ok(Self {
81            catalog,
82            namespace: config.catalog.namespace,
83            table_name: config.catalog.table_name,
84            schema,
85            aligner,
86        })
87    }
88
89    /// Convert one PK cell into an Iceberg [`Datum`], or `None` when NULL.
90    fn cell_to_datum(
91        col_name: &str,
92        array: &dyn Array,
93        row: usize,
94    ) -> Result<Option<Datum>, LookupError> {
95        use arrow_array::{
96            BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array,
97            Int8Array, LargeStringArray, StringArray, StringViewArray,
98        };
99        use arrow_schema::DataType;
100
101        if array.is_null(row) {
102            return Ok(None);
103        }
104
105        macro_rules! downcast {
106            ($ty:ty) => {
107                array.as_any().downcast_ref::<$ty>().ok_or_else(|| {
108                    LookupError::Internal(format!("pk column '{col_name}' downcast failed"))
109                })?
110            };
111        }
112
113        let datum = match array.data_type() {
114            DataType::Int8 => Datum::int(i32::from(downcast!(Int8Array).value(row))),
115            DataType::Int16 => Datum::int(i32::from(downcast!(Int16Array).value(row))),
116            DataType::Int32 => Datum::int(downcast!(Int32Array).value(row)),
117            DataType::Int64 => Datum::long(downcast!(Int64Array).value(row)),
118            DataType::Float32 => Datum::float(downcast!(Float32Array).value(row)),
119            DataType::Float64 => Datum::double(downcast!(Float64Array).value(row)),
120            DataType::Boolean => Datum::bool(downcast!(BooleanArray).value(row)),
121            DataType::Utf8 => Datum::string(downcast!(StringArray).value(row)),
122            DataType::LargeUtf8 => Datum::string(downcast!(LargeStringArray).value(row)),
123            DataType::Utf8View => Datum::string(downcast!(StringViewArray).value(row)),
124            dt => {
125                return Err(LookupError::Internal(format!(
126                    "unsupported PK data type for iceberg lookup: {dt} (column \"{col_name}\")"
127                )));
128            }
129        };
130        Ok(Some(datum))
131    }
132
133    /// Build a single Iceberg predicate over the decoded key columns: a
134    /// single-column PK folds into `pk IN (...)`; a composite PK becomes an OR
135    /// of per-key AND-groups.
136    fn build_key_predicate(
137        pk_cols: &[String],
138        pk_arrays: &[ArrayRef],
139        n_keys: usize,
140    ) -> Result<IcebergPredicate, LookupError> {
141        if pk_cols.len() == 1 {
142            let col = &pk_cols[0];
143            let array = pk_arrays[0].as_ref();
144            let mut datums = Vec::with_capacity(n_keys);
145            let mut has_null = false;
146            for row in 0..n_keys {
147                match Self::cell_to_datum(col, array, row)? {
148                    Some(d) => datums.push(d),
149                    None => has_null = true,
150                }
151            }
152            let mut pred: Option<IcebergPredicate> =
153                (!datums.is_empty()).then(|| Reference::new(col.clone()).is_in(datums));
154            if has_null {
155                let null_pred = Reference::new(col.clone()).is_null();
156                pred = Some(match pred {
157                    Some(p) => p.or(null_pred),
158                    None => null_pred,
159                });
160            }
161            return pred.ok_or_else(|| LookupError::Internal("no keys to look up".into()));
162        }
163
164        let mut groups: Vec<IcebergPredicate> = Vec::with_capacity(n_keys);
165        for row in 0..n_keys {
166            let mut conj: Option<IcebergPredicate> = None;
167            for (ci, col) in pk_cols.iter().enumerate() {
168                let term = match Self::cell_to_datum(col, pk_arrays[ci].as_ref(), row)? {
169                    Some(d) => Reference::new(col.clone()).equal_to(d),
170                    None => Reference::new(col.clone()).is_null(),
171                };
172                conj = Some(match conj {
173                    Some(c) => c.and(term),
174                    None => term,
175                });
176            }
177            if let Some(c) = conj {
178                groups.push(c);
179            }
180        }
181        let mut it = groups.into_iter();
182        it.next()
183            .map(|first| it.fold(first, IcebergPredicate::or))
184            .ok_or_else(|| LookupError::Internal("no keys to look up".into()))
185    }
186}
187
188impl LookupSource for IcebergLookupSource {
189    async fn query(
190        &self,
191        keys: &[&[u8]],
192        _predicates: &[Predicate],
193        projection: &[ColumnId],
194    ) -> Result<Vec<Option<RecordBatch>>, LookupError> {
195        use tokio_stream::StreamExt;
196
197        if keys.is_empty() {
198            return Ok(Vec::new());
199        }
200
201        let pk_arrays = self.aligner.decode_keys(keys)?;
202        let predicate =
203            Self::build_key_predicate(self.aligner.pk_columns(), &pk_arrays, keys.len())?;
204
205        // Reload the table per query so lookups see the latest snapshot
206        // (eventually-consistent freshness; the cache layer adds TTL on top).
207        let table = crate::lakehouse::iceberg_io::load_table(
208            self.catalog.as_ref(),
209            &self.namespace,
210            &self.table_name,
211        )
212        .await
213        .map_err(|e| LookupError::Query(format!("load iceberg table: {e}")))?;
214
215        // Projection pushdown: select only the requested columns (always incl.
216        // the key, so realignment still works), else every column.
217        let mut builder = table.scan().with_filter(predicate);
218        builder = if projection.is_empty() {
219            builder.select_all()
220        } else {
221            builder.select(projection_names(&self.schema, projection)?)
222        };
223        let scan = builder
224            .build()
225            .map_err(|e| LookupError::Query(format!("build iceberg scan: {e}")))?;
226        let stream = scan
227            .to_arrow()
228            .await
229            .map_err(|e| LookupError::Query(format!("iceberg scan to arrow: {e}")))?;
230
231        let mut batches = Vec::new();
232        let mut stream = std::pin::pin!(stream);
233        while let Some(result) = stream.next().await {
234            batches
235                .push(result.map_err(|e| LookupError::Query(format!("read iceberg batch: {e}")))?);
236        }
237
238        self.aligner.align(keys, &batches)
239    }
240
241    fn capabilities(&self) -> LookupSourceCapabilities {
242        LookupSourceCapabilities {
243            supports_batch_lookup: true,
244            supports_projection_pushdown: true,
245            ..LookupSourceCapabilities::none()
246        }
247    }
248
249    fn source_name(&self) -> &'static str {
250        "iceberg"
251    }
252
253    fn schema(&self) -> SchemaRef {
254        Arc::clone(&self.schema)
255    }
256
257    async fn health_check(&self) -> Result<(), LookupError> {
258        crate::lakehouse::iceberg_io::load_table(
259            self.catalog.as_ref(),
260            &self.namespace,
261            &self.table_name,
262        )
263        .await
264        .map(|_| ())
265        .map_err(|e| LookupError::Connection(format!("health check: {e}")))
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272    use arrow_array::{Int64Array, StringArray};
273
274    #[test]
275    fn cell_to_datum_null_and_unsupported() {
276        let arr = Int64Array::from(vec![None, Some(1)]);
277        assert!(IcebergLookupSource::cell_to_datum("id", &arr, 0)
278            .unwrap()
279            .is_none());
280        assert!(IcebergLookupSource::cell_to_datum("id", &arr, 1)
281            .unwrap()
282            .is_some());
283        // Binary keys are not supported as Iceberg predicates.
284        let bin = arrow_array::BinaryArray::from(vec![b"x".as_ref()]);
285        assert!(IcebergLookupSource::cell_to_datum("k", &bin, 0).is_err());
286    }
287
288    #[test]
289    fn single_col_predicate_is_in_list() {
290        let cols = vec!["id".to_string()];
291        let arrays: Vec<ArrayRef> = vec![Arc::new(Int64Array::from(vec![1, 2, 3]))];
292        let s = format!(
293            "{}",
294            IcebergLookupSource::build_key_predicate(&cols, &arrays, 3).unwrap()
295        )
296        .to_uppercase();
297        assert!(s.contains("IN") && s.contains("id".to_uppercase().as_str()));
298    }
299
300    #[test]
301    fn composite_predicate_is_or_of_and() {
302        let cols = vec!["a".to_string(), "b".to_string()];
303        let arrays: Vec<ArrayRef> = vec![
304            Arc::new(Int64Array::from(vec![1, 2])),
305            Arc::new(StringArray::from(vec!["x", "y"])),
306        ];
307        let s = format!(
308            "{}",
309            IcebergLookupSource::build_key_predicate(&cols, &arrays, 2).unwrap()
310        )
311        .to_uppercase();
312        assert!(s.contains("AND") && s.contains("OR"));
313    }
314
315    #[test]
316    fn null_key_adds_is_null_term() {
317        let cols = vec!["id".to_string()];
318        let arrays: Vec<ArrayRef> = vec![Arc::new(Int64Array::from(vec![Some(1), None]))];
319        let s = format!(
320            "{}",
321            IcebergLookupSource::build_key_predicate(&cols, &arrays, 2).unwrap()
322        )
323        .to_uppercase();
324        assert!(s.contains("NULL"));
325    }
326}
laminar_connectors/lakehouse/iceberg_lookup.rs

laminar_connectors/lakehouse/
iceberg_lookup.rs
