laminar_sql/datafusion/lookup_join_exec/
mod.rs

1//! Physical execution plan for lookup joins.
2//!
3//! Bridges `LookupJoinNode` (logical) to a hash-probe executor that
4//! joins streaming input against a pre-indexed lookup table snapshot.
5//!
6//! ## Data flow
7//!
8//! ```text
9//! Stream input ──► LookupJoinExec ──► Output (stream + lookup columns)
10//!                       │
11//!                  HashIndex probe
12//!                       │
13//!                  LookupSnapshot (pre-indexed RecordBatch)
14//! ```
15
16use std::collections::HashMap;
17use std::fmt::{self, Debug, Formatter};
18use std::sync::Arc;
19
20use parking_lot::RwLock;
21
22use arrow::compute::take;
23use arrow::row::{RowConverter, SortField};
24use arrow_array::{Array, RecordBatch, UInt32Array};
25use arrow_schema::{Schema, SchemaRef};
26use async_trait::async_trait;
27use datafusion::execution::{SendableRecordBatchStream, SessionState, TaskContext};
28use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode};
29use datafusion::physical_expr::{EquivalenceProperties, LexOrdering, Partitioning};
30use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
31use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
32use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
33use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
34use datafusion_common::{DataFusionError, Result};
35use datafusion_expr::Expr;
36use futures::StreamExt;
37use laminar_core::lookup::lookup_cache::LookupMemoryCache;
38use laminar_core::lookup::source::{ColumnId, LookupSourceDyn};
39use tokio::sync::Semaphore;
40
41use super::lookup_join::{LookupJoinNode, LookupJoinType};
42
43mod partial;
44
45pub use partial::PartialLookupJoinExec;
46
47/// Deadline for a cache-miss source fetch. The fetch is awaited inline on the
48/// single compute thread, so without a bound a hung source wedges the pipeline.
49const LOOKUP_SOURCE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
50
51// ── Registry ─────────────────────────────────────────────────────
52
53/// Thread-safe registry of lookup table entries (snapshot or partial).
54///
55/// The db layer populates this when `CREATE LOOKUP TABLE` executes;
56/// the [`LookupJoinExtensionPlanner`] reads it at physical plan time.
57#[derive(Default)]
58pub struct LookupTableRegistry {
59    tables: RwLock<HashMap<String, RegisteredLookup>>,
60}
61
62/// A registered lookup table entry — snapshot or partial (on-demand).
63pub enum RegisteredLookup {
64    /// Full snapshot: all rows pre-loaded in a single batch.
65    Snapshot(Arc<LookupSnapshot>),
66    /// Partial (on-demand): bounded lookup cache with S3-FIFO eviction.
67    Partial(Arc<PartialLookupState>),
68}
69
70/// Point-in-time snapshot of a lookup table for join execution.
71pub struct LookupSnapshot {
72    /// All rows concatenated into a single batch.
73    pub batch: RecordBatch,
74}
75
76/// State for a partial (on-demand) lookup table.
77pub struct PartialLookupState {
78    /// Bounded in-memory cache with S3-FIFO eviction.
79    pub lookup_cache: Arc<LookupMemoryCache>,
80    /// Schema of the lookup table.
81    pub schema: SchemaRef,
82    /// Key column names for row encoding.
83    pub key_columns: Vec<String>,
84    /// `SortField` descriptors for key encoding via `RowConverter`.
85    pub key_sort_fields: Vec<SortField>,
86    /// Async source for cache miss fallback (None = cache-only mode).
87    pub source: Option<Arc<dyn LookupSourceDyn>>,
88    /// Limits concurrent source queries to avoid overloading the source.
89    pub fetch_semaphore: Arc<Semaphore>,
90    /// Column indices (into `schema`) to fetch from the source — the union of
91    /// every column any query references, plus the key. Empty = fetch all.
92    /// Per-table (the cache is shared across queries), so it must be a superset.
93    pub projection: Vec<ColumnId>,
94}
95
96impl LookupTableRegistry {
97    /// Creates an empty registry.
98    #[must_use]
99    pub fn new() -> Self {
100        Self::default()
101    }
102
103    /// Registers or replaces a lookup table snapshot.
104    ///
105    /// # Panics
106    ///
107    /// Panics if the internal lock is poisoned.
108    pub fn register(&self, name: &str, snapshot: LookupSnapshot) {
109        self.tables.write().insert(
110            name.to_lowercase(),
111            RegisteredLookup::Snapshot(Arc::new(snapshot)),
112        );
113    }
114
115    /// Registers or replaces a partial (on-demand) lookup table.
116    ///
117    /// # Panics
118    ///
119    /// Panics if the internal lock is poisoned.
120    pub fn register_partial(&self, name: &str, state: PartialLookupState) {
121        self.tables.write().insert(
122            name.to_lowercase(),
123            RegisteredLookup::Partial(Arc::new(state)),
124        );
125    }
126
127    /// Removes a lookup table from the registry.
128    ///
129    /// # Panics
130    ///
131    /// Panics if the internal lock is poisoned.
132    pub fn unregister(&self, name: &str) {
133        self.tables.write().remove(&name.to_lowercase());
134    }
135
136    /// Returns the current snapshot for a table, if registered as a snapshot.
137    ///
138    /// # Panics
139    ///
140    /// Panics if the internal lock is poisoned.
141    #[must_use]
142    pub fn get(&self, name: &str) -> Option<Arc<LookupSnapshot>> {
143        let tables = self.tables.read();
144        match tables.get(&name.to_lowercase())? {
145            RegisteredLookup::Snapshot(s) => Some(Arc::clone(s)),
146            RegisteredLookup::Partial(_) => None,
147        }
148    }
149
150    /// Returns the registered lookup entry (snapshot or partial).
151    ///
152    /// # Panics
153    ///
154    /// Panics if the internal lock is poisoned.
155    pub fn get_entry(&self, name: &str) -> Option<RegisteredLookup> {
156        let tables = self.tables.read();
157        tables.get(&name.to_lowercase()).map(|e| match e {
158            RegisteredLookup::Snapshot(s) => RegisteredLookup::Snapshot(Arc::clone(s)),
159            RegisteredLookup::Partial(p) => RegisteredLookup::Partial(Arc::clone(p)),
160        })
161    }
162}
163
164// ── Hash Index ───────────────────────────────────────────────────
165
166/// Pre-built hash index mapping encoded key bytes to row indices.
167struct HashIndex {
168    map: HashMap<Box<[u8]>, Vec<u32>>,
169}
170
171impl HashIndex {
172    /// Builds an index over `key_indices` columns in `batch`.
173    ///
174    /// Uses Arrow's `RowConverter` for binary-comparable key encoding
175    /// so any Arrow data type is handled without manual serialization.
176    fn build(batch: &RecordBatch, key_indices: &[usize]) -> Result<Self> {
177        if batch.num_rows() == 0 {
178            return Ok(Self {
179                map: HashMap::new(),
180            });
181        }
182
183        let sort_fields: Vec<SortField> = key_indices
184            .iter()
185            .map(|&i| SortField::new(batch.schema().field(i).data_type().clone()))
186            .collect();
187        let converter = RowConverter::new(sort_fields)?;
188
189        let key_cols: Vec<_> = key_indices
190            .iter()
191            .map(|&i| batch.column(i).clone())
192            .collect();
193        let rows = converter.convert_columns(&key_cols)?;
194
195        let num_rows = batch.num_rows();
196        let mut map: HashMap<Box<[u8]>, Vec<u32>> = HashMap::with_capacity(num_rows);
197        #[allow(clippy::cast_possible_truncation)] // batch row count fits u32
198        for i in 0..num_rows {
199            map.entry(Box::from(rows.row(i).as_ref()))
200                .or_default()
201                .push(i as u32);
202        }
203
204        Ok(Self { map })
205    }
206
207    fn probe(&self, key: &[u8]) -> Option<&[u32]> {
208        self.map.get(key).map(Vec::as_slice)
209    }
210}
211
212// ── Physical Execution Plan ──────────────────────────────────────
213
214/// Physical plan that hash-probes a pre-indexed lookup table for
215/// each batch from the streaming input.
216pub struct LookupJoinExec {
217    input: Arc<dyn ExecutionPlan>,
218    index: Arc<HashIndex>,
219    lookup_batch: Arc<RecordBatch>,
220    stream_key_indices: Vec<usize>,
221    join_type: LookupJoinType,
222    schema: SchemaRef,
223    properties: Arc<PlanProperties>,
224    /// Prebuilt `RowConverter` for encoding probe keys. Shared across
225    /// every `execute()` call so we don't rebuild per-type encoders on
226    /// every cycle of a cached physical plan.
227    converter: Arc<RowConverter>,
228    stream_field_count: usize,
229    projection: Vec<usize>,
230}
231
232impl LookupJoinExec {
233    /// Creates a new lookup join executor.
234    ///
235    /// `stream_key_indices` and `lookup_key_indices` must be the same
236    /// length and correspond pairwise (stream key 0 matches lookup key 0).
237    ///
238    /// # Errors
239    ///
240    /// Returns an error if the hash index cannot be built (e.g., unsupported key type).
241    #[allow(clippy::needless_pass_by_value)] // lookup_batch is moved into Arc
242    pub fn try_new(
243        input: Arc<dyn ExecutionPlan>,
244        lookup_batch: RecordBatch,
245        stream_key_indices: Vec<usize>,
246        lookup_key_indices: Vec<usize>,
247        join_type: LookupJoinType,
248        output_schema: SchemaRef,
249    ) -> Result<Self> {
250        let index = HashIndex::build(&lookup_batch, &lookup_key_indices)?;
251
252        let key_sort_fields: Vec<SortField> = lookup_key_indices
253            .iter()
254            .map(|&i| SortField::new(lookup_batch.schema().field(i).data_type().clone()))
255            .collect();
256        let converter = Arc::new(RowConverter::new(key_sort_fields)?);
257
258        // Left outer joins produce NULLs for non-matching lookup rows,
259        // so force all lookup columns nullable in the output schema.
260        let output_schema = if join_type == LookupJoinType::LeftOuter {
261            let stream_count = input.schema().fields().len();
262            let mut fields = output_schema.fields().to_vec();
263            for f in &mut fields[stream_count..] {
264                if !f.is_nullable() {
265                    *f = Arc::new(f.as_ref().clone().with_nullable(true));
266                }
267            }
268            Arc::new(Schema::new_with_metadata(
269                fields,
270                output_schema.metadata().clone(),
271            ))
272        } else {
273            output_schema
274        };
275
276        let properties = Arc::new(PlanProperties::new(
277            EquivalenceProperties::new(Arc::clone(&output_schema)),
278            Partitioning::UnknownPartitioning(1),
279            EmissionType::Incremental,
280            Boundedness::Unbounded {
281                requires_infinite_memory: false,
282            },
283        ));
284
285        let stream_field_count = input.schema().fields().len();
286        let projection = (0..(stream_field_count + lookup_batch.num_columns())).collect();
287
288        Ok(Self {
289            input,
290            index: Arc::new(index),
291            lookup_batch: Arc::new(lookup_batch),
292            stream_key_indices,
293            join_type,
294            schema: output_schema,
295            properties,
296            converter,
297            stream_field_count,
298            projection,
299        })
300    }
301
302    /// Sets the projection list for output columns.
303    #[must_use]
304    pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
305        self.projection = projection;
306        self
307    }
308}
309
310impl Debug for LookupJoinExec {
311    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
312        f.debug_struct("LookupJoinExec")
313            .field("join_type", &self.join_type)
314            .field("stream_keys", &self.stream_key_indices)
315            .field("lookup_rows", &self.lookup_batch.num_rows())
316            .finish_non_exhaustive()
317    }
318}
319
320impl DisplayAs for LookupJoinExec {
321    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result {
322        match t {
323            DisplayFormatType::Default | DisplayFormatType::Verbose => {
324                write!(
325                    f,
326                    "LookupJoinExec: type={}, stream_keys={:?}, lookup_rows={}",
327                    self.join_type,
328                    self.stream_key_indices,
329                    self.lookup_batch.num_rows(),
330                )
331            }
332            DisplayFormatType::TreeRender => write!(f, "LookupJoinExec"),
333        }
334    }
335}
336
337impl ExecutionPlan for LookupJoinExec {
338    fn as_any(&self) -> &dyn std::any::Any {
339        self
340    }
341
342    fn name(&self) -> &'static str {
343        "LookupJoinExec"
344    }
345
346    fn schema(&self) -> SchemaRef {
347        Arc::clone(&self.schema)
348    }
349
350    fn properties(&self) -> &Arc<PlanProperties> {
351        &self.properties
352    }
353
354    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
355        vec![&self.input]
356    }
357
358    fn with_new_children(
359        self: Arc<Self>,
360        mut children: Vec<Arc<dyn ExecutionPlan>>,
361    ) -> Result<Arc<dyn ExecutionPlan>> {
362        if children.len() != 1 {
363            return Err(DataFusionError::Plan(
364                "LookupJoinExec requires exactly one child".into(),
365            ));
366        }
367        Ok(Arc::new(Self {
368            input: children.swap_remove(0),
369            index: Arc::clone(&self.index),
370            lookup_batch: Arc::clone(&self.lookup_batch),
371            stream_key_indices: self.stream_key_indices.clone(),
372            join_type: self.join_type,
373            schema: Arc::clone(&self.schema),
374            properties: self.properties.clone(),
375            converter: Arc::clone(&self.converter),
376            stream_field_count: self.stream_field_count,
377            projection: self.projection.clone(),
378        }))
379    }
380
381    fn execute(
382        &self,
383        partition: usize,
384        context: Arc<TaskContext>,
385    ) -> Result<SendableRecordBatchStream> {
386        let input_stream = self.input.execute(partition, context)?;
387        let converter = Arc::clone(&self.converter);
388        let index = Arc::clone(&self.index);
389        let lookup_batch = Arc::clone(&self.lookup_batch);
390        let stream_key_indices = self.stream_key_indices.clone();
391        let join_type = self.join_type;
392        let schema = self.schema();
393        let stream_field_count = self.stream_field_count;
394        let projection = self.projection.clone();
395
396        let output = input_stream.map(move |result| {
397            let batch = result?;
398            if batch.num_rows() == 0 {
399                return Ok(RecordBatch::new_empty(Arc::clone(&schema)));
400            }
401            probe_batch(
402                &batch,
403                &converter,
404                &index,
405                &lookup_batch,
406                &stream_key_indices,
407                join_type,
408                &schema,
409                stream_field_count,
410                &projection,
411            )
412        });
413
414        Ok(Box::pin(RecordBatchStreamAdapter::new(
415            self.schema(),
416            output,
417        )))
418    }
419}
420
421impl datafusion::physical_plan::ExecutionPlanProperties for LookupJoinExec {
422    fn output_partitioning(&self) -> &Partitioning {
423        self.properties.output_partitioning()
424    }
425
426    fn output_ordering(&self) -> Option<&LexOrdering> {
427        self.properties.output_ordering()
428    }
429
430    fn boundedness(&self) -> Boundedness {
431        Boundedness::Unbounded {
432            requires_infinite_memory: false,
433        }
434    }
435
436    fn pipeline_behavior(&self) -> EmissionType {
437        EmissionType::Incremental
438    }
439
440    fn equivalence_properties(&self) -> &EquivalenceProperties {
441        self.properties.equivalence_properties()
442    }
443}
444
445// ── Probe Logic ──────────────────────────────────────────────────
446
447/// Probes the hash index for each row in `stream_batch` and builds
448/// the joined output batch.
449#[allow(clippy::too_many_arguments)]
450fn probe_batch(
451    stream_batch: &RecordBatch,
452    converter: &RowConverter,
453    index: &HashIndex,
454    lookup_batch: &RecordBatch,
455    stream_key_indices: &[usize],
456    join_type: LookupJoinType,
457    output_schema: &SchemaRef,
458    stream_field_count: usize,
459    projection: &[usize],
460) -> Result<RecordBatch> {
461    let key_cols: Vec<_> = stream_key_indices
462        .iter()
463        .map(|&i| stream_batch.column(i).clone())
464        .collect();
465    let rows = converter.convert_columns(&key_cols)?;
466
467    let num_rows = stream_batch.num_rows();
468    let mut stream_indices: Vec<u32> = Vec::with_capacity(num_rows);
469    let mut lookup_indices: Vec<Option<u32>> = Vec::with_capacity(num_rows);
470
471    #[allow(clippy::cast_possible_truncation)] // batch row count fits u32
472    for row in 0..num_rows {
473        // SQL semantics: NULL != NULL, so rows with any null key never match.
474        if key_cols.iter().any(|c| c.is_null(row)) {
475            if join_type == LookupJoinType::LeftOuter {
476                stream_indices.push(row as u32);
477                lookup_indices.push(None);
478            }
479            continue;
480        }
481
482        let key = rows.row(row);
483        match index.probe(key.as_ref()) {
484            Some(matches) => {
485                for &lookup_row in matches {
486                    stream_indices.push(row as u32);
487                    lookup_indices.push(Some(lookup_row));
488                }
489            }
490            None if join_type == LookupJoinType::LeftOuter => {
491                stream_indices.push(row as u32);
492                lookup_indices.push(None);
493            }
494            None => {}
495        }
496    }
497
498    if stream_indices.is_empty() {
499        return Ok(RecordBatch::new_empty(Arc::clone(output_schema)));
500    }
501
502    // Gather stream-side columns
503    let take_stream = UInt32Array::from(stream_indices);
504    let mut columns = Vec::with_capacity(stream_field_count + lookup_batch.num_columns());
505
506    for col in stream_batch.columns() {
507        columns.push(take(col.as_ref(), &take_stream, None)?);
508    }
509
510    // Gather lookup-side columns (None → null in output)
511    let take_lookup: UInt32Array = lookup_indices.into_iter().collect();
512    for col in lookup_batch.columns() {
513        columns.push(take(col.as_ref(), &take_lookup, None)?);
514    }
515
516    debug_assert_eq!(
517        columns.len(),
518        stream_field_count + lookup_batch.num_columns(),
519        "output column count mismatch"
520    );
521
522    let projected_columns: Vec<_> = projection.iter().map(|&idx| columns[idx].clone()).collect();
523
524    debug_assert_eq!(
525        projected_columns.len(),
526        output_schema.fields().len(),
527        "projected column count mismatch"
528    );
529
530    Ok(RecordBatch::try_new(
531        Arc::clone(output_schema),
532        projected_columns,
533    )?)
534}
535
536fn resolve_physical_projection(
537    logical_schema: &datafusion::common::DFSchema,
538    stream_schema: &arrow::datatypes::Schema,
539    lookup_schema: &arrow::datatypes::Schema,
540    lookup_table: &str,
541    lookup_alias: Option<&str>,
542) -> Result<Vec<usize>> {
543    let mut projection = Vec::with_capacity(logical_schema.fields().len());
544    for idx in 0..logical_schema.fields().len() {
545        let (qualifier, field) = logical_schema.qualified_field(idx);
546        let name = field.name();
547        let is_lookup = if let Some(relation) = qualifier {
548            let table = relation.table();
549            table == lookup_table || Some(table) == lookup_alias
550        } else {
551            // No qualifier: fallback to name-based detection.
552            if stream_schema.index_of(name).is_ok() && lookup_schema.index_of(name).is_err() {
553                false
554            } else if lookup_schema.index_of(name).is_ok() && stream_schema.index_of(name).is_err()
555            {
556                true
557            } else {
558                stream_schema.index_of(name).is_ok()
559            }
560        };
561
562        if is_lookup {
563            let lookup_idx = lookup_schema.index_of(name).map_err(|_| {
564                DataFusionError::Plan(format!(
565                    "lookup join projection: output field '{name}' not found in lookup \
566                     schema {lookup_schema:?}"
567                ))
568            })?;
569            projection.push(stream_schema.fields().len() + lookup_idx);
570        } else {
571            let stream_idx = stream_schema.index_of(name).map_err(|_| {
572                DataFusionError::Plan(format!(
573                    "lookup join projection: output field '{name}' not found in stream \
574                     schema {stream_schema:?}"
575                ))
576            })?;
577            projection.push(stream_idx);
578        }
579    }
580    Ok(projection)
581}
582
583// ── Extension Planner ────────────────────────────────────────────
584
585/// Converts `LookupJoinNode` logical plans to [`LookupJoinExec`]
586/// or [`PartialLookupJoinExec`] physical plans by resolving table
587/// data from the registry.
588pub struct LookupJoinExtensionPlanner {
589    registry: Arc<LookupTableRegistry>,
590}
591
592impl LookupJoinExtensionPlanner {
593    /// Creates a planner backed by the given registry.
594    pub fn new(registry: Arc<LookupTableRegistry>) -> Self {
595        Self { registry }
596    }
597}
598
599#[async_trait]
600impl ExtensionPlanner for LookupJoinExtensionPlanner {
601    #[allow(clippy::too_many_lines)]
602    async fn plan_extension(
603        &self,
604        _planner: &dyn PhysicalPlanner,
605        node: &dyn UserDefinedLogicalNode,
606        _logical_inputs: &[&LogicalPlan],
607        physical_inputs: &[Arc<dyn ExecutionPlan>],
608        session_state: &SessionState,
609    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
610        let Some(lookup_node) = node.as_any().downcast_ref::<LookupJoinNode>() else {
611            return Ok(None);
612        };
613
614        let entry = self
615            .registry
616            .get_entry(lookup_node.lookup_table_name())
617            .ok_or_else(|| {
618                DataFusionError::Plan(format!(
619                    "lookup table '{}' not registered",
620                    lookup_node.lookup_table_name()
621                ))
622            })?;
623
624        let input = Arc::clone(&physical_inputs[0]);
625        let stream_schema = input.schema();
626
627        match entry {
628            RegisteredLookup::Partial(partial_state) => {
629                let stream_key_indices = resolve_stream_keys(lookup_node, &stream_schema)?;
630
631                let mut output_fields = stream_schema.fields().to_vec();
632                output_fields.extend(partial_state.schema.fields().iter().cloned());
633                let output_schema = Arc::new(Schema::new(output_fields));
634
635                let exec = PartialLookupJoinExec::try_new_with_source(
636                    input,
637                    Arc::clone(&partial_state.lookup_cache),
638                    stream_key_indices,
639                    partial_state.key_sort_fields.clone(),
640                    lookup_node.join_type(),
641                    Arc::clone(&partial_state.schema),
642                    output_schema,
643                    partial_state.source.clone(),
644                    Arc::clone(&partial_state.fetch_semaphore),
645                    partial_state.projection.clone(),
646                )?;
647                Ok(Some(Arc::new(exec)))
648            }
649            RegisteredLookup::Snapshot(snapshot) => {
650                let lookup_schema = snapshot.batch.schema();
651                let lookup_key_indices = resolve_lookup_keys(lookup_node, &lookup_schema)?;
652
653                let lookup_batch = if lookup_node.pushdown_predicates().is_empty()
654                    || snapshot.batch.num_rows() == 0
655                {
656                    snapshot.batch.clone()
657                } else {
658                    apply_pushdown_predicates(
659                        &snapshot.batch,
660                        lookup_node.pushdown_predicates(),
661                        session_state,
662                    )?
663                };
664
665                let stream_key_indices = resolve_stream_keys(lookup_node, &stream_schema)?;
666
667                // Validate join key types are compatible
668                for (si, li) in stream_key_indices.iter().zip(&lookup_key_indices) {
669                    let st = stream_schema.field(*si).data_type();
670                    let lt = lookup_schema.field(*li).data_type();
671                    if st != lt {
672                        return Err(DataFusionError::Plan(format!(
673                            "Lookup join key type mismatch: stream '{}' is {st:?} \
674                             but lookup '{}' is {lt:?}",
675                            stream_schema.field(*si).name(),
676                            lookup_schema.field(*li).name(),
677                        )));
678                    }
679                }
680
681                let logical_schema = lookup_node.schema();
682                let physical_projection = resolve_physical_projection(
683                    logical_schema,
684                    &stream_schema,
685                    &lookup_schema,
686                    lookup_node.lookup_table_name(),
687                    lookup_node.lookup_alias(),
688                )?;
689
690                let output_schema = Arc::new(Schema::new(logical_schema.fields().to_vec()));
691
692                let exec = LookupJoinExec::try_new(
693                    input,
694                    lookup_batch,
695                    stream_key_indices,
696                    lookup_key_indices,
697                    lookup_node.join_type(),
698                    output_schema,
699                )?
700                .with_projection(physical_projection);
701
702                Ok(Some(Arc::new(exec)))
703            }
704        }
705    }
706}
707
708/// Evaluates pushdown predicates against the lookup snapshot, returning
709/// only the rows that pass all predicates. This shrinks the hash index.
710fn apply_pushdown_predicates(
711    batch: &RecordBatch,
712    predicates: &[Expr],
713    session_state: &SessionState,
714) -> Result<RecordBatch> {
715    use arrow::compute::filter_record_batch;
716    use datafusion::physical_expr::create_physical_expr;
717
718    let schema = batch.schema();
719    let df_schema = datafusion::common::DFSchema::try_from(schema.as_ref().clone())?;
720
721    let mut mask = None::<arrow_array::BooleanArray>;
722    for pred in predicates {
723        let phys_expr = create_physical_expr(pred, &df_schema, session_state.execution_props())?;
724        let result = phys_expr.evaluate(batch)?;
725        let bool_arr = result
726            .into_array(batch.num_rows())?
727            .as_any()
728            .downcast_ref::<arrow_array::BooleanArray>()
729            .ok_or_else(|| {
730                DataFusionError::Internal("pushdown predicate did not evaluate to boolean".into())
731            })?
732            .clone();
733        mask = Some(match mask {
734            Some(existing) => arrow::compute::and(&existing, &bool_arr)?,
735            None => bool_arr,
736        });
737    }
738
739    match mask {
740        Some(m) => Ok(filter_record_batch(batch, &m)?),
741        None => Ok(batch.clone()),
742    }
743}
744
745fn resolve_stream_keys(node: &LookupJoinNode, schema: &SchemaRef) -> Result<Vec<usize>> {
746    node.join_keys()
747        .iter()
748        .map(|pair| match &pair.stream_expr {
749            Expr::Column(col) => schema.index_of(&col.name).map_err(|_| {
750                DataFusionError::Plan(format!(
751                    "stream key column '{}' not found in physical schema",
752                    col.name
753                ))
754            }),
755            other => Err(DataFusionError::NotImplemented(format!(
756                "lookup join requires column references as stream keys, got: {other}"
757            ))),
758        })
759        .collect()
760}
761
762fn resolve_lookup_keys(node: &LookupJoinNode, schema: &SchemaRef) -> Result<Vec<usize>> {
763    node.join_keys()
764        .iter()
765        .map(|pair| {
766            schema.index_of(&pair.lookup_column).map_err(|_| {
767                DataFusionError::Plan(format!(
768                    "lookup key column '{}' not found in lookup table schema",
769                    pair.lookup_column
770                ))
771            })
772        })
773        .collect()
774}
775
776// ── Tests ────────────────────────────────────────────────────────
777
778#[cfg(test)]
779mod tests;
laminar_sql/datafusion/lookup_join_exec/mod.rs

laminar_sql/datafusion/lookup_join_exec/
mod.rs
