laminar_connectors/lakehouse/delta_io/
read.rs

1//! Bounded snapshot and change-data-feed reads.
2
3use super::{
4    debug, from_delta_version, to_delta_version, Arc, ConnectorError, DeltaTable, RecordBatch,
5    SchemaRef,
6};
7
8/// Returns the table's partition columns, or an empty list if the snapshot is
9/// unavailable. Best-effort: used for clustering diagnostics, never for
10/// correctness, so a missing snapshot is not an error.
11#[cfg(feature = "delta-lake")]
12#[must_use]
13pub fn get_partition_columns(table: &DeltaTable) -> Vec<String> {
14    match table.snapshot() {
15        Ok(snapshot) => snapshot.snapshot().metadata().partition_columns().to_vec(),
16        Err(_) => Vec::new(),
17    }
18}
19
20/// Extracts the Arrow schema from a Delta Lake table.
21///
22/// # Arguments
23///
24/// * `table` - The Delta Lake table handle
25///
26/// # Returns
27///
28/// The table's Arrow schema.
29///
30/// # Errors
31///
32/// Returns `ConnectorError::SchemaMismatch` if schema extraction fails.
33#[cfg(feature = "delta-lake")]
34pub fn get_table_schema(table: &DeltaTable) -> Result<SchemaRef, ConnectorError> {
35    let state = table
36        .snapshot()
37        .map_err(|e| ConnectorError::SchemaMismatch(format!("table has no snapshot: {e}")))?;
38
39    // Use the pre-computed Arrow schema from the EagerSnapshot.
40    Ok(state.snapshot().arrow_schema())
41}
42
43/// Returns the latest committed version via the log store.
44///
45/// # Errors
46///
47/// Returns `ConnectorError::ReadError` on failure.
48#[cfg(feature = "delta-lake")]
49pub async fn get_latest_version(table: &mut DeltaTable) -> Result<i64, ConnectorError> {
50    let log_store = table.log_store();
51    let current = table.version().unwrap_or(0);
52    let version = log_store
53        .get_latest_version(current)
54        .await
55        .map_err(|e| ConnectorError::ReadError(format!("failed to get latest version: {e}")))?;
56    from_delta_version(version)
57}
58
59/// Streams the snapshot at `version` through the delta-rs table provider, which applies
60/// deletion vectors and column mapping. Batches are produced on demand, so a caller that
61/// consumes them one at a time never holds the whole snapshot.
62///
63/// # Errors
64///
65/// Returns `ConnectorError::ReadError` if the version cannot be loaded or planned.
66#[cfg(feature = "delta-lake")]
67pub(crate) async fn snapshot_stream(
68    table: &mut DeltaTable,
69    version: i64,
70) -> Result<datafusion::physical_plan::SendableRecordBatchStream, ConnectorError> {
71    use datafusion::prelude::SessionContext;
72
73    let delta_version = to_delta_version(version)?;
74    table
75        .load_version(delta_version)
76        .await
77        .map_err(|e| ConnectorError::ReadError(format!("failed to load version {version}: {e}")))?;
78    debug!(version, "Delta Lake: loaded version for reading");
79
80    let provider =
81        table.table_provider().build().await.map_err(|e| {
82            ConnectorError::ReadError(format!("failed to build table provider: {e}"))
83        })?;
84    let ctx = SessionContext::new();
85    ctx.register_table("delta_source_scan", Arc::new(provider))
86        .map_err(|e| ConnectorError::ReadError(format!("failed to register scan table: {e}")))?;
87    ctx.sql("SELECT * FROM delta_source_scan")
88        .await
89        .map_err(|e| ConnectorError::ReadError(format!("scan query failed: {e}")))?
90        .execute_stream()
91        .await
92        .map_err(|e| ConnectorError::ReadError(format!("stream execution failed: {e}")))
93}
94
95/// Collects up to `max_records` rows of the snapshot at `version`.
96#[cfg(all(feature = "delta-lake", test))]
97pub(crate) async fn read_batches_at_version(
98    table: &mut DeltaTable,
99    version: i64,
100    max_records: usize,
101) -> Result<(Vec<RecordBatch>, bool), ConnectorError> {
102    use tokio_stream::StreamExt;
103
104    let mut stream = snapshot_stream(table, version).await?;
105    let mut batches = Vec::new();
106    let mut total_rows: usize = 0;
107    while let Some(result) = stream.next().await {
108        let batch =
109            result.map_err(|e| ConnectorError::ReadError(format!("stream batch failed: {e}")))?;
110        if batch.num_rows() == 0 {
111            continue;
112        }
113        total_rows += batch.num_rows();
114        batches.push(batch);
115        if total_rows >= max_records {
116            break;
117        }
118    }
119    let fully_consumed = total_rows < max_records || stream.next().await.is_none();
120    Ok((batches, fully_consumed))
121}
122
123/// Maps Delta CDF rows to the canonical signed-weight changelog and drops CDF
124/// metadata columns (`_change_type`, `_commit_version`, `_commit_timestamp`).
125///
126/// # Errors
127///
128/// Returns `ConnectorError::ReadError` when CDF metadata is missing or invalid,
129/// or when the output batch cannot be built.
130#[cfg(feature = "delta-lake")]
131pub fn map_cdf_to_changelog(batch: &RecordBatch) -> Result<RecordBatch, ConnectorError> {
132    use arrow_array::{Array, Int64Array, StringArray};
133    use laminar_core::changelog::WEIGHT_COLUMN;
134
135    let schema = batch.schema();
136    let ct_idx = schema
137        .index_of("_change_type")
138        .map_err(|_| ConnectorError::ReadError("CDF batch is missing _change_type".into()))?;
139    if let Some(field) = schema.fields().iter().find(|field| {
140        ["_op", "__op", WEIGHT_COLUMN]
141            .iter()
142            .any(|reserved| field.name().eq_ignore_ascii_case(reserved))
143    }) {
144        return Err(ConnectorError::ReadError(format!(
145            "CDF batch contains reserved mutation column '{}'",
146            field.name()
147        )));
148    }
149
150    let change_type = batch
151        .column(ct_idx)
152        .as_any()
153        .downcast_ref::<StringArray>()
154        .ok_or_else(|| ConnectorError::ReadError("_change_type is not Utf8".into()))?;
155
156    let mut weights = Vec::with_capacity(batch.num_rows());
157    for row in 0..batch.num_rows() {
158        if change_type.is_null(row) {
159            return Err(ConnectorError::ReadError(format!(
160                "CDF _change_type is null at row {row}"
161            )));
162        }
163        weights.push(match change_type.value(row) {
164            "insert" | "update_postimage" => 1,
165            "delete" | "update_preimage" => -1,
166            value => {
167                return Err(ConnectorError::ReadError(format!(
168                    "unknown CDF _change_type '{value}' at row {row}"
169                )));
170            }
171        });
172    }
173
174    // Rebuild batch: keep user columns, drop CDF metadata, append __weight.
175    let cdf_meta = ["_change_type", "_commit_version", "_commit_timestamp"];
176    let mut fields = Vec::new();
177    let mut columns: Vec<Arc<dyn arrow_array::Array>> = Vec::new();
178    for (i, field) in schema.fields().iter().enumerate() {
179        if !cdf_meta.contains(&field.name().as_str()) {
180            fields.push(field.clone());
181            columns.push(batch.column(i).clone());
182        }
183    }
184    fields.push(Arc::new(arrow_schema::Field::new(
185        WEIGHT_COLUMN,
186        arrow_schema::DataType::Int64,
187        false,
188    )));
189    columns.push(Arc::new(Int64Array::from(weights)));
190
191    RecordBatch::try_new(
192        Arc::new(arrow_schema::Schema::new_with_metadata(
193            fields,
194            schema.metadata().clone(),
195        )),
196        columns,
197    )
198    .map_err(|e| ConnectorError::ReadError(format!("CDF batch rebuild: {e}")))
199}
laminar_connectors/lakehouse/delta_io/read.rs

laminar_connectors/lakehouse/delta_io/
read.rs
