laminar_connectors/lakehouse/delta_reference/
mod.rs

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

laminar_connectors/lakehouse/delta_reference/
mod.rs
