laminar_core/lookup/source/
mod.rs

1//! Async lookup source trait with predicate and projection pushdown.
2
3use std::future::Future;
4use std::time::Duration;
5
6use arrow::compute::filter_record_batch;
7use arrow_array::BooleanArray;
8use arrow_array::RecordBatch;
9use arrow_schema::SchemaRef;
10
11use crate::lookup::predicate::{split_predicates, Predicate, ScalarValue, SourceCapabilities};
12
13/// Column identifier for projection pushdown.
14pub type ColumnId = u32;
15
16/// Errors from lookup source operations.
17#[derive(Debug, thiserror::Error)]
18pub enum LookupError {
19    /// Connection to the external system failed.
20    #[error("connection failed: {0}")]
21    Connection(String),
22
23    /// Query execution failed.
24    #[error("query failed: {0}")]
25    Query(String),
26
27    /// The operation timed out.
28    #[error("timeout after {0:?}")]
29    Timeout(Duration),
30
31    /// The source is not available (e.g., not initialized).
32    #[error("not available: {0}")]
33    NotAvailable(String),
34
35    /// Internal error (cache I/O, codec failure, etc.).
36    #[error("internal: {0}")]
37    Internal(String),
38}
39
40/// Capabilities that a lookup source advertises.
41///
42/// This describes source-level capabilities (batch support, pushdown
43/// support) rather than per-column capabilities (which are described
44/// by [`SourceCapabilities`]).
45#[derive(Debug, Clone, Default)]
46pub struct LookupSourceCapabilities {
47    /// Whether the source supports predicate pushdown.
48    pub supports_predicate_pushdown: bool,
49    /// Whether the source supports projection pushdown.
50    pub supports_projection_pushdown: bool,
51    /// Whether the source supports batch lookups.
52    pub supports_batch_lookup: bool,
53    /// Maximum batch size for batch lookups (0 = unlimited).
54    pub max_batch_size: usize,
55}
56
57impl LookupSourceCapabilities {
58    /// Create capabilities with no pushdown support.
59    #[must_use]
60    pub fn none() -> Self {
61        Self::default()
62    }
63}
64
65/// Resolve projection column indices to their names in `schema`, in projection
66/// order. An empty projection means "all columns" and returns every name in
67/// schema order. Shared by the lookup backends to turn the `query` projection
68/// into a column selection.
69///
70/// # Errors
71/// Returns [`LookupError::Internal`] if a projection index is out of range.
72pub fn projection_names(
73    schema: &SchemaRef,
74    projection: &[ColumnId],
75) -> Result<Vec<String>, LookupError> {
76    if projection.is_empty() {
77        return Ok(schema.fields().iter().map(|f| f.name().clone()).collect());
78    }
79    projection
80        .iter()
81        .map(|&c| {
82            schema
83                .fields()
84                .get(c as usize)
85                .map(|f| f.name().clone())
86                .ok_or_else(|| LookupError::Internal(format!("projection column {c} out of range")))
87        })
88        .collect()
89}
90
91/// Async data source for lookup table refresh and query.
92///
93/// This trait uses RPITIT (return-position `impl Trait` in traits,
94/// stabilized in Rust 1.75) for zero-overhead async dispatch.
95///
96/// ## Implementing
97///
98/// Sources that support predicate/projection pushdown should set the
99/// corresponding flags in [`capabilities()`](Self::capabilities) and
100/// handle filtered queries in [`query()`](Self::query). Sources that
101/// do not support pushdown can be wrapped in [`PushdownAdapter`] to
102/// get automatic local evaluation.
103pub trait LookupSource: Send + Sync {
104    /// Query the source by keys, predicates, and/or projection.
105    ///
106    /// Returns a `Vec<Option<RecordBatch>>` aligned with the input `keys`:
107    /// - `Some(batch)` — key found, value is a single-row `RecordBatch`
108    /// - `None` — key not found
109    fn query(
110        &self,
111        keys: &[&[u8]],
112        predicates: &[Predicate],
113        projection: &[ColumnId],
114    ) -> impl Future<Output = Result<Vec<Option<RecordBatch>>, LookupError>> + Send;
115
116    /// Capabilities this source advertises.
117    fn capabilities(&self) -> LookupSourceCapabilities;
118
119    /// Source name for logging and metrics.
120    fn source_name(&self) -> &'static str;
121
122    /// Arrow schema of the data this source returns.
123    fn schema(&self) -> SchemaRef;
124
125    /// Optional row count estimate for query planning.
126    fn estimated_row_count(&self) -> Option<u64> {
127        None
128    }
129
130    /// Health check. Default: always healthy.
131    fn health_check(&self) -> impl Future<Output = Result<(), LookupError>> + Send {
132        async { Ok(()) }
133    }
134}
135
136/// Dyn-compatible version of [`LookupSource`] for use as `Arc<dyn LookupSourceDyn>`.
137///
138/// `LookupSource` uses RPITIT which is not dyn-compatible. This trait
139/// uses `async_trait` boxing instead, suitable for the cold path
140/// (cache miss → source query).
141#[async_trait::async_trait]
142pub trait LookupSourceDyn: Send + Sync {
143    /// Query the source by keys, predicates, and/or projection.
144    async fn query_batch(
145        &self,
146        keys: &[&[u8]],
147        predicates: &[Predicate],
148        projection: &[ColumnId],
149    ) -> Result<Vec<Option<RecordBatch>>, LookupError>;
150
151    /// Arrow schema of the data this source returns.
152    fn schema(&self) -> SchemaRef;
153}
154
155#[async_trait::async_trait]
156impl<T: LookupSource> LookupSourceDyn for T {
157    async fn query_batch(
158        &self,
159        keys: &[&[u8]],
160        predicates: &[Predicate],
161        projection: &[ColumnId],
162    ) -> Result<Vec<Option<RecordBatch>>, LookupError> {
163        self.query(keys, predicates, projection).await
164    }
165
166    fn schema(&self) -> SchemaRef {
167        LookupSource::schema(self)
168    }
169}
170
171/// Wraps a [`LookupSource`] that doesn't support pushdown.
172///
173/// Predicates that can be pushed down (according to [`SourceCapabilities`])
174/// are forwarded to the inner source. Remaining predicates are evaluated
175/// locally after fetching results using Arrow SIMD filtering.
176pub struct PushdownAdapter<S> {
177    inner: S,
178    column_capabilities: SourceCapabilities,
179}
180
181impl<S: LookupSource> PushdownAdapter<S> {
182    /// Create a new adapter.
183    ///
184    /// * `inner` — the underlying source
185    /// * `column_capabilities` — per-column pushdown capabilities used
186    ///   by [`split_predicates`]
187    pub fn new(inner: S, column_capabilities: SourceCapabilities) -> Self {
188        Self {
189            inner,
190            column_capabilities,
191        }
192    }
193
194    /// Split predicates into pushable and local sets.
195    fn split(&self, predicates: &[Predicate]) -> (Vec<Predicate>, Vec<Predicate>) {
196        let split = split_predicates(predicates.to_vec(), &self.column_capabilities);
197        (split.pushable, split.local)
198    }
199}
200
201/// Apply a comparison from `arrow::compute::kernels::cmp` between a column
202/// and a scalar value. Builds a typed single-element array for the scalar
203/// side so `Scalar<T>` implements `Datum`.
204fn compare_column_scalar(
205    batch: &RecordBatch,
206    column: &str,
207    value: &ScalarValue,
208    cmp_fn: fn(
209        &dyn arrow_array::Datum,
210        &dyn arrow_array::Datum,
211    ) -> Result<BooleanArray, arrow::error::ArrowError>,
212) -> Option<BooleanArray> {
213    use arrow_array::types::{TimestampMicrosecondType, TimestampMillisecondType};
214    use arrow_array::{Float64Array, Int64Array, PrimitiveArray, Scalar, StringArray};
215
216    let idx = batch.schema().index_of(column).ok()?;
217    let col = batch.column(idx);
218    match value {
219        ScalarValue::Int64(v) => cmp_fn(col, &Scalar::new(Int64Array::from(vec![*v]))).ok(),
220        ScalarValue::Float64(v) => cmp_fn(col, &Scalar::new(Float64Array::from(vec![*v]))).ok(),
221        ScalarValue::Utf8(v) => cmp_fn(col, &Scalar::new(StringArray::from(vec![v.as_str()]))).ok(),
222        ScalarValue::Bool(v) => cmp_fn(col, &Scalar::new(BooleanArray::from(vec![*v]))).ok(),
223        ScalarValue::Timestamp(us) => {
224            if col
225                .as_any()
226                .is::<PrimitiveArray<TimestampMicrosecondType>>()
227            {
228                let scalar = PrimitiveArray::<TimestampMicrosecondType>::from(vec![*us]);
229                cmp_fn(col, &Scalar::new(scalar)).ok()
230            } else if col
231                .as_any()
232                .is::<PrimitiveArray<TimestampMillisecondType>>()
233            {
234                let scalar = PrimitiveArray::<TimestampMillisecondType>::from(vec![*us / 1000]);
235                cmp_fn(col, &Scalar::new(scalar)).ok()
236            } else {
237                None
238            }
239        }
240        _ => None,
241    }
242}
243
244/// Evaluate a single predicate against a `RecordBatch`, returning a boolean mask.
245fn evaluate_predicate(batch: &RecordBatch, predicate: &Predicate) -> Option<BooleanArray> {
246    use arrow::compute::kernels::cmp;
247
248    match predicate {
249        Predicate::Eq { column, value } => compare_column_scalar(batch, column, value, cmp::eq),
250        Predicate::NotEq { column, value } => compare_column_scalar(batch, column, value, cmp::neq),
251        Predicate::Lt { column, value } => compare_column_scalar(batch, column, value, cmp::lt),
252        Predicate::LtEq { column, value } => {
253            compare_column_scalar(batch, column, value, cmp::lt_eq)
254        }
255        Predicate::Gt { column, value } => compare_column_scalar(batch, column, value, cmp::gt),
256        Predicate::GtEq { column, value } => {
257            compare_column_scalar(batch, column, value, cmp::gt_eq)
258        }
259        Predicate::IsNull { column } => {
260            let idx = batch.schema().index_of(column).ok()?;
261            let col = batch.column(idx);
262            Some(arrow::compute::is_null(col).ok()?)
263        }
264        Predicate::IsNotNull { column } => {
265            let idx = batch.schema().index_of(column).ok()?;
266            let col = batch.column(idx);
267            Some(arrow::compute::is_not_null(col).ok()?)
268        }
269        Predicate::In { column, values } => {
270            let idx = batch.schema().index_of(column).ok()?;
271            let col = batch.column(idx);
272            let mut mask: Option<BooleanArray> = None;
273            for v in values {
274                let eq_mask = evaluate_predicate(
275                    batch,
276                    &Predicate::Eq {
277                        column: column.clone(),
278                        value: v.clone(),
279                    },
280                )?;
281                mask = Some(match mask {
282                    Some(existing) => arrow::compute::or(&existing, &eq_mask).ok()?,
283                    None => eq_mask,
284                });
285            }
286            mask.or_else(|| Some(BooleanArray::from(vec![false; col.len()])))
287        }
288    }
289}
290
291/// Apply local predicates to a `RecordBatch`, filtering out non-matching rows.
292fn apply_local_predicates(batch: &RecordBatch, predicates: &[Predicate]) -> Option<RecordBatch> {
293    if predicates.is_empty() {
294        return Some(batch.clone());
295    }
296    let mut combined: Option<BooleanArray> = None;
297    for pred in predicates {
298        let mask = evaluate_predicate(batch, pred)?;
299        combined = Some(match combined {
300            Some(existing) => arrow::compute::and(&existing, &mask).ok()?,
301            None => mask,
302        });
303    }
304    match combined {
305        Some(mask) => filter_record_batch(batch, &mask).ok(),
306        None => Some(batch.clone()),
307    }
308}
309
310impl<S: LookupSource> LookupSource for PushdownAdapter<S> {
311    async fn query(
312        &self,
313        keys: &[&[u8]],
314        predicates: &[Predicate],
315        projection: &[ColumnId],
316    ) -> Result<Vec<Option<RecordBatch>>, LookupError> {
317        let (pushable, local) = self.split(predicates);
318        let results = self.inner.query(keys, &pushable, projection).await?;
319
320        if local.is_empty() {
321            return Ok(results);
322        }
323
324        Ok(results
325            .into_iter()
326            .map(|opt| {
327                opt.and_then(|batch| {
328                    let filtered = apply_local_predicates(&batch, &local)?;
329                    if filtered.num_rows() == 0 {
330                        None
331                    } else {
332                        Some(filtered)
333                    }
334                })
335            })
336            .collect())
337    }
338
339    fn capabilities(&self) -> LookupSourceCapabilities {
340        self.inner.capabilities()
341    }
342
343    fn source_name(&self) -> &'static str {
344        self.inner.source_name()
345    }
346
347    fn schema(&self) -> SchemaRef {
348        self.inner.schema()
349    }
350
351    fn estimated_row_count(&self) -> Option<u64> {
352        self.inner.estimated_row_count()
353    }
354
355    fn health_check(&self) -> impl Future<Output = Result<(), LookupError>> + Send {
356        self.inner.health_check()
357    }
358}
359
360#[cfg(test)]
361#[allow(clippy::disallowed_types)] // test code
362mod tests;
laminar_core/lookup/source/mod.rs

laminar_core/lookup/source/
mod.rs
