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