delta_arrow_reader/
datafusion_provider.rs

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

delta_arrow_reader/
datafusion_provider.rs
