Skip to main content

delta_arrow_reader/
datafusion_provider.rs

1//! Optional DataFusion table-provider and registration surface.
2
3use std::{collections::HashSet, fmt, sync::Arc};
4
5use arrow::datatypes::SchemaRef;
6use async_trait::async_trait;
7use datafusion::{
8    catalog::Session,
9    common::{DataFusionError, Result as DataFusionResult},
10    datasource::{TableProvider, TableType},
11    execution::context::SessionContext,
12    logical_expr::{Expr, TableProviderFilterPushDown},
13    physical_plan::ExecutionPlan,
14};
15
16use crate::{
17    DeltaReaderBackend, DeltaReaderError, DeltaReaderExecutionOptions, DeltaTable,
18    datafusion_execution::create_datafusion_execution_plan,
19    datafusion_planning::{
20        DataFusionFilterCapabilities, plan_datafusion_filters, plan_datafusion_scan,
21    },
22    kernel::delta_predicate_to_kernel_pruning,
23    planning::{
24        DeltaScanPartitionTargetOptions, plan_row_predicate, plan_scan, validate_backend_available,
25    },
26};
27
28const TRACING_TARGET: &str = "delta_arrow_reader::datafusion";
29
30/// DataFusion-specific scan settings for one provider.
31#[derive(Debug, Clone, Default)]
32pub struct DeltaDataFusionScanOptions {
33    /// Reader execution settings used by each provider scan.
34    pub execution_options: DeltaReaderExecutionOptions,
35    /// Optional explicit scan partition target.
36    pub target_partitions: Option<usize>,
37}
38
39/// Immutable DataFusion provider for one loaded Delta table snapshot.
40///
41/// ```no_run
42/// use std::sync::Arc;
43/// use datafusion::prelude::SessionContext;
44/// use delta_arrow_reader::{
45///     DeltaDataFusionScanOptions, DeltaTableBuilder, DeltaTableProvider,
46/// };
47///
48/// # fn build_provider() -> Result<(), Box<dyn std::error::Error>> {
49/// let table = DeltaTableBuilder::new("/tmp/example-delta-table").load()?;
50/// let provider = DeltaTableProvider::try_new(
51///     table,
52///     DeltaDataFusionScanOptions::default(),
53/// )?;
54/// SessionContext::new().register_table("orders", Arc::new(provider))?;
55/// # Ok(())
56/// # }
57/// ```
58#[derive(Clone)]
59pub struct DeltaTableProvider {
60    table: DeltaTable,
61    options: DeltaDataFusionScanOptions,
62    source_name: Option<String>,
63}
64
65impl DeltaTableProvider {
66    /// Creates a provider after validating its options and table protocol.
67    pub fn try_new(
68        table: DeltaTable,
69        options: DeltaDataFusionScanOptions,
70    ) -> Result<Self, DeltaReaderError> {
71        Self::try_new_with_source_name(table, options, None)
72    }
73
74    fn try_new_with_source_name(
75        table: DeltaTable,
76        options: DeltaDataFusionScanOptions,
77        source_name: Option<String>,
78    ) -> Result<Self, DeltaReaderError> {
79        options.execution_options.validate()?;
80        validate_backend_available(options.execution_options)?;
81        if options.target_partitions == Some(0) {
82            return Err(DeltaReaderError::InvalidConfiguration {
83                reason: "scan_partition_target_must_be_positive",
84            });
85        }
86        table.validate_protocol()?;
87        Ok(Self {
88            table,
89            options,
90            source_name,
91        })
92    }
93
94    fn plan(
95        &self,
96        state: &dyn Session,
97        projection: Option<&[usize]>,
98        filters: &[Expr],
99    ) -> Result<(Arc<dyn ExecutionPlan>, usize), DeltaReaderError> {
100        let partition_columns = self
101            .table
102            .partition_columns()
103            .iter()
104            .cloned()
105            .collect::<HashSet<_>>();
106        let filter_refs = filters.iter().collect::<Vec<_>>();
107        let planning = plan_datafusion_scan(
108            self.table.schema(),
109            &partition_columns,
110            projection,
111            &filter_refs,
112            DataFusionFilterCapabilities {
113                exact_predicate_evaluation: self.options.execution_options.reader_backend()
114                    == DeltaReaderBackend::NativeAsync,
115            },
116        )?;
117        if planning
118            .filters
119            .decisions
120            .iter()
121            .any(|decision| decision.pushdown == TableProviderFilterPushDown::Unsupported)
122        {
123            return Err(DeltaReaderError::UnsupportedPredicate {
124                reason: "datafusion_scan_contains_unsupported_filter",
125            });
126        }
127        let physical_projection = planning.projection.physical_projection.clone();
128        let hidden_columns = planning.projection.hidden_columns.clone();
129        let kernel_predicate = planning
130            .filters
131            .predicate
132            .as_ref()
133            .map(|predicate| {
134                delta_predicate_to_kernel_pruning(predicate).ok_or(
135                    DeltaReaderError::UnsupportedPredicate {
136                        reason: "datafusion_predicate_not_kernel_safe",
137                    },
138                )
139            })
140            .transpose()?;
141        let row_predicate = planning
142            .filters
143            .row_predicate
144            .as_ref()
145            .map(|predicate| {
146                delta_predicate_to_kernel_pruning(predicate).ok_or(
147                    DeltaReaderError::UnsupportedPredicate {
148                        reason: "exact_row_predicate_not_kernel_safe",
149                    },
150                )
151            })
152            .transpose()?;
153        let row_predicate = plan_row_predicate(
154            self.table.snapshot(),
155            physical_projection.as_deref(),
156            &hidden_columns,
157            row_predicate,
158        )?;
159        let core = plan_scan(
160            self.table.snapshot(),
161            physical_projection.as_deref(),
162            &hidden_columns,
163            kernel_predicate,
164            planning.filters.requires_statistics,
165            self.options.execution_options,
166            DeltaScanPartitionTargetOptions {
167                explicit_target_partitions: self.options.target_partitions,
168                caller_target_partitions: Some(state.config().target_partitions()),
169            },
170        )?;
171        let partition_count = core.partitions.len();
172        Ok((
173            create_datafusion_execution_plan(
174                core,
175                planning,
176                row_predicate,
177                self.source_name.clone(),
178            ),
179            partition_count,
180        ))
181    }
182}
183
184impl fmt::Debug for DeltaTableProvider {
185    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
186        formatter
187            .debug_struct("DeltaTableProvider")
188            .field("snapshot_version", &self.table.version())
189            .finish_non_exhaustive()
190    }
191}
192
193#[async_trait]
194impl TableProvider for DeltaTableProvider {
195    fn schema(&self) -> SchemaRef {
196        Arc::clone(self.table.schema())
197    }
198
199    fn table_type(&self) -> TableType {
200        TableType::Base
201    }
202
203    async fn scan(
204        &self,
205        state: &dyn Session,
206        projection: Option<&Vec<usize>>,
207        filters: &[Expr],
208        _limit: Option<usize>,
209    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
210        match self.plan(state, projection.map(Vec::as_slice), filters) {
211            Ok((plan, partition_count)) => {
212                tracing::debug!(
213                    target: TRACING_TARGET,
214                    event = "provider_scan.planned",
215                    snapshot_version = self.table.version(),
216                    partition_count,
217                    backend = ?self.options.execution_options.reader_backend(),
218                    outcome = "planned"
219                );
220                Ok(plan)
221            }
222            Err(error) => {
223                trace_failure(
224                    "provider_scan.failed",
225                    self.table.version(),
226                    self.options.execution_options.reader_backend(),
227                    &error,
228                );
229                Err(DataFusionError::External(Box::new(error)))
230            }
231        }
232    }
233
234    fn supports_filters_pushdown(
235        &self,
236        filters: &[&Expr],
237    ) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
238        let partition_columns = self
239            .table
240            .partition_columns()
241            .iter()
242            .cloned()
243            .collect::<HashSet<_>>();
244        let planning = plan_datafusion_filters(
245            self.table.schema(),
246            &partition_columns,
247            filters,
248            DataFusionFilterCapabilities {
249                exact_predicate_evaluation: self.options.execution_options.reader_backend()
250                    == DeltaReaderBackend::NativeAsync,
251            },
252        );
253        Ok(planning
254            .decisions
255            .iter()
256            .map(|decision| decision.pushdown.clone())
257            .collect())
258    }
259}
260
261/// Result of registering one loaded Delta table in a DataFusion context.
262#[derive(Debug, Clone, PartialEq, Eq)]
263pub struct RegisteredDeltaTable {
264    /// Caller-supplied DataFusion table name.
265    pub name: String,
266    /// Loaded Delta snapshot version.
267    pub version: u64,
268}
269
270/// Registers one loaded Delta table in a DataFusion session.
271///
272/// Registration performs no scan. Existing registrations are preserved and
273/// reported through [`DeltaReaderError`].
274///
275/// ```no_run
276/// use datafusion::prelude::SessionContext;
277/// use delta_arrow_reader::{
278///     DeltaDataFusionScanOptions, DeltaTableBuilder, register_delta_table,
279/// };
280///
281/// # fn register() -> Result<(), Box<dyn std::error::Error>> {
282/// let context = SessionContext::new();
283/// let table = DeltaTableBuilder::new("/tmp/example-delta-table").load()?;
284/// let registered = register_delta_table(
285///     &context,
286///     "orders",
287///     table,
288///     DeltaDataFusionScanOptions::default(),
289/// )?;
290/// assert_eq!(registered.name, "orders");
291/// # Ok(())
292/// # }
293/// ```
294pub fn register_delta_table(
295    context: &SessionContext,
296    name: impl Into<String>,
297    table: DeltaTable,
298    options: DeltaDataFusionScanOptions,
299) -> Result<RegisteredDeltaTable, DeltaReaderError> {
300    let name = name.into();
301    let version = table.version();
302    let backend = options.execution_options.reader_backend();
303    let result = (|| {
304        validate_registration_name(&name)?;
305        let provider =
306            DeltaTableProvider::try_new_with_source_name(table, options, Some(name.clone()))?;
307        context
308            .register_table(name.as_str(), Arc::new(provider))
309            .map_err(|source| DeltaReaderError::DataFusionAdapter {
310                reason: "table_registration_failed",
311                source: Box::new(source),
312            })?;
313        Ok(RegisteredDeltaTable { name, version })
314    })();
315    match result {
316        Ok(registered) => {
317            tracing::debug!(
318                target: TRACING_TARGET,
319                event = "provider_registration.registered",
320                snapshot_version = version,
321                partition_count = tracing::field::Empty,
322                backend = ?backend,
323                outcome = "registered"
324            );
325            Ok(registered)
326        }
327        Err(error) => {
328            trace_failure("provider_registration.failed", version, backend, &error);
329            Err(error)
330        }
331    }
332}
333
334fn validate_registration_name(name: &str) -> Result<(), DeltaReaderError> {
335    let mut chars = name.chars();
336    let valid = chars
337        .next()
338        .is_some_and(|first| first == '_' || first.is_ascii_alphabetic())
339        && chars.all(|value| value == '_' || value.is_ascii_alphanumeric());
340    if !valid || is_reserved_sql_keyword(name) {
341        let reason = if name.is_empty() {
342            "table_registration_name_empty"
343        } else {
344            "table_registration_name_invalid"
345        };
346        return Err(DeltaReaderError::DataFusionAdapter {
347            reason,
348            source: Box::new(DataFusionError::Plan(reason.to_owned())),
349        });
350    }
351    Ok(())
352}
353
354fn is_reserved_sql_keyword(name: &str) -> bool {
355    const KEYWORDS: &[&str] = &[
356        "all",
357        "alter",
358        "analyze",
359        "and",
360        "anti",
361        "as",
362        "asof",
363        "by",
364        "case",
365        "connect",
366        "cross",
367        "delete",
368        "distinct",
369        "distribute",
370        "drop",
371        "else",
372        "end",
373        "except",
374        "exists",
375        "explain",
376        "false",
377        "fetch",
378        "for",
379        "format",
380        "from",
381        "full",
382        "global",
383        "group",
384        "having",
385        "in",
386        "inner",
387        "insert",
388        "intersect",
389        "into",
390        "is",
391        "join",
392        "lateral",
393        "left",
394        "like",
395        "limit",
396        "minus",
397        "natural",
398        "not",
399        "null",
400        "offset",
401        "on",
402        "open",
403        "or",
404        "order",
405        "outer",
406        "partition",
407        "pivot",
408        "prewhere",
409        "qualify",
410        "returning",
411        "right",
412        "sample",
413        "select",
414        "semi",
415        "set",
416        "settings",
417        "sort",
418        "start",
419        "table",
420        "tablesample",
421        "then",
422        "top",
423        "true",
424        "union",
425        "unpivot",
426        "update",
427        "using",
428        "values",
429        "view",
430        "when",
431        "where",
432        "window",
433        "with",
434    ];
435    KEYWORDS
436        .iter()
437        .any(|keyword| name.eq_ignore_ascii_case(keyword))
438}
439
440fn trace_failure(
441    event: &'static str,
442    snapshot_version: u64,
443    backend: DeltaReaderBackend,
444    error: &DeltaReaderError,
445) {
446    tracing::debug!(
447        target: TRACING_TARGET,
448        event,
449        snapshot_version,
450        partition_count = tracing::field::Empty,
451        backend = ?backend,
452        outcome = "failed",
453        error_variant = error.as_str(),
454        error_phase = error.phase().as_str()
455    );
456}
457
458#[cfg(test)]
459mod tests {
460    use super::validate_registration_name;
461
462    #[test]
463    fn registration_names_preserve_the_frozen_unquoted_identifier_boundary() {
464        for name in ["orders", "_customers", "Regions_2026", "line_items"] {
465            assert!(validate_registration_name(name).is_ok(), "{name}");
466        }
467
468        for name in [
469            "",
470            "2026_orders",
471            "orders.latest",
472            "line-items",
473            "line items",
474            "\"orders\"",
475            "orders$",
476            "ordérs",
477            "select",
478            "FROM",
479            "Join",
480            "where",
481            "table",
482        ] {
483            assert!(validate_registration_name(name).is_err(), "{name}");
484        }
485    }
486}