laminar_connectors/lakehouse/
delta_lookup.rs

1//! Delta Lake on-demand lookup source for cache-miss fallback.
2//!
3//! Implements `LookupSource` backed by a `DataFusion` `TableProvider`. A
4//! batched, typed `pk IN (...)` filter folds all missed keys of a probe into
5//! one file-/partition-pruned scan; [`KeyAligner`] handles key decode and
6//! result realignment.
7
8use std::sync::Arc;
9
10use arrow_array::{Array, ArrayRef, RecordBatch};
11use arrow_row::SortField;
12use arrow_schema::SchemaRef;
13use datafusion::common::ScalarValue;
14use datafusion::prelude::{col, Expr, SessionContext};
15
16use laminar_core::lookup::predicate::Predicate;
17use laminar_core::lookup::source::{
18    projection_names, ColumnId, LookupError, LookupSource, LookupSourceCapabilities,
19};
20use laminar_core::lookup::KeyAligner;
21
22/// Configuration for [`DeltaLookupSource`].
23#[derive(Debug, Clone)]
24pub struct DeltaLookupSourceConfig {
25    /// Table path (resolved, post-catalog).
26    pub table_path: String,
27    /// Storage options (credentials, etc.).
28    pub storage_options: std::collections::HashMap<String, String>,
29    /// Primary key column names.
30    pub primary_key_columns: Vec<String>,
31    /// `DataFusion` table name (registered in session context).
32    pub table_name: String,
33}
34
35/// Delta Lake lookup source for on-demand/partial cache mode.
36pub struct DeltaLookupSource {
37    ctx: Arc<SessionContext>,
38    table_name: String,
39    schema: SchemaRef,
40    aligner: KeyAligner,
41}
42
43impl DeltaLookupSource {
44    /// Opens the Delta table and registers it as a `DataFusion` `TableProvider`.
45    ///
46    /// # Errors
47    ///
48    /// Returns `LookupError` if the table cannot be opened/registered or a
49    /// primary key column is missing from the schema.
50    pub async fn open(config: DeltaLookupSourceConfig) -> Result<Self, LookupError> {
51        let ctx = SessionContext::new();
52        crate::lakehouse::delta_table_provider::register_delta_table(
53            &ctx,
54            &config.table_name,
55            &config.table_path,
56            config.storage_options.clone(),
57        )
58        .await
59        .map_err(|e| LookupError::Connection(format!("register delta table: {e}")))?;
60
61        let table = ctx
62            .table(&config.table_name)
63            .await
64            .map_err(|e| LookupError::Internal(format!("get table: {e}")))?;
65        let schema: SchemaRef = Arc::new(table.schema().as_arrow().clone());
66
67        let pk_sort_fields = pk_sort_fields(&schema, &config.primary_key_columns)?;
68        let aligner = KeyAligner::new(pk_sort_fields, config.primary_key_columns.clone())?;
69
70        warn_if_unclustered(&config).await;
71
72        Ok(Self {
73            ctx: Arc::new(ctx),
74            table_name: config.table_name,
75            schema,
76            aligner,
77        })
78    }
79}
80
81/// Resolve the `RowConverter` sort fields for the primary-key columns.
82fn pk_sort_fields(
83    schema: &SchemaRef,
84    pk_columns: &[String],
85) -> Result<Vec<SortField>, LookupError> {
86    pk_columns
87        .iter()
88        .map(|name| {
89            let idx = schema
90                .index_of(name)
91                .map_err(|_| LookupError::Internal(format!("pk column not found: {name}")))?;
92            Ok(SortField::new(schema.field(idx).data_type().clone()))
93        })
94        .collect()
95}
96
97/// Build a typed `pk IN (...)` (single column) or OR-of-AND-groups (composite)
98/// filter from the decoded primary-key columns. Using typed `Expr` literals
99/// (not string SQL) keeps type handling and escaping correct.
100fn build_in_list_filter(
101    pk_columns: &[String],
102    pk_arrays: &[ArrayRef],
103) -> Result<Expr, LookupError> {
104    let n = if pk_arrays.is_empty() {
105        0
106    } else {
107        pk_arrays[0].len()
108    };
109    let scalar = |arr: &ArrayRef, row: usize| {
110        ScalarValue::try_from_array(arr, row)
111            .map(|sv| Expr::Literal(sv, None))
112            .map_err(|e| LookupError::Internal(format!("scalar from key: {e}")))
113    };
114
115    if pk_columns.len() == 1 {
116        let column = col(&pk_columns[0]);
117        let arr = &pk_arrays[0];
118        let mut lits = Vec::new();
119        let mut has_null = false;
120        for row in 0..n {
121            if arr.is_null(row) {
122                has_null = true;
123            } else {
124                lits.push(scalar(arr, row)?);
125            }
126        }
127        let mut filter = (!lits.is_empty()).then(|| column.clone().in_list(lits, false));
128        if has_null {
129            let is_null = column.is_null();
130            filter = Some(match filter {
131                Some(f) => f.or(is_null),
132                None => is_null,
133            });
134        }
135        return filter.ok_or_else(|| LookupError::Internal("no keys to look up".into()));
136    }
137
138    let mut groups: Vec<Expr> = Vec::with_capacity(n);
139    for row in 0..n {
140        let mut conj: Option<Expr> = None;
141        for (ci, name) in pk_columns.iter().enumerate() {
142            let term = if pk_arrays[ci].is_null(row) {
143                col(name).is_null()
144            } else {
145                col(name).eq(scalar(&pk_arrays[ci], row)?)
146            };
147            conj = Some(match conj {
148                Some(c) => c.and(term),
149                None => term,
150            });
151        }
152        if let Some(c) = conj {
153            groups.push(c);
154        }
155    }
156    let mut it = groups.into_iter();
157    it.next()
158        .map(|first| it.fold(first, Expr::or))
159        .ok_or_else(|| LookupError::Internal("no keys to look up".into()))
160}
161
162/// Best-effort clustering diagnostic: an on-demand lookup is only cheap if the
163/// dimension is partitioned/clustered on the key. Delta exposes partition
164/// columns (not Z-ORDER), so this is a warning, never an error.
165async fn warn_if_unclustered(config: &DeltaLookupSourceConfig) {
166    let Ok(table) = crate::lakehouse::delta_io::open_or_create_table(
167        &config.table_path,
168        config.storage_options.clone(),
169        None,
170    )
171    .await
172    else {
173        return;
174    };
175    let partition_columns = crate::lakehouse::delta_io::get_partition_columns(&table);
176    if !config
177        .primary_key_columns
178        .iter()
179        .any(|k| partition_columns.contains(k))
180    {
181        tracing::warn!(
182            table = %config.table_path,
183            primary_key = ?config.primary_key_columns,
184            partition_columns = ?partition_columns,
185            "delta lookup table is not partitioned on the lookup key; unless it is \
186             Z-ORDER clustered on the key, every cache-miss fetch will full-scan the \
187             table. Cluster the dimension on the lookup key for bounded per-fetch cost."
188        );
189    }
190}
191
192impl LookupSource for DeltaLookupSource {
193    async fn query(
194        &self,
195        keys: &[&[u8]],
196        _predicates: &[Predicate],
197        projection: &[ColumnId],
198    ) -> Result<Vec<Option<RecordBatch>>, LookupError> {
199        if keys.is_empty() {
200            return Ok(Vec::new());
201        }
202        let pk_arrays = self.aligner.decode_keys(keys)?;
203        let filter = build_in_list_filter(self.aligner.pk_columns(), &pk_arrays)?;
204
205        let mut df = self
206            .ctx
207            .table(&self.table_name)
208            .await
209            .map_err(|e| LookupError::Query(format!("open delta table: {e}")))?
210            .filter(filter)
211            .map_err(|e| LookupError::Query(format!("apply lookup filter: {e}")))?;
212        let original_names = if projection.is_empty() {
213            None
214        } else {
215            Some(projection_names(&self.schema, projection)?)
216        };
217        // Projection pushdown: select only the requested columns (the optimizer
218        // pushes this into the Parquet scan). The projection must carry the
219        // key columns so realignment works, then we project them out if unrequested.
220        if !projection.is_empty() {
221            let mut names = projection_names(&self.schema, projection)?;
222            for pk in self.aligner.pk_columns() {
223                if !names.contains(pk) {
224                    names.push(pk.clone());
225                }
226            }
227            let refs: Vec<&str> = names.iter().map(String::as_str).collect();
228            df = df
229                .select_columns(&refs)
230                .map_err(|e| LookupError::Query(format!("apply lookup projection: {e}")))?;
231        }
232        let batches = df
233            .collect()
234            .await
235            .map_err(|e| LookupError::Query(format!("collect lookup results: {e}")))?;
236
237        let aligned = self
238            .aligner
239            .align(keys, &batches)
240            .map_err(|e| LookupError::Internal(format!("align lookup results: {e}")))?;
241
242        if let Some(orig_names) = original_names {
243            let mut projected_aligned = Vec::with_capacity(aligned.len());
244            for maybe_batch in aligned {
245                if let Some(batch) = maybe_batch {
246                    let indices: Vec<usize> = orig_names
247                        .iter()
248                        .map(|name| {
249                            batch.schema().index_of(name).map_err(|e| {
250                                LookupError::Internal(format!(
251                                    "column not found in aligned schema: {e}"
252                                ))
253                            })
254                        })
255                        .collect::<Result<Vec<usize>, LookupError>>()?;
256                    let projected = batch.project(&indices).map_err(|e| {
257                        LookupError::Internal(format!("project aligned batch: {e}"))
258                    })?;
259                    projected_aligned.push(Some(projected));
260                } else {
261                    projected_aligned.push(None);
262                }
263            }
264            Ok(projected_aligned)
265        } else {
266            Ok(aligned)
267        }
268    }
269
270    fn capabilities(&self) -> LookupSourceCapabilities {
271        LookupSourceCapabilities {
272            supports_batch_lookup: true,
273            supports_projection_pushdown: true,
274            ..LookupSourceCapabilities::none()
275        }
276    }
277
278    fn source_name(&self) -> &'static str {
279        "delta-lake"
280    }
281
282    fn schema(&self) -> SchemaRef {
283        Arc::clone(&self.schema)
284    }
285
286    async fn health_check(&self) -> Result<(), LookupError> {
287        self.ctx
288            .table(&self.table_name)
289            .await
290            .map(|_| ())
291            .map_err(|e| LookupError::Connection(format!("health check: {e}")))
292    }
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298    use arrow_array::{Int64Array, StringArray};
299    use arrow_row::RowConverter;
300    use arrow_schema::{DataType, Field, Schema};
301    use std::collections::HashMap;
302    use tempfile::TempDir;
303
304    fn test_schema() -> SchemaRef {
305        Arc::new(Schema::new(vec![
306            Field::new("id", DataType::Int64, false),
307            Field::new("name", DataType::Utf8, true),
308        ]))
309    }
310
311    fn test_batch(ids: &[i64], names: &[&str]) -> RecordBatch {
312        RecordBatch::try_new(
313            test_schema(),
314            vec![
315                Arc::new(Int64Array::from(ids.to_vec())),
316                Arc::new(StringArray::from(names.to_vec())),
317            ],
318        )
319        .unwrap()
320    }
321
322    fn int_keys(ids: &[i64]) -> Vec<Vec<u8>> {
323        let converter = RowConverter::new(vec![SortField::new(DataType::Int64)]).unwrap();
324        let rows = converter
325            .convert_columns(&[Arc::new(Int64Array::from(ids.to_vec()))])
326            .unwrap();
327        (0..ids.len())
328            .map(|i| rows.row(i).as_ref().to_vec())
329            .collect()
330    }
331
332    async fn create_delta_table(path: &str, batches: Vec<RecordBatch>) {
333        use crate::lakehouse::delta_io;
334        use deltalake::protocol::SaveMode;
335
336        let table = delta_io::open_or_create_table(path, HashMap::new(), Some(&test_schema()))
337            .await
338            .unwrap();
339        delta_io::write_batches(table, batches, SaveMode::Append, None, false, None, None)
340            .await
341            .unwrap();
342    }
343
344    async fn open_source(path: &str, table_name: &str) -> DeltaLookupSource {
345        DeltaLookupSource::open(DeltaLookupSourceConfig {
346            table_path: path.to_string(),
347            storage_options: HashMap::new(),
348            primary_key_columns: vec!["id".into()],
349            table_name: table_name.to_string(),
350        })
351        .await
352        .unwrap()
353    }
354
355    fn id_at(batch: &RecordBatch) -> i64 {
356        batch
357            .column(0)
358            .as_any()
359            .downcast_ref::<Int64Array>()
360            .unwrap()
361            .value(0)
362    }
363
364    #[tokio::test]
365    async fn batched_lookup_aligns_hits_and_misses() {
366        let temp_dir = TempDir::new().unwrap();
367        let path = temp_dir.path().to_str().unwrap();
368        create_delta_table(path, vec![test_batch(&[1, 2, 3], &["A", "B", "C"])]).await;
369        let source = open_source(path, "lk").await;
370
371        // Out-of-table-order with a miss; one batched fetch.
372        let keys = int_keys(&[3, 1, 999, 2]);
373        let key_refs: Vec<&[u8]> = keys.iter().map(Vec::as_slice).collect();
374        let results = source.query(&key_refs, &[], &[]).await.unwrap();
375
376        assert_eq!(results.len(), 4);
377        assert_eq!(id_at(results[0].as_ref().unwrap()), 3);
378        assert_eq!(id_at(results[1].as_ref().unwrap()), 1);
379        assert!(results[2].is_none());
380        assert_eq!(id_at(results[3].as_ref().unwrap()), 2);
381    }
382
383    /// The Phase 3 exit criterion: a batched `pk IN (...)` must prune
384    /// non-matching partition files (per-key cost O(matching files), not
385    /// O(table)) on a table clustered on the key.
386    #[tokio::test]
387    async fn in_list_prunes_partition_files() {
388        use datafusion::physical_plan::collect;
389
390        let temp_dir = TempDir::new().unwrap();
391        let path = temp_dir.path().to_str().unwrap();
392
393        // 8 distinct keys → 8 partition directories (one Parquet file each).
394        {
395            use crate::lakehouse::delta_io;
396            use deltalake::protocol::SaveMode;
397            let t = delta_io::open_or_create_table(path, HashMap::new(), None)
398                .await
399                .unwrap();
400            delta_io::write_batches(
401                t,
402                vec![test_batch(
403                    &[0, 1, 2, 3, 4, 5, 6, 7],
404                    &["a", "b", "c", "d", "e", "f", "g", "h"],
405                )],
406                SaveMode::Append,
407                Some(&["id".to_string()]),
408                false,
409                None,
410                None,
411            )
412            .await
413            .unwrap();
414        }
415
416        // Correctness across partition files via the source.
417        let source = open_source(path, "lk").await;
418        let keys = int_keys(&[5, 2, 100]);
419        let key_refs: Vec<&[u8]> = keys.iter().map(Vec::as_slice).collect();
420        let results = source.query(&key_refs, &[], &[]).await.unwrap();
421        assert!(results[0].is_some() && results[1].is_some() && results[2].is_none());
422
423        // Pruning: the IN-list query reads fewer than all 8 files. (`next`
424        // provider reports files read via `count_files_scanned`.)
425        let ctx = SessionContext::new();
426        crate::lakehouse::delta_table_provider::register_delta_table(
427            &ctx,
428            "lk",
429            path,
430            HashMap::new(),
431        )
432        .await
433        .unwrap();
434        let plan = ctx
435            .sql("SELECT * FROM \"lk\" WHERE \"id\" IN (2, 5)")
436            .await
437            .unwrap()
438            .create_physical_plan()
439            .await
440            .unwrap();
441        let _ = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap();
442        let scanned = sum_plan_metric(&plan, "count_files_scanned");
443        assert!(
444            scanned > 0 && scanned < 8,
445            "expected pruning, scanned={scanned}"
446        );
447    }
448
449    fn sum_plan_metric(
450        plan: &Arc<dyn datafusion::physical_plan::ExecutionPlan>,
451        name: &str,
452    ) -> usize {
453        let mut total = plan
454            .metrics()
455            .and_then(|m| m.sum_by_name(name))
456            .map_or(0, |v| v.as_usize());
457        for child in plan.children() {
458            total += sum_plan_metric(child, name);
459        }
460        total
461    }
462}
laminar_connectors/lakehouse/delta_lookup.rs

laminar_connectors/lakehouse/
delta_lookup.rs
