laminar_connectors/lakehouse/delta_lookup/
mod.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            primary_key = ?config.primary_key_columns,
183            partition_columns = ?partition_columns,
184            "delta lookup table is not partitioned on the lookup key; unless it is \
185             Z-ORDER clustered on the key, every cache-miss fetch will full-scan the \
186             table. Cluster the dimension on the lookup key for bounded per-fetch cost."
187        );
188    }
189}
190
191impl LookupSource for DeltaLookupSource {
192    async fn query(
193        &self,
194        keys: &[&[u8]],
195        _predicates: &[Predicate],
196        projection: &[ColumnId],
197    ) -> Result<Vec<Option<RecordBatch>>, LookupError> {
198        if keys.is_empty() {
199            return Ok(Vec::new());
200        }
201        let pk_arrays = self.aligner.decode_keys(keys)?;
202        let filter = build_in_list_filter(self.aligner.pk_columns(), &pk_arrays)?;
203
204        let mut df = self
205            .ctx
206            .table(&self.table_name)
207            .await
208            .map_err(|e| LookupError::Query(format!("open delta table: {e}")))?
209            .filter(filter)
210            .map_err(|e| LookupError::Query(format!("apply lookup filter: {e}")))?;
211        let original_names = if projection.is_empty() {
212            None
213        } else {
214            Some(projection_names(&self.schema, projection)?)
215        };
216        // Projection pushdown: select only the requested columns (the optimizer
217        // pushes this into the Parquet scan). The projection must carry the
218        // key columns so realignment works, then we project them out if unrequested.
219        if !projection.is_empty() {
220            let mut names = projection_names(&self.schema, projection)?;
221            for pk in self.aligner.pk_columns() {
222                if !names.contains(pk) {
223                    names.push(pk.clone());
224                }
225            }
226            let refs: Vec<&str> = names.iter().map(String::as_str).collect();
227            df = df
228                .select_columns(&refs)
229                .map_err(|e| LookupError::Query(format!("apply lookup projection: {e}")))?;
230        }
231        let batches = df
232            .collect()
233            .await
234            .map_err(|e| LookupError::Query(format!("collect lookup results: {e}")))?;
235
236        let aligned = self
237            .aligner
238            .align(keys, &batches)
239            .map_err(|e| LookupError::Internal(format!("align lookup results: {e}")))?;
240
241        if let Some(orig_names) = original_names {
242            let mut projected_aligned = Vec::with_capacity(aligned.len());
243            for maybe_batch in aligned {
244                if let Some(batch) = maybe_batch {
245                    let indices: Vec<usize> = orig_names
246                        .iter()
247                        .map(|name| {
248                            batch.schema().index_of(name).map_err(|e| {
249                                LookupError::Internal(format!(
250                                    "column not found in aligned schema: {e}"
251                                ))
252                            })
253                        })
254                        .collect::<Result<Vec<usize>, LookupError>>()?;
255                    let projected = batch.project(&indices).map_err(|e| {
256                        LookupError::Internal(format!("project aligned batch: {e}"))
257                    })?;
258                    projected_aligned.push(Some(projected));
259                } else {
260                    projected_aligned.push(None);
261                }
262            }
263            Ok(projected_aligned)
264        } else {
265            Ok(aligned)
266        }
267    }
268
269    fn capabilities(&self) -> LookupSourceCapabilities {
270        LookupSourceCapabilities {
271            supports_batch_lookup: true,
272            supports_projection_pushdown: true,
273            ..LookupSourceCapabilities::none()
274        }
275    }
276
277    fn source_name(&self) -> &'static str {
278        "delta-lake"
279    }
280
281    fn schema(&self) -> SchemaRef {
282        Arc::clone(&self.schema)
283    }
284
285    async fn health_check(&self) -> Result<(), LookupError> {
286        self.ctx
287            .table(&self.table_name)
288            .await
289            .map(|_| ())
290            .map_err(|e| LookupError::Connection(format!("health check: {e}")))
291    }
292}
293
294#[cfg(test)]
295mod tests;
laminar_connectors/lakehouse/delta_lookup/mod.rs

laminar_connectors/lakehouse/delta_lookup/
mod.rs
