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