laminar_connectors/lakehouse/
delta_reference.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 = if self.config.partition_filter.is_some() {
89            self.load_filtered_snapshot(version).await?
90        } else {
91            delta_io::read_batches_at_version(
92                self.table
93                    .as_mut()
94                    .ok_or_else(|| ConnectorError::Internal("Delta table is not open".into()))?,
95                version,
96                usize::MAX,
97            )
98            .await?
99            .0
100        };
101
102        if batches.is_empty() {
103            let table_schema = delta_io::get_table_schema(
104                self.table
105                    .as_ref()
106                    .ok_or_else(|| ConnectorError::Internal("Delta table is not open".into()))?,
107            )
108            .map_err(|error| {
109                ConnectorError::ReadError(format!("read Delta snapshot schema: {error}"))
110            })?;
111            validate_snapshot_schema(table_schema.as_ref(), self.declared_schema.as_ref())?;
112        }
113
114        let batches = batches
115            .into_iter()
116            .map(|batch| conform_snapshot_batch(&batch, &self.declared_schema))
117            .collect::<Result<Vec<_>, _>>()?;
118        let rows = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
119        info!(version, rows, "Delta reference snapshot loaded");
120        self.pending_batches = batches.into();
121        Ok(())
122    }
123
124    async fn load_filtered_snapshot(
125        &mut self,
126        version: i64,
127    ) -> Result<Vec<RecordBatch>, ConnectorError> {
128        use tokio_stream::StreamExt;
129
130        let table = self
131            .table
132            .as_mut()
133            .ok_or_else(|| ConnectorError::Internal("Delta table is not open".into()))?;
134        table
135            .load_version(version)
136            .await
137            .map_err(|error| ConnectorError::ReadError(format!("load Delta version: {error}")))?;
138        let provider = table
139            .table_provider()
140            .build()
141            .await
142            .map_err(|error| ConnectorError::ReadError(format!("build Delta scan: {error}")))?;
143
144        let context = datafusion::prelude::SessionContext::new();
145        context
146            .register_table("delta_reference_scan", std::sync::Arc::new(provider))
147            .map_err(|error| ConnectorError::ReadError(format!("register Delta scan: {error}")))?;
148        let filter = self.config.partition_filter.as_deref().ok_or_else(|| {
149            ConnectorError::Internal("filtered Delta snapshot has no filter".into())
150        })?;
151        let dataframe = context
152            .sql(&format!(
153                "SELECT * FROM delta_reference_scan WHERE {filter}"
154            ))
155            .await
156            .map_err(|error| {
157                ConnectorError::ReadError(format!("plan filtered Delta scan: {error}"))
158            })?;
159        let mut stream = dataframe.execute_stream().await.map_err(|error| {
160            ConnectorError::ReadError(format!("execute filtered Delta scan: {error}"))
161        })?;
162        let mut batches = Vec::new();
163        while let Some(batch) = stream.next().await {
164            let batch = batch.map_err(|error| {
165                ConnectorError::ReadError(format!("read filtered Delta batch: {error}"))
166            })?;
167            if batch.num_rows() != 0 {
168                batches.push(batch);
169            }
170        }
171        Ok(batches)
172    }
173}
174
175#[async_trait::async_trait]
176impl ReferenceTableSource for DeltaReferenceTableSource {
177    async fn poll_snapshot(&mut self) -> Result<Option<RecordBatch>, ConnectorError> {
178        match self.phase {
179            Phase::Closed => {
180                return Err(ConnectorError::InvalidState {
181                    expected: "open reference snapshot source".into(),
182                    actual: "closed".into(),
183                });
184            }
185            Phase::Done => return Ok(None),
186            Phase::Ready => {
187                self.open_table().await?;
188                self.load_snapshot().await?;
189                self.phase = Phase::Draining;
190            }
191            Phase::Draining => {}
192        }
193
194        if let Some(batch) = self.pending_batches.pop_front() {
195            return Ok(Some(batch));
196        }
197        self.phase = Phase::Done;
198        Ok(None)
199    }
200
201    async fn close(&mut self) -> Result<(), ConnectorError> {
202        self.phase = Phase::Closed;
203        self.pending_batches.clear();
204        self.table = None;
205        Ok(())
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use std::sync::Arc;
212
213    use arrow_schema::{DataType, Field, Schema};
214
215    use super::*;
216
217    fn declared_schema() -> SchemaRef {
218        Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]))
219    }
220
221    #[test]
222    fn construction_carries_declared_schema() {
223        let source = DeltaReferenceTableSource::from_source_config(
224            DeltaSourceConfig::new("/tmp/test_delta"),
225            declared_schema(),
226        );
227        assert_eq!(source.declared_schema.field(0).name(), "id");
228        assert!(!source.declared_schema.field(0).is_nullable());
229    }
230
231    #[test]
232    fn missing_table_path_is_rejected() {
233        let config = ConnectorConfig::new("delta-lake");
234        assert!(
235            DeltaReferenceTableSource::from_connector_config(&config, declared_schema()).is_err()
236        );
237    }
238
239    #[tokio::test]
240    async fn close_is_idempotent_and_prevents_reads() {
241        let mut source = DeltaReferenceTableSource::from_source_config(
242            DeltaSourceConfig::new("/tmp/test_delta"),
243            declared_schema(),
244        );
245        source.close().await.unwrap();
246        source.close().await.unwrap();
247        assert!(source.poll_snapshot().await.is_err());
248    }
249
250    mod integration {
251        use std::collections::HashMap;
252
253        use arrow_array::{Int64Array, StringArray};
254        use deltalake::protocol::SaveMode;
255        use tempfile::TempDir;
256
257        use super::*;
258
259        fn physical_schema() -> SchemaRef {
260            Arc::new(Schema::new(vec![
261                Field::new("id", DataType::Int64, true),
262                Field::new("name", DataType::Utf8, true),
263            ]))
264        }
265
266        fn declared_schema() -> SchemaRef {
267            Arc::new(Schema::new(vec![
268                Field::new("id", DataType::Int64, false),
269                Field::new("name", DataType::Utf8, true),
270            ]))
271        }
272
273        async fn write_table(path: &str) {
274            use crate::lakehouse::delta_io;
275
276            let schema = physical_schema();
277            let batch = RecordBatch::try_new(
278                schema.clone(),
279                vec![
280                    Arc::new(Int64Array::from(vec![1, 2])),
281                    Arc::new(StringArray::from(vec!["one", "two"])),
282                ],
283            )
284            .unwrap();
285            let table = delta_io::open_or_create_table(path, HashMap::new(), Some(&schema))
286                .await
287                .unwrap();
288            delta_io::write_batches(
289                table,
290                vec![batch],
291                SaveMode::Append,
292                None,
293                false,
294                None,
295                None,
296            )
297            .await
298            .unwrap();
299        }
300
301        #[tokio::test]
302        async fn snapshot_uses_declared_schema_and_exhausts() {
303            let directory = TempDir::new().unwrap();
304            let path = directory.path().to_str().unwrap();
305            write_table(path).await;
306            let mut source = DeltaReferenceTableSource::from_source_config(
307                DeltaSourceConfig::new(path),
308                declared_schema(),
309            );
310
311            let batch = source.poll_snapshot().await.unwrap().unwrap();
312            assert_eq!(batch.schema(), declared_schema());
313            assert_eq!(batch.num_rows(), 2);
314            assert!(source.poll_snapshot().await.unwrap().is_none());
315            assert!(source.poll_snapshot().await.unwrap().is_none());
316        }
317
318        #[tokio::test]
319        async fn incompatible_declared_schema_fails_closed() {
320            let directory = TempDir::new().unwrap();
321            let path = directory.path().to_str().unwrap();
322            write_table(path).await;
323            let incompatible = Arc::new(Schema::new(vec![
324                Field::new("id", DataType::Utf8, false),
325                Field::new("name", DataType::Utf8, true),
326            ]));
327            let mut source = DeltaReferenceTableSource::from_source_config(
328                DeltaSourceConfig::new(path),
329                incompatible,
330            );
331
332            assert!(source.poll_snapshot().await.is_err());
333        }
334    }
335}
laminar_connectors/lakehouse/delta_reference.rs

laminar_connectors/lakehouse/
delta_reference.rs
