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

delta_arrow_reader/reader/
datafusion.rs
