laminar_connectors/lakehouse/delta_reference/
mod.rs

1//! Delta Lake startup snapshot source for reference tables.
2
3use std::collections::VecDeque;
4
5use arrow_array::RecordBatch;
6use arrow_schema::SchemaRef;
7use deltalake::DeltaTable;
8use tracing::info;
9
10use crate::config::ConnectorConfig;
11use crate::error::ConnectorError;
12use crate::lakehouse::delta_source_config::DeltaSourceConfig;
13use crate::reference::ReferenceTableSource;
14
15use super::snapshot_schema::{conform_snapshot_batch, validate_snapshot_schema};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18enum Phase {
19    Ready,
20    Draining,
21    Done,
22    Closed,
23}
24
25/// A finite snapshot of one Delta Lake table.
26pub struct DeltaReferenceTableSource {
27    config: DeltaSourceConfig,
28    declared_schema: SchemaRef,
29    phase: Phase,
30    table: Option<DeltaTable>,
31    pending_batches: VecDeque<RecordBatch>,
32}
33
34impl DeltaReferenceTableSource {
35    /// Creates a source from parsed connector configuration and the declared table schema.
36    #[must_use]
37    pub fn from_source_config(config: DeltaSourceConfig, declared_schema: SchemaRef) -> Self {
38        Self {
39            config,
40            declared_schema,
41            phase: Phase::Ready,
42            table: None,
43            pending_batches: VecDeque::new(),
44        }
45    }
46
47    /// Creates a source from SQL connector properties and the declared table schema.
48    ///
49    /// # Errors
50    ///
51    /// Returns an error when the Delta connector configuration is invalid.
52    pub fn from_connector_config(
53        config: &ConnectorConfig,
54        declared_schema: SchemaRef,
55    ) -> Result<Self, ConnectorError> {
56        Ok(Self::from_source_config(
57            DeltaSourceConfig::from_config(config)?,
58            declared_schema,
59        ))
60    }
61
62    async fn open_table(&mut self) -> Result<(), ConnectorError> {
63        use crate::lakehouse::delta_io;
64
65        let (resolved_path, resolved_options) = delta_io::resolve_catalog_options(
66            &self.config.catalog_type,
67            self.config.catalog_database.as_deref(),
68            self.config.catalog_name.as_deref(),
69            self.config.catalog_schema.as_deref(),
70            &self.config.table_path,
71            &self.config.storage_options,
72        )
73        .await?;
74        self.table =
75            Some(delta_io::open_or_create_table(&resolved_path, resolved_options, None).await?);
76        Ok(())
77    }
78
79    async fn load_snapshot(&mut self) -> Result<(), ConnectorError> {
80        use crate::lakehouse::delta_io;
81
82        let version = delta_io::get_latest_version(
83            self.table
84                .as_mut()
85                .ok_or_else(|| ConnectorError::Internal("Delta table is not open".into()))?,
86        )
87        .await?;
88        let batches = delta_io::read_batches_at_version(
89            self.table
90                .as_mut()
91                .ok_or_else(|| ConnectorError::Internal("Delta table is not open".into()))?,
92            version,
93            usize::MAX,
94        )
95        .await?
96        .0;
97
98        if batches.is_empty() {
99            let table_schema = delta_io::get_table_schema(
100                self.table
101                    .as_ref()
102                    .ok_or_else(|| ConnectorError::Internal("Delta table is not open".into()))?,
103            )
104            .map_err(|error| {
105                ConnectorError::ReadError(format!("read Delta snapshot schema: {error}"))
106            })?;
107            validate_snapshot_schema(table_schema.as_ref(), self.declared_schema.as_ref())?;
108        }
109
110        let batches = batches
111            .into_iter()
112            .map(|batch| conform_snapshot_batch(&batch, &self.declared_schema))
113            .collect::<Result<Vec<_>, _>>()?;
114        let rows = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
115        info!(version, rows, "Delta reference snapshot loaded");
116        self.pending_batches = batches.into();
117        Ok(())
118    }
119}
120
121#[async_trait::async_trait]
122impl ReferenceTableSource for DeltaReferenceTableSource {
123    async fn poll_snapshot(&mut self) -> Result<Option<RecordBatch>, ConnectorError> {
124        match self.phase {
125            Phase::Closed => {
126                return Err(ConnectorError::InvalidState {
127                    expected: "open reference snapshot source".into(),
128                    actual: "closed".into(),
129                });
130            }
131            Phase::Done => return Ok(None),
132            Phase::Ready => {
133                self.open_table().await?;
134                self.load_snapshot().await?;
135                self.phase = Phase::Draining;
136            }
137            Phase::Draining => {}
138        }
139
140        if let Some(batch) = self.pending_batches.pop_front() {
141            return Ok(Some(batch));
142        }
143        self.phase = Phase::Done;
144        Ok(None)
145    }
146
147    async fn close(&mut self) -> Result<(), ConnectorError> {
148        self.phase = Phase::Closed;
149        self.pending_batches.clear();
150        self.table = None;
151        Ok(())
152    }
153}
154
155#[cfg(test)]
156mod tests;
laminar_connectors/lakehouse/delta_reference/mod.rs

laminar_connectors/lakehouse/delta_reference/
mod.rs
