Skip to main content

delta_arrow_reader/reader/
datafusion.rs

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