laminar_core/lookup/align/
mod.rs

1//! Shared key decoding + result realignment for on-demand lookup sources.
2//!
3//! Every [`LookupSource`](crate::lookup::source::LookupSource) has the same
4//! shape around its backend-specific fetch: the operator hands keys as
5//! `RowConverter`-encoded bytes, the source fetches matching rows in arbitrary
6//! order, and each fetched row is matched back to its key by re-encoding its
7//! primary-key columns with the *same* converter. [`KeyAligner`] owns that
8//! converter and performs the decode + realign so each source only writes its
9//! fetch.
10
11use std::sync::Arc;
12
13use arrow::row::{RowConverter, SortField};
14use arrow_array::{ArrayRef, RecordBatch};
15use rustc_hash::FxHashMap;
16
17use crate::lookup::source::LookupError;
18
19/// Decodes opaque lookup keys and realigns fetched rows to the input key order.
20pub struct KeyAligner {
21    converter: RowConverter,
22    pk_columns: Vec<String>,
23}
24
25impl KeyAligner {
26    /// Build an aligner from the primary-key sort fields and column names
27    /// (same order).
28    ///
29    /// # Errors
30    ///
31    /// Returns `LookupError::Internal` if the key is empty or the converter
32    /// cannot be built for the given sort fields.
33    pub fn new(
34        pk_sort_fields: Vec<SortField>,
35        pk_columns: Vec<String>,
36    ) -> Result<Self, LookupError> {
37        if pk_columns.is_empty() {
38            return Err(LookupError::Internal(
39                "primary_key_columns must not be empty".into(),
40            ));
41        }
42        let converter = RowConverter::new(pk_sort_fields)
43            .map_err(|e| LookupError::Internal(format!("row converter: {e}")))?;
44        Ok(Self {
45            converter,
46            pk_columns,
47        })
48    }
49
50    /// The primary-key column names, in key order.
51    #[must_use]
52    pub fn pk_columns(&self) -> &[String] {
53        &self.pk_columns
54    }
55
56    /// Decode opaque key bytes into the primary-key columns (column-major: one
57    /// array per PK column, each `keys.len()` long) for building a source
58    /// filter.
59    ///
60    /// # Errors
61    ///
62    /// Returns `LookupError::Internal` if the bytes cannot be decoded.
63    pub fn decode_keys(&self, keys: &[&[u8]]) -> Result<Vec<ArrayRef>, LookupError> {
64        let parser = self.converter.parser();
65        let parsed = keys.iter().map(|k| parser.parse(k));
66        self.converter
67            .convert_rows(parsed)
68            .map_err(|e| LookupError::Internal(format!("decode keys: {e}")))
69    }
70
71    /// Realign fetched rows to the input key order. Each fetched row is matched
72    /// to its key by re-encoding its PK columns. Duplicate input keys each
73    /// resolve to their own single-row slice; duplicate fetched keys are
74    /// rejected, and misses are `None`.
75    ///
76    /// # Errors
77    ///
78    /// Returns `LookupError::Internal` if a PK column is absent from a fetched
79    /// batch, cannot be re-encoded, or occurs more than once in the fetched rows.
80    pub fn align(
81        &self,
82        keys: &[&[u8]],
83        fetched: &[RecordBatch],
84    ) -> Result<Vec<Option<RecordBatch>>, LookupError> {
85        let mut index: FxHashMap<Vec<u8>, (usize, usize)> = FxHashMap::default();
86        for (batch_idx, batch) in fetched.iter().enumerate() {
87            if batch.num_rows() == 0 {
88                continue;
89            }
90            let pk_cols = self
91                .pk_columns
92                .iter()
93                .map(|name| {
94                    let idx = batch.schema().index_of(name).map_err(|_| {
95                        LookupError::Internal(format!("pk column not found in result: {name}"))
96                    })?;
97                    Ok(Arc::clone(batch.column(idx)))
98                })
99                .collect::<Result<Vec<ArrayRef>, LookupError>>()?;
100            let rows = self
101                .converter
102                .convert_columns(&pk_cols)
103                .map_err(|e| LookupError::Internal(format!("encode result keys: {e}")))?;
104            for row in 0..batch.num_rows() {
105                if index
106                    .insert(rows.row(row).as_ref().to_vec(), (batch_idx, row))
107                    .is_some()
108                {
109                    return Err(LookupError::Internal(
110                        "lookup source returned multiple rows for one key".into(),
111                    ));
112                }
113            }
114        }
115        Ok(keys
116            .iter()
117            .map(|key| index.get(*key).map(|&(bi, row)| fetched[bi].slice(row, 1)))
118            .collect())
119    }
120}
121
122#[cfg(test)]
123mod tests;
laminar_core/lookup/align/mod.rs

laminar_core/lookup/align/
mod.rs
