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::{DataType, Schema, SchemaRef};
6use async_trait::async_trait;
7use datafusion::{
8    catalog::Session,
9    common::{DataFusionError, Result as DataFusionResult},
10    datasource::{TableProvider, TableType, physical_plan::wrap_partition_type_in_dict},
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    transform::schema_with_view_types,
27};
28
29const TRACING_TARGET: &str = "delta_arrow_reader::datafusion";
30
31/// DataFusion-specific scan settings for one provider.
32#[derive(Debug, Clone)]
33pub struct DeltaDataFusionScanOptions {
34    /// Reader execution settings used by each provider scan.
35    pub execution_options: DeltaReaderExecutionOptions,
36    /// Optional explicit scan partition target.
37    pub target_partitions: Option<usize>,
38    /// Decode string and binary data-file columns into Arrow view arrays.
39    pub use_view_types: bool,
40}
41
42impl Default for DeltaDataFusionScanOptions {
43    fn default() -> Self {
44        Self {
45            execution_options: DeltaReaderExecutionOptions::default(),
46            target_partitions: None,
47            use_view_types: true,
48        }
49    }
50}
51
52/// Immutable DataFusion provider for one loaded Delta table snapshot.
53///
54/// ```no_run
55/// use std::sync::Arc;
56/// use datafusion::prelude::SessionContext;
57/// use delta_arrow_reader::{
58///     DeltaDataFusionScanOptions, DeltaTableBuilder, DeltaTableProvider,
59/// };
60///
61/// # fn build_provider() -> Result<(), Box<dyn std::error::Error>> {
62/// let table = DeltaTableBuilder::new("/tmp/example-delta-table").load()?;
63/// let provider = DeltaTableProvider::try_new(
64///     table,
65///     DeltaDataFusionScanOptions::default(),
66/// )?;
67/// SessionContext::new().register_table("orders", Arc::new(provider))?;
68/// # Ok(())
69/// # }
70/// ```
71#[derive(Clone)]
72pub struct DeltaTableProvider {
73    table: DeltaTable,
74    schema: SchemaRef,
75    options: DeltaDataFusionScanOptions,
76    source_name: Option<String>,
77}
78
79impl DeltaTableProvider {
80    /// Creates a provider after validating its options and table protocol.
81    pub fn try_new(
82        table: DeltaTable,
83        options: DeltaDataFusionScanOptions,
84    ) -> Result<Self, DeltaReaderError> {
85        Self::try_new_with_source_name(table, options, None)
86    }
87
88    fn try_new_with_source_name(
89        table: DeltaTable,
90        options: DeltaDataFusionScanOptions,
91        source_name: Option<String>,
92    ) -> Result<Self, DeltaReaderError> {
93        options.execution_options.validate()?;
94        validate_backend_available(options.execution_options)?;
95        if options.target_partitions == Some(0) {
96            return Err(DeltaReaderError::InvalidConfiguration {
97                reason: "scan_partition_target_must_be_positive",
98            });
99        }
100        table.validate_protocol()?;
101        let partition_columns = table.partition_columns().iter().cloned().collect();
102        let schema = datafusion_schema(table.schema(), &partition_columns, options.use_view_types);
103        Ok(Self {
104            table,
105            schema,
106            options,
107            source_name,
108        })
109    }
110
111    fn plan(
112        &self,
113        state: &dyn Session,
114        projection: Option<&[usize]>,
115        filters: &[Expr],
116    ) -> Result<(Arc<dyn ExecutionPlan>, usize), DeltaReaderError> {
117        let _planning = tracing::debug_span!(
118            target: "delta_arrow_reader::profile",
119            "Delta scan planning"
120        )
121        .entered();
122        let partition_columns = self
123            .table
124            .partition_columns()
125            .iter()
126            .cloned()
127            .collect::<HashSet<_>>();
128        let filter_refs = filters.iter().collect::<Vec<_>>();
129        let mut planning = plan_datafusion_scan(
130            self.table.schema(),
131            &partition_columns,
132            projection,
133            &filter_refs,
134            DataFusionFilterCapabilities {
135                exact_predicate_evaluation: self.options.execution_options.reader_backend()
136                    == DeltaReaderBackend::NativeAsync,
137            },
138        )?;
139        if planning
140            .filters
141            .decisions
142            .iter()
143            .any(|decision| decision.pushdown == TableProviderFilterPushDown::Unsupported)
144        {
145            return Err(DeltaReaderError::UnsupportedPredicate {
146                reason: "datafusion_scan_contains_unsupported_filter",
147            });
148        }
149        let physical_projection = planning.projection.physical_projection.clone();
150        let hidden_columns = planning.projection.hidden_columns.clone();
151        let kernel_predicate = planning
152            .filters
153            .predicate
154            .as_ref()
155            .map(|predicate| {
156                delta_predicate_to_kernel_pruning(predicate).ok_or(
157                    DeltaReaderError::UnsupportedPredicate {
158                        reason: "datafusion_predicate_not_kernel_safe",
159                    },
160                )
161            })
162            .transpose()?;
163        let row_predicate = planning
164            .filters
165            .row_predicate
166            .as_ref()
167            .map(|predicate| {
168                delta_predicate_to_kernel_pruning(predicate).ok_or(
169                    DeltaReaderError::UnsupportedPredicate {
170                        reason: "exact_row_predicate_not_kernel_safe",
171                    },
172                )
173            })
174            .transpose()?;
175        let row_predicate = plan_row_predicate(
176            self.table.snapshot(),
177            physical_projection.as_deref(),
178            &hidden_columns,
179            row_predicate,
180        )?;
181        let mut core = plan_scan(
182            self.table.snapshot(),
183            physical_projection.as_deref(),
184            &hidden_columns,
185            kernel_predicate,
186            planning.filters.requires_statistics,
187            self.options.execution_options,
188            DeltaScanPartitionTargetOptions {
189                explicit_target_partitions: self.options.target_partitions,
190                caller_target_partitions: Some(state.config().target_partitions()),
191            },
192        )?;
193        core.logical_schema = datafusion_schema(
194            &core.logical_schema,
195            &partition_columns,
196            self.options.use_view_types,
197        );
198        core.physical_schema = datafusion_schema(
199            &core.physical_schema,
200            &partition_columns,
201            self.options.use_view_types,
202        );
203        core.projected_schema = datafusion_schema(
204            &core.projected_schema,
205            &partition_columns,
206            self.options.use_view_types,
207        );
208        planning.projection.output_schema = datafusion_schema(
209            &planning.projection.output_schema,
210            &partition_columns,
211            self.options.use_view_types,
212        );
213        let partition_count = core.partitions.len();
214        let plan = {
215            let _setup = tracing::debug_span!(
216                target: "delta_arrow_reader::profile",
217                "Delta scan execution setup"
218            )
219            .entered();
220            create_datafusion_execution_plan(
221                core,
222                planning,
223                row_predicate,
224                self.source_name.clone(),
225                self.options.use_view_types,
226            )
227        };
228        Ok((plan, partition_count))
229    }
230}
231
232fn datafusion_schema(
233    schema: &Schema,
234    partition_columns: &HashSet<String>,
235    use_view_types: bool,
236) -> SchemaRef {
237    let view_schema = schema_with_view_types(schema);
238    Arc::new(Schema::new_with_metadata(
239        schema
240            .fields()
241            .iter()
242            .zip(view_schema.fields())
243            .map(|(logical, view)| {
244                if partition_columns.contains(logical.name())
245                    && matches!(
246                        logical.data_type(),
247                        DataType::Utf8
248                            | DataType::LargeUtf8
249                            | DataType::Binary
250                            | DataType::LargeBinary
251                    )
252                {
253                    Arc::new(
254                        logical
255                            .as_ref()
256                            .clone()
257                            .with_data_type(wrap_partition_type_in_dict(
258                                logical.data_type().clone(),
259                            )),
260                    )
261                } else if use_view_types {
262                    Arc::clone(view)
263                } else {
264                    Arc::clone(logical)
265                }
266            })
267            .collect::<Vec<_>>(),
268        schema.metadata().clone(),
269    ))
270}
271
272impl fmt::Debug for DeltaTableProvider {
273    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
274        formatter
275            .debug_struct("DeltaTableProvider")
276            .field("snapshot_version", &self.table.version())
277            .finish_non_exhaustive()
278    }
279}
280
281#[async_trait]
282impl TableProvider for DeltaTableProvider {
283    fn schema(&self) -> SchemaRef {
284        Arc::clone(&self.schema)
285    }
286
287    fn table_type(&self) -> TableType {
288        TableType::Base
289    }
290
291    async fn scan(
292        &self,
293        state: &dyn Session,
294        projection: Option<&Vec<usize>>,
295        filters: &[Expr],
296        _limit: Option<usize>,
297    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
298        match self.plan(state, projection.map(Vec::as_slice), filters) {
299            Ok((plan, partition_count)) => {
300                tracing::debug!(
301                    target: TRACING_TARGET,
302                    event = "provider_scan.planned",
303                    snapshot_version = self.table.version(),
304                    partition_count,
305                    backend = ?self.options.execution_options.reader_backend(),
306                    outcome = "planned"
307                );
308                Ok(plan)
309            }
310            Err(error) => {
311                trace_failure(
312                    "provider_scan.failed",
313                    self.table.version(),
314                    self.options.execution_options.reader_backend(),
315                    &error,
316                );
317                Err(DataFusionError::External(Box::new(error)))
318            }
319        }
320    }
321
322    fn supports_filters_pushdown(
323        &self,
324        filters: &[&Expr],
325    ) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
326        let partition_columns = self
327            .table
328            .partition_columns()
329            .iter()
330            .cloned()
331            .collect::<HashSet<_>>();
332        let planning = plan_datafusion_filters(
333            self.table.schema(),
334            &partition_columns,
335            filters,
336            DataFusionFilterCapabilities {
337                exact_predicate_evaluation: self.options.execution_options.reader_backend()
338                    == DeltaReaderBackend::NativeAsync,
339            },
340        );
341        Ok(planning
342            .decisions
343            .iter()
344            .map(|decision| decision.pushdown.clone())
345            .collect())
346    }
347}
348
349/// Result of registering one loaded Delta table in a DataFusion context.
350#[derive(Debug, Clone, PartialEq, Eq)]
351pub struct RegisteredDeltaTable {
352    /// Caller-supplied DataFusion table name.
353    pub name: String,
354    /// Loaded Delta snapshot version.
355    pub version: u64,
356}
357
358/// Registers one loaded Delta table in a DataFusion session.
359///
360/// Registration performs no scan. Existing registrations are preserved and
361/// reported through [`DeltaReaderError`].
362///
363/// ```no_run
364/// use datafusion::prelude::SessionContext;
365/// use delta_arrow_reader::{
366///     DeltaDataFusionScanOptions, DeltaTableBuilder, register_delta_table,
367/// };
368///
369/// # fn register() -> Result<(), Box<dyn std::error::Error>> {
370/// let context = SessionContext::new();
371/// let table = DeltaTableBuilder::new("/tmp/example-delta-table").load()?;
372/// let registered = register_delta_table(
373///     &context,
374///     "orders",
375///     table,
376///     DeltaDataFusionScanOptions::default(),
377/// )?;
378/// assert_eq!(registered.name, "orders");
379/// # Ok(())
380/// # }
381/// ```
382pub fn register_delta_table(
383    context: &SessionContext,
384    name: impl Into<String>,
385    table: DeltaTable,
386    options: DeltaDataFusionScanOptions,
387) -> Result<RegisteredDeltaTable, DeltaReaderError> {
388    let name = name.into();
389    let version = table.version();
390    let backend = options.execution_options.reader_backend();
391    let result = (|| {
392        validate_registration_name(&name)?;
393        let provider =
394            DeltaTableProvider::try_new_with_source_name(table, options, Some(name.clone()))?;
395        context
396            .register_table(name.as_str(), Arc::new(provider))
397            .map_err(|source| DeltaReaderError::DataFusionAdapter {
398                reason: "table_registration_failed",
399                source: Box::new(source),
400            })?;
401        Ok(RegisteredDeltaTable { name, version })
402    })();
403    match result {
404        Ok(registered) => {
405            tracing::debug!(
406                target: TRACING_TARGET,
407                event = "provider_registration.registered",
408                snapshot_version = version,
409                partition_count = tracing::field::Empty,
410                backend = ?backend,
411                outcome = "registered"
412            );
413            Ok(registered)
414        }
415        Err(error) => {
416            trace_failure("provider_registration.failed", version, backend, &error);
417            Err(error)
418        }
419    }
420}
421
422fn validate_registration_name(name: &str) -> Result<(), DeltaReaderError> {
423    let mut chars = name.chars();
424    let valid = chars
425        .next()
426        .is_some_and(|first| first == '_' || first.is_ascii_alphabetic())
427        && chars.all(|value| value == '_' || value.is_ascii_alphanumeric());
428    if !valid || is_reserved_sql_keyword(name) {
429        let reason = if name.is_empty() {
430            "table_registration_name_empty"
431        } else {
432            "table_registration_name_invalid"
433        };
434        return Err(DeltaReaderError::DataFusionAdapter {
435            reason,
436            source: Box::new(DataFusionError::Plan(reason.to_owned())),
437        });
438    }
439    Ok(())
440}
441
442fn is_reserved_sql_keyword(name: &str) -> bool {
443    const KEYWORDS: &[&str] = &[
444        "all",
445        "alter",
446        "analyze",
447        "and",
448        "anti",
449        "as",
450        "asof",
451        "by",
452        "case",
453        "connect",
454        "cross",
455        "delete",
456        "distinct",
457        "distribute",
458        "drop",
459        "else",
460        "end",
461        "except",
462        "exists",
463        "explain",
464        "false",
465        "fetch",
466        "for",
467        "format",
468        "from",
469        "full",
470        "global",
471        "group",
472        "having",
473        "in",
474        "inner",
475        "insert",
476        "intersect",
477        "into",
478        "is",
479        "join",
480        "lateral",
481        "left",
482        "like",
483        "limit",
484        "minus",
485        "natural",
486        "not",
487        "null",
488        "offset",
489        "on",
490        "open",
491        "or",
492        "order",
493        "outer",
494        "partition",
495        "pivot",
496        "prewhere",
497        "qualify",
498        "returning",
499        "right",
500        "sample",
501        "select",
502        "semi",
503        "set",
504        "settings",
505        "sort",
506        "start",
507        "table",
508        "tablesample",
509        "then",
510        "top",
511        "true",
512        "union",
513        "unpivot",
514        "update",
515        "using",
516        "values",
517        "view",
518        "when",
519        "where",
520        "window",
521        "with",
522    ];
523    KEYWORDS
524        .iter()
525        .any(|keyword| name.eq_ignore_ascii_case(keyword))
526}
527
528fn trace_failure(
529    event: &'static str,
530    snapshot_version: u64,
531    backend: DeltaReaderBackend,
532    error: &DeltaReaderError,
533) {
534    tracing::debug!(
535        target: TRACING_TARGET,
536        event,
537        snapshot_version,
538        partition_count = tracing::field::Empty,
539        backend = ?backend,
540        outcome = "failed",
541        error_variant = error.as_str(),
542        error_phase = error.phase().as_str()
543    );
544}
545
546#[cfg(test)]
547mod tests {
548    use std::collections::{HashMap, HashSet};
549
550    use arrow::datatypes::{DataType, Field, Schema};
551
552    use super::{datafusion_schema, validate_registration_name};
553
554    #[test]
555    fn registration_names_preserve_the_frozen_unquoted_identifier_boundary() {
556        for name in ["orders", "_customers", "Regions_2026", "line_items"] {
557            assert!(validate_registration_name(name).is_ok(), "{name}");
558        }
559
560        for name in [
561            "",
562            "2026_orders",
563            "orders.latest",
564            "line-items",
565            "line items",
566            "\"orders\"",
567            "orders$",
568            "ordérs",
569            "select",
570            "FROM",
571            "Join",
572            "where",
573            "table",
574        ] {
575            assert!(validate_registration_name(name).is_err(), "{name}");
576        }
577    }
578
579    #[test]
580    fn datafusion_schema_uses_views_except_for_dictionary_partitions() {
581        let field_metadata = HashMap::from([("field-key".to_owned(), "field-value".to_owned())]);
582        let schema_metadata = HashMap::from([("schema-key".to_owned(), "schema-value".to_owned())]);
583        let schema = Schema::new_with_metadata(
584            vec![
585                Field::new("text", DataType::Utf8, true).with_metadata(field_metadata.clone()),
586                Field::new("payload", DataType::Binary, true),
587                Field::new("region", DataType::Utf8, true),
588                Field::new("partition_payload", DataType::LargeBinary, true),
589                Field::new("id", DataType::Int32, false),
590            ],
591            schema_metadata.clone(),
592        );
593        let partitions = HashSet::from(["region".to_owned(), "partition_payload".to_owned()]);
594
595        let mapped = datafusion_schema(&schema, &partitions, true);
596
597        assert_eq!(
598            mapped.as_ref(),
599            &Schema::new_with_metadata(
600                vec![
601                    Field::new("text", DataType::Utf8View, true).with_metadata(field_metadata),
602                    Field::new("payload", DataType::BinaryView, true),
603                    Field::new(
604                        "region",
605                        DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
606                        true,
607                    ),
608                    Field::new(
609                        "partition_payload",
610                        DataType::Dictionary(
611                            Box::new(DataType::UInt16),
612                            Box::new(DataType::LargeBinary),
613                        ),
614                        true,
615                    ),
616                    Field::new("id", DataType::Int32, false),
617                ],
618                schema_metadata.clone(),
619            )
620        );
621
622        let standard = datafusion_schema(&schema, &partitions, false);
623        assert_eq!(standard.field(0).data_type(), &DataType::Utf8);
624        assert_eq!(standard.field(1).data_type(), &DataType::Binary);
625        assert_eq!(
626            standard.field(2).data_type(),
627            &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8))
628        );
629        assert_eq!(
630            standard.field(3).data_type(),
631            &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::LargeBinary),)
632        );
633        assert_eq!(standard.field(4).data_type(), &DataType::Int32);
634        assert_eq!(standard.metadata(), &schema_metadata);
635    }
636}
delta_arrow_reader/datafusion_provider.rs

delta_arrow_reader/
datafusion_provider.rs
