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