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