laminar_connectors/lakehouse/
delta_table_provider.rs

1//! Delta Lake table provider integration with `DataFusion`.
2//!
3//! This module provides a thin helper to open a Delta Lake table and
4//! register it as a `TableProvider` in a `SessionContext`.
5//!
6//! # Usage
7//!
8//! ```rust,ignore
9//! use laminar_connectors::lakehouse::delta_table_provider::register_delta_table;
10//! use datafusion::prelude::SessionContext;
11//! use std::collections::HashMap;
12//!
13//! let ctx = SessionContext::new();
14//! register_delta_table(&ctx, "my_table", "/path/to/delta/table", HashMap::new()).await?;
15//!
16//! // Now query it:
17//! let df = ctx.sql("SELECT * FROM my_table").await?;
18//! ```
19
20#[cfg(feature = "delta-lake")]
21use std::collections::HashMap;
22
23#[cfg(feature = "delta-lake")]
24use std::sync::Arc;
25
26#[cfg(feature = "delta-lake")]
27use datafusion::prelude::SessionContext;
28
29#[cfg(feature = "delta-lake")]
30use tracing::info;
31
32#[cfg(feature = "delta-lake")]
33use crate::error::ConnectorError;
34
35/// Opens a Delta Lake table and registers it as a table provider in the
36/// given `DataFusion` `SessionContext`.
37///
38/// # Arguments
39///
40/// * `ctx` - The `DataFusion` session context to register in
41/// * `name` - The SQL table name (e.g., `"trades"`)
42/// * `table_uri` - Path to the Delta Lake table (local, `s3://`, `az://`, `gs://`)
43/// * `storage_options` - Storage credentials and configuration
44///
45/// # Errors
46///
47/// Returns `ConnectorError::ConnectionFailed` if the table cannot be opened,
48/// or `ConnectorError::Internal` if registration fails.
49#[cfg(feature = "delta-lake")]
50#[allow(clippy::implicit_hasher)]
51pub async fn register_delta_table(
52    ctx: &SessionContext,
53    name: &str,
54    table_uri: &str,
55    storage_options: HashMap<String, String>,
56) -> Result<(), ConnectorError> {
57    use super::delta_io;
58
59    info!(
60        name,
61        table_uri, "registering Delta Lake table as TableProvider"
62    );
63
64    // Open the existing table.
65    let table = delta_io::open_or_create_table(table_uri, storage_options, None).await?;
66
67    // Register the table's object store with the session so scans can resolve
68    // non-local URLs (s3://, az://, gs://); without it, reading an
69    // object-store-backed table fails with "No suitable object store found".
70    // (Local-filesystem tables use DataFusion's built-in store.)
71    table
72        .update_datafusion_session(&ctx.state())
73        .map_err(|e| ConnectorError::Internal(format!("register Delta object store: {e}")))?;
74
75    // Build a DeltaTableProvider (which implements TableProvider) from the table.
76    let provider =
77        table.table_provider().build().await.map_err(|e| {
78            ConnectorError::Internal(format!("failed to build table provider: {e}"))
79        })?;
80
81    ctx.register_table(
82        datafusion::common::TableReference::bare(name),
83        Arc::new(provider),
84    )
85    .map_err(|e| {
86        ConnectorError::Internal(format!("failed to register Delta table '{name}': {e}"))
87    })?;
88
89    info!(name, table_uri, "Delta Lake table registered successfully");
90
91    Ok(())
92}
93
94#[cfg(all(test, feature = "delta-lake"))]
95mod tests {
96    use super::*;
97    use arrow_array::{Float64Array, Int64Array, StringArray};
98    use arrow_schema::{DataType, Field, Schema, SchemaRef};
99    use tempfile::TempDir;
100
101    fn test_schema() -> SchemaRef {
102        Arc::new(Schema::new(vec![
103            Field::new("id", DataType::Int64, false),
104            Field::new("name", DataType::Utf8, true),
105            Field::new("value", DataType::Float64, true),
106        ]))
107    }
108
109    #[allow(clippy::cast_precision_loss)]
110    fn test_batch(n: usize) -> arrow_array::RecordBatch {
111        let ids: Vec<i64> = (0..n as i64).collect();
112        let names: Vec<&str> = (0..n).map(|_| "test").collect();
113        let values: Vec<f64> = (0..n).map(|i| i as f64 * 1.5).collect();
114
115        arrow_array::RecordBatch::try_new(
116            test_schema(),
117            vec![
118                Arc::new(Int64Array::from(ids)),
119                Arc::new(StringArray::from(names)),
120                Arc::new(Float64Array::from(values)),
121            ],
122        )
123        .unwrap()
124    }
125
126    #[tokio::test]
127    async fn test_register_and_query_delta_table() {
128        use super::super::delta_io;
129        use deltalake::protocol::SaveMode;
130
131        let temp_dir = TempDir::new().unwrap();
132        let table_path = temp_dir.path().to_str().unwrap();
133
134        // Create a Delta table with some data.
135        let schema = test_schema();
136        let table = delta_io::open_or_create_table(table_path, HashMap::new(), Some(&schema))
137            .await
138            .unwrap();
139
140        let batch = test_batch(10);
141        let (_table, version) = delta_io::write_batches(
142            table,
143            vec![batch],
144            SaveMode::Append,
145            None,
146            false,
147            None,
148            None,
149        )
150        .await
151        .unwrap();
152        assert_eq!(version, 1);
153
154        // Register as TableProvider and query.
155        let ctx = SessionContext::new();
156        register_delta_table(&ctx, "MixedDelta", table_path, HashMap::new())
157            .await
158            .unwrap();
159
160        let df = ctx
161            .sql("SELECT COUNT(*) AS cnt FROM \"MixedDelta\"")
162            .await
163            .unwrap();
164        let results = df.collect().await.unwrap();
165
166        assert_eq!(results.len(), 1);
167        let count = results[0]
168            .column(0)
169            .as_any()
170            .downcast_ref::<Int64Array>()
171            .unwrap()
172            .value(0);
173        assert_eq!(count, 10);
174    }
175}
laminar_connectors/lakehouse/delta_table_provider.rs

laminar_connectors/lakehouse/
delta_table_provider.rs
