Skip to main content

delta_arrow_reader/reader/datafusion/
execution.rs

1//! Optional DataFusion physical execution adapter.
2//!
3//! # Intra-file repartitioning
4//!
5//! Scan planning first balances whole Delta data files across the resolved
6//! partition target. By default, DataFusion's `FileGroupPartitioner` only
7//! receives those groups when they do not fill the target. The provider can
8//! instead allow repartitioning at any partition count. In either mode, the
9//! partitioner flattens the groups, divides their total bytes across the target,
10//! and returns new groups containing whole files or byte ranges. Its
11//! `repartition_file_min_size` setting is the minimum total input size needed
12//! to attempt this operation, not the size of each generated range.
13//!
14//! Each returned range is stored on its `DeltaScanFileTask`. During direct
15//! Parquet execution, the range containing a row group's first column chunk
16//! page offset owns that complete row group. Range ownership and
17//! footer-statistics pruning are intersected before the selected row groups
18//! are passed to the Parquet reader, so a byte boundary never splits a row
19//! group.
20
21use std::{
22    collections::HashSet,
23    fmt,
24    sync::{
25        Arc,
26        atomic::{AtomicU64, Ordering},
27    },
28};
29
30use arrow::{datatypes::SchemaRef, record_batch::RecordBatch};
31use datafusion::{
32    common::{DataFusionError, Result as DataFusionResult, config::ConfigOptions},
33    datasource::{
34        listing::{FileRange, PartitionedFile},
35        physical_plan::{FileGroup, FileGroupPartitioner},
36    },
37    execution::TaskContext,
38    physical_expr::EquivalenceProperties,
39    physical_plan::{
40        DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
41        SendableRecordBatchStream,
42        execution_plan::{Boundedness, EmissionType, SchedulingType},
43        filter_pushdown::{
44            ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, PushedDown,
45        },
46        stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter},
47    },
48};
49use futures_util::{StreamExt, stream};
50
51use super::{
52    dynamic_filters::{DynamicFilterClassification, DynamicFilterDecision, RetainedDynamicFilter},
53    dynamic_partition_pruning::{
54        DynamicPartitionKeepReason, DynamicPartitionPruningDecision,
55        evaluate_dynamic_partition_filter,
56    },
57    planning::DataFusionScanPlan,
58};
59
60use crate::reader::backend::direct_parquet::{
61    ParquetMetadataCache, ParquetRangeReadEstimator, direct_parquet_file_executor,
62};
63use crate::{
64    DeltaReaderError, DeltaScanMetrics, DeltaScanMetricsSnapshot, ParquetReaderBackend,
65    delta::kernel::DeltaKernelPredicate,
66    reader::{
67        delta_kernel_executor,
68        metrics::saturating_fetch_add,
69        planning::{DeltaScanFileTask, DeltaScanPartition, DeltaScanPlan},
70        scheduling::{
71            DeltaScanScheduler, FileAdmissionDecision, FileAdmissionPolicy, ScanReadLimiter,
72        },
73    },
74};
75
76/// Controls when DataFusion may split direct Parquet reads into ranged scan tasks.
77#[non_exhaustive]
78#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
79pub enum IntraFileRepartitioning {
80    /// Allow splitting only below the target partition count.
81    #[default]
82    WhenBelowTarget,
83    /// Allow splitting at any partition count.
84    Always,
85}
86
87impl IntraFileRepartitioning {
88    fn allows_repartitioning(self, current_partitions: usize, target_partitions: usize) -> bool {
89        match self {
90            Self::WhenBelowTarget => current_partitions < target_partitions,
91            Self::Always => true,
92        }
93    }
94}
95
96/// Immutable point-in-time copy of one DataFusion scan's metrics.
97///
98/// A snapshot taken while the scan is running may contain partial progress. It does not change
99/// after creation; call [`ScanMetrics::snapshot`] again to observe later progress.
100#[non_exhaustive]
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct ScanMetricsSnapshot {
103    /// Delta reader planning and execution metrics.
104    pub reader_metrics: DeltaScanMetricsSnapshot,
105    /// Whether the provider requested Arrow view arrays for string and binary data columns.
106    pub uses_arrow_view_types: bool,
107    /// Configured DataFusion batch row target observed at execution.
108    pub configured_batch_size_rows: Option<u64>,
109    /// File tasks pruned before admission by a dynamic partition filter.
110    /// A task is either a whole physical file or one independently read file range.
111    pub dynamic_partition_tasks_pruned: u64,
112    /// File tasks kept after consulting retained dynamic partition filters.
113    /// A task is either a whole physical file or one independently read file range.
114    pub dynamic_partition_tasks_kept: u64,
115    /// Physical filters offered to the post-optimization hook.
116    pub dynamic_filters_received: u64,
117    /// Offered filters retained for dynamic partition pruning.
118    pub dynamic_filters_accepted: u64,
119    /// Offered filters rejected by the dynamic partition policy.
120    pub dynamic_filters_rejected: u64,
121    /// Dynamic partition filters checked against file tasks during admission.
122    pub dynamic_partition_filter_checks: u64,
123    /// Kept file tasks with missing, invalid, or unparsable partition metadata.
124    pub dynamic_partition_tasks_kept_unusable_metadata: u64,
125    /// Kept file tasks whose dynamic filter was unavailable or unevaluable.
126    pub dynamic_partition_tasks_kept_unevaluable_filter: u64,
127}
128
129/// Shared live metrics for one DataFusion physical scan plan.
130#[derive(Clone)]
131pub struct ScanMetrics {
132    inner: Arc<MetricsInner>,
133}
134
135struct MetricsInner {
136    registration_name: Option<String>,
137    reader_metrics: DeltaScanMetrics,
138    uses_arrow_view_types: bool,
139    configured_batch_size_rows: AtomicU64,
140    dynamic_partition_tasks_pruned: AtomicU64,
141    dynamic_partition_tasks_kept: AtomicU64,
142    dynamic_filters_received: AtomicU64,
143    dynamic_filters_accepted: AtomicU64,
144    dynamic_filters_rejected: AtomicU64,
145    dynamic_partition_filter_checks: AtomicU64,
146    dynamic_partition_tasks_kept_unusable_metadata: AtomicU64,
147    dynamic_partition_tasks_kept_unevaluable_filter: AtomicU64,
148}
149
150impl ScanMetrics {
151    pub(super) fn new(
152        registration_name: Option<String>,
153        reader_metrics: DeltaScanMetrics,
154        use_arrow_view_types: bool,
155    ) -> Self {
156        Self {
157            inner: Arc::new(MetricsInner {
158                registration_name,
159                reader_metrics,
160                uses_arrow_view_types: use_arrow_view_types,
161                configured_batch_size_rows: AtomicU64::new(0),
162                dynamic_partition_tasks_pruned: AtomicU64::new(0),
163                dynamic_partition_tasks_kept: AtomicU64::new(0),
164                dynamic_filters_received: AtomicU64::new(0),
165                dynamic_filters_accepted: AtomicU64::new(0),
166                dynamic_filters_rejected: AtomicU64::new(0),
167                dynamic_partition_filter_checks: AtomicU64::new(0),
168                dynamic_partition_tasks_kept_unusable_metadata: AtomicU64::new(0),
169                dynamic_partition_tasks_kept_unevaluable_filter: AtomicU64::new(0),
170            }),
171        }
172    }
173
174    /// Returns the optional registration label supplied by the DataFusion provider.
175    pub fn registration_name(&self) -> Option<&str> {
176        self.inner.registration_name.as_deref()
177    }
178
179    /// Returns an immutable point-in-time copy of all DataFusion scan metrics.
180    pub fn snapshot(&self) -> ScanMetricsSnapshot {
181        let inner = self.inner.as_ref();
182        ScanMetricsSnapshot {
183            reader_metrics: inner.reader_metrics.snapshot(),
184            uses_arrow_view_types: inner.uses_arrow_view_types,
185            configured_batch_size_rows: nonzero_load(&inner.configured_batch_size_rows),
186            dynamic_partition_tasks_pruned: load(&inner.dynamic_partition_tasks_pruned),
187            dynamic_partition_tasks_kept: load(&inner.dynamic_partition_tasks_kept),
188            dynamic_filters_received: load(&inner.dynamic_filters_received),
189            dynamic_filters_accepted: load(&inner.dynamic_filters_accepted),
190            dynamic_filters_rejected: load(&inner.dynamic_filters_rejected),
191            dynamic_partition_filter_checks: load(&inner.dynamic_partition_filter_checks),
192            dynamic_partition_tasks_kept_unusable_metadata: load(
193                &inner.dynamic_partition_tasks_kept_unusable_metadata,
194            ),
195            dynamic_partition_tasks_kept_unevaluable_filter: load(
196                &inner.dynamic_partition_tasks_kept_unevaluable_filter,
197            ),
198        }
199    }
200
201    fn record_configured_batch_size_rows(&self, value: usize) {
202        self.inner
203            .configured_batch_size_rows
204            .store(u64::try_from(value).unwrap_or(u64::MAX), Ordering::Relaxed);
205    }
206
207    fn record_dynamic_partition_task_pruned(&self) {
208        saturating_fetch_add(&self.inner.dynamic_partition_tasks_pruned, 1);
209    }
210
211    fn record_dynamic_partition_task_kept(&self) {
212        saturating_fetch_add(&self.inner.dynamic_partition_tasks_kept, 1);
213    }
214
215    fn record_dynamic_filters_received(&self, value: usize) {
216        saturating_fetch_add(
217            &self.inner.dynamic_filters_received,
218            u64::try_from(value).unwrap_or(u64::MAX),
219        );
220    }
221
222    fn record_dynamic_filters_accepted(&self, value: usize) {
223        saturating_fetch_add(
224            &self.inner.dynamic_filters_accepted,
225            u64::try_from(value).unwrap_or(u64::MAX),
226        );
227    }
228
229    fn record_dynamic_filters_rejected(&self, value: usize) {
230        saturating_fetch_add(
231            &self.inner.dynamic_filters_rejected,
232            u64::try_from(value).unwrap_or(u64::MAX),
233        );
234    }
235
236    fn record_dynamic_partition_filter_check(&self) {
237        saturating_fetch_add(&self.inner.dynamic_partition_filter_checks, 1);
238    }
239
240    fn record_dynamic_partition_task_kept_unusable_metadata(&self) {
241        saturating_fetch_add(
242            &self.inner.dynamic_partition_tasks_kept_unusable_metadata,
243            1,
244        );
245    }
246
247    fn record_dynamic_partition_task_kept_unevaluable_filter(&self) {
248        saturating_fetch_add(
249            &self.inner.dynamic_partition_tasks_kept_unevaluable_filter,
250            1,
251        );
252    }
253
254    fn identity(&self) -> usize {
255        Arc::as_ptr(&self.inner) as usize
256    }
257}
258
259impl fmt::Debug for ScanMetrics {
260    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
261        formatter
262            .debug_struct("ScanMetrics")
263            .finish_non_exhaustive()
264    }
265}
266
267fn load(counter: &AtomicU64) -> u64 {
268    counter.load(Ordering::Relaxed)
269}
270
271fn nonzero_load(counter: &AtomicU64) -> Option<u64> {
272    match load(counter) {
273        0 => None,
274        value => Some(value),
275    }
276}
277
278/// Collects distinct Delta DataFusion scan metrics in depth-first plan order.
279pub fn collect_scan_metrics(plan: &dyn ExecutionPlan) -> Vec<ScanMetrics> {
280    fn collect(
281        plan: &dyn ExecutionPlan,
282        seen_plans: &mut HashSet<usize>,
283        seen_metrics: &mut HashSet<usize>,
284        metrics: &mut Vec<ScanMetrics>,
285    ) {
286        let plan_identity = plan as *const dyn ExecutionPlan as *const () as usize;
287        if !seen_plans.insert(plan_identity) {
288            return;
289        }
290        if let Some(scan) = plan.downcast_ref::<DeltaScanExec>() {
291            let handle = scan.metrics.clone();
292            if seen_metrics.insert(handle.identity()) {
293                metrics.push(handle);
294            }
295        }
296        for child in plan.children() {
297            collect(child.as_ref(), seen_plans, seen_metrics, metrics);
298        }
299    }
300
301    let mut metrics = Vec::new();
302    collect(plan, &mut HashSet::new(), &mut HashSet::new(), &mut metrics);
303    metrics
304}
305
306#[allow(dead_code)]
307pub(crate) fn create_datafusion_execution_plan(
308    reader_plan: DeltaScanPlan,
309    datafusion_plan: DataFusionScanPlan,
310    exact_row_predicate: Option<DeltaKernelPredicate>,
311    range_read_estimator: Arc<ParquetRangeReadEstimator>,
312    metrics: ScanMetrics,
313    intra_file_repartitioning: IntraFileRepartitioning,
314) -> Arc<dyn ExecutionPlan> {
315    Arc::new(DeltaScanExec::new(
316        reader_plan,
317        datafusion_plan,
318        exact_row_predicate,
319        range_read_estimator,
320        metrics,
321        intra_file_repartitioning,
322    ))
323}
324
325#[derive(Clone)]
326struct DeltaScanExec {
327    reader_plan: Arc<DeltaScanPlan>,
328    schema: SchemaRef,
329    output_projection: Option<Arc<[usize]>>,
330    exact_row_predicate: Option<DeltaKernelPredicate>,
331    properties: Arc<PlanProperties>,
332    metrics: ScanMetrics,
333    limiter: Arc<ScanReadLimiter>,
334    dynamic_filters: Arc<[RetainedDynamicFilter]>,
335    intra_file_repartitioning: IntraFileRepartitioning,
336    range_read_estimator: Arc<ParquetRangeReadEstimator>,
337    // Created only when DataFusion repartitions files into ranged tasks.
338    parquet_metadata_cache: Option<Arc<ParquetMetadataCache>>,
339    file_tasks_repartitioned: bool,
340}
341
342impl DeltaScanExec {
343    #[allow(dead_code)]
344    fn new(
345        reader_plan: DeltaScanPlan,
346        datafusion_plan: DataFusionScanPlan,
347        exact_row_predicate: Option<DeltaKernelPredicate>,
348        range_read_estimator: Arc<ParquetRangeReadEstimator>,
349        metrics: ScanMetrics,
350        intra_file_repartitioning: IntraFileRepartitioning,
351    ) -> Self {
352        let schema = datafusion_plan.projection.output_schema;
353        let output_projection = datafusion_plan.projection.output_projection.map(Arc::from);
354        let properties = scan_properties(&schema, reader_plan.partitions.len());
355        let limiter = ScanReadLimiter::new(
356            reader_plan.execution_options,
357            reader_plan.partition_target_diagnostic.target_partitions,
358            reader_plan.partitions.len(),
359        );
360
361        Self {
362            reader_plan: Arc::new(reader_plan),
363            schema,
364            output_projection,
365            exact_row_predicate,
366            properties,
367            metrics,
368            limiter,
369            dynamic_filters: Arc::from([]),
370            intra_file_repartitioning,
371            range_read_estimator,
372            parquet_metadata_cache: None,
373            file_tasks_repartitioned: false,
374        }
375    }
376
377    fn with_dynamic_filters(
378        &self,
379        dynamic_filters: Vec<RetainedDynamicFilter>,
380    ) -> Arc<dyn ExecutionPlan> {
381        Arc::new(Self {
382            dynamic_filters: Arc::from(dynamic_filters),
383            ..self.clone()
384        })
385    }
386
387    fn with_repartitioned_partitions(
388        &self,
389        partitions: Vec<DeltaScanPartition>,
390    ) -> Arc<dyn ExecutionPlan> {
391        let partition_count = partitions.len();
392        let mut reader_plan = (*self.reader_plan).clone();
393        reader_plan.partitions = partitions;
394        reader_plan
395            .metrics
396            .record_scan_partitions_planned(partition_count);
397        let target_partitions = reader_plan.partition_target_diagnostic.target_partitions;
398        let limiter = ScanReadLimiter::new(
399            reader_plan.execution_options,
400            target_partitions,
401            partition_count,
402        );
403        Arc::new(Self {
404            reader_plan: Arc::new(reader_plan),
405            properties: scan_properties(&self.schema, partition_count),
406            limiter,
407            // Ranged tasks within this plan share their first metadata load.
408            parquet_metadata_cache: Some(Arc::new(ParquetMetadataCache::default())),
409            file_tasks_repartitioned: true,
410            ..self.clone()
411        })
412    }
413}
414
415fn scan_properties(schema: &SchemaRef, file_partition_count: usize) -> Arc<PlanProperties> {
416    Arc::new(
417        PlanProperties::new(
418            EquivalenceProperties::new(Arc::clone(schema)),
419            // DataFusion requires an executable partition even when there are no file tasks.
420            Partitioning::UnknownPartitioning(file_partition_count.max(1)),
421            EmissionType::Incremental,
422            Boundedness::Bounded,
423        )
424        .with_scheduling_type(SchedulingType::Cooperative),
425    )
426}
427
428/// Uses DataFusion to split and regroup file tasks across a partition target.
429///
430/// DataFusion flattens the current groups and aims for
431/// `ceil(total_input_bytes / target_partitions)` bytes per output group. The
432/// `minimum_total_bytes` argument controls whether the complete input is large
433/// enough to repartition; it is not a generated range size.
434///
435/// `Ok(None)` preserves the existing plan when the selected policy does not
436/// apply, file sizes are unavailable, or DataFusion finds no useful
437/// repartitioning.
438fn repartition_file_tasks(
439    partitions: &[DeltaScanPartition],
440    target_partitions: usize,
441    minimum_total_bytes: usize,
442    policy: IntraFileRepartitioning,
443) -> DataFusionResult<Option<Vec<DeltaScanPartition>>> {
444    if target_partitions == 0 {
445        return Err(adapter_error("scan_partition_target_must_be_positive"));
446    }
447    // Whole-file planning already balances up to the requested partition count.
448    // The default avoids extra ranged reads unless the scan lacks parallelism.
449    if !policy.allows_repartitioning(partitions.len(), target_partitions) {
450        return Ok(None);
451    }
452
453    // DataFusion cannot safely split tasks whose physical file size is unknown.
454    let Some(file_groups) = file_groups_from_partitions(partitions)? else {
455        return Ok(None);
456    };
457    let Some(file_groups) = FileGroupPartitioner::new()
458        .with_target_partitions(target_partitions)
459        .with_repartition_file_min_size(minimum_total_bytes)
460        .repartition_file_groups(&file_groups)
461    else {
462        // Preserve the original plan when DataFusion finds no useful split.
463        return Ok(None);
464    };
465
466    partitions_from_file_groups(file_groups).map(Some)
467}
468
469fn file_groups_from_partitions(
470    partitions: &[DeltaScanPartition],
471) -> DataFusionResult<Option<Vec<FileGroup>>> {
472    let mut groups = Vec::with_capacity(partitions.len());
473    for partition in partitions {
474        let mut files = Vec::with_capacity(partition.file_tasks.len());
475        for task in &partition.file_tasks {
476            let Some(file) = partitioned_file_from_task(task)? else {
477                return Ok(None);
478            };
479            files.push(file);
480        }
481        groups.push(FileGroup::new(files));
482    }
483    Ok(Some(groups))
484}
485
486fn partitioned_file_from_task(
487    task: &DeltaScanFileTask,
488) -> DataFusionResult<Option<PartitionedFile>> {
489    let Some(file_size) = task.file_size.filter(|size| *size > 0) else {
490        return Ok(None);
491    };
492    let mut file = PartitionedFile::new(&task.path, file_size);
493    if let Some(range) = &task.parquet_byte_range {
494        if range.start >= range.end || range.end > file_size {
495            return Err(adapter_error("scan_file_range_invalid"));
496        }
497        file.range = Some(FileRange {
498            start: i64::try_from(range.start)
499                .map_err(|_| adapter_error("scan_file_range_invalid"))?,
500            end: i64::try_from(range.end).map_err(|_| adapter_error("scan_file_range_invalid"))?,
501        });
502    }
503    // DataFusion copies extensions to every output range. Carry the Delta task
504    // so its schema, partition, transform, and deletion-vector metadata survive.
505    Ok(Some(file.with_extension(task.clone())))
506}
507
508fn partitions_from_file_groups(
509    groups: Vec<FileGroup>,
510) -> DataFusionResult<Vec<DeltaScanPartition>> {
511    let mut partitions = Vec::with_capacity(groups.len());
512    for group in groups {
513        let tasks = group
514            .into_inner()
515            .into_iter()
516            .map(task_from_partitioned_file)
517            .collect::<DataFusionResult<Vec<_>>>()?;
518        partitions.push(DeltaScanPartition { file_tasks: tasks });
519    }
520    Ok(partitions)
521}
522
523fn task_from_partitioned_file(file: PartitionedFile) -> DataFusionResult<DeltaScanFileTask> {
524    let mut task = file
525        .extension::<DeltaScanFileTask>()
526        .cloned()
527        .ok_or_else(|| adapter_error("scan_file_task_extension_missing"))?;
528    let range = file
529        .range
530        .ok_or_else(|| adapter_error("scan_file_range_missing"))?;
531    let start = u64::try_from(range.start).map_err(|_| adapter_error("scan_file_range_invalid"))?;
532    let end = u64::try_from(range.end).map_err(|_| adapter_error("scan_file_range_invalid"))?;
533    let file_size = task
534        .file_size
535        .ok_or_else(|| adapter_error("scan_file_size_missing"))?;
536    if start >= end || end > file_size {
537        return Err(adapter_error("scan_file_range_invalid"));
538    }
539    task.parquet_byte_range = (start != 0 || end != file_size).then_some(start..end);
540    Ok(task)
541}
542
543impl fmt::Debug for DeltaScanExec {
544    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
545        formatter
546            .debug_struct("DeltaScanExec")
547            .field("snapshot_version", &self.reader_plan.snapshot_version)
548            .field(
549                "partition_count",
550                &self.properties.output_partitioning().partition_count(),
551            )
552            .field("dynamic_filter_count", &self.dynamic_filters.len())
553            .finish_non_exhaustive()
554    }
555}
556
557impl DisplayAs for DeltaScanExec {
558    fn fmt_as(
559        &self,
560        display_type: DisplayFormatType,
561        formatter: &mut fmt::Formatter,
562    ) -> fmt::Result {
563        match display_type {
564            DisplayFormatType::Default | DisplayFormatType::Verbose => write!(
565                formatter,
566                "DeltaScanExec: snapshot_version={}, partitions={}",
567                self.reader_plan.snapshot_version,
568                self.properties.output_partitioning().partition_count()
569            ),
570            DisplayFormatType::TreeRender => write!(formatter, "DeltaScanExec"),
571        }
572    }
573}
574
575impl ExecutionPlan for DeltaScanExec {
576    fn name(&self) -> &str {
577        "DeltaScanExec"
578    }
579
580    fn properties(&self) -> &Arc<PlanProperties> {
581        &self.properties
582    }
583
584    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
585        vec![]
586    }
587
588    fn with_new_children(
589        self: Arc<Self>,
590        children: Vec<Arc<dyn ExecutionPlan>>,
591    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
592        if children.is_empty() {
593            Ok(self)
594        } else {
595            Err(DataFusionError::Internal(
596                "DeltaScanExec does not accept child execution plans".to_owned(),
597            ))
598        }
599    }
600
601    fn repartitioned(
602        &self,
603        _datafusion_target_partitions: usize,
604        config: &ConfigOptions,
605    ) -> DataFusionResult<Option<Arc<dyn ExecutionPlan>>> {
606        if self.file_tasks_repartitioned
607            || self.reader_plan.execution_options.parquet_backend() != ParquetReaderBackend::Direct
608        {
609            return Ok(None);
610        }
611        // Scan planning already applied the provider override and resource caps.
612        let target_partitions = self
613            .reader_plan
614            .partition_target_diagnostic
615            .target_partitions;
616        let Some(partitions) = repartition_file_tasks(
617            &self.reader_plan.partitions,
618            target_partitions,
619            config.optimizer.repartition_file_min_size,
620            self.intra_file_repartitioning,
621        )?
622        else {
623            return Ok(None);
624        };
625        Ok(Some(self.with_repartitioned_partitions(partitions)))
626    }
627
628    fn execute(
629        &self,
630        partition: usize,
631        context: Arc<TaskContext>,
632    ) -> DataFusionResult<SendableRecordBatchStream> {
633        if partition >= self.properties.output_partitioning().partition_count() {
634            return Err(adapter_error("scan_partition_index_out_of_range"));
635        }
636        if self.reader_plan.partitions.is_empty() {
637            return Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone(
638                &self.schema,
639            ))));
640        }
641
642        let configured_batch_size_rows = context.session_config().batch_size();
643        self.metrics
644            .record_configured_batch_size_rows(configured_batch_size_rows);
645        let admission = dynamic_partition_admission_policy(
646            self.metrics.clone(),
647            Arc::clone(&self.dynamic_filters),
648        );
649        let executor = match self.reader_plan.execution_options.parquet_backend() {
650            ParquetReaderBackend::Direct => direct_parquet_file_executor(
651                &self.reader_plan,
652                Some(configured_batch_size_rows),
653                self.exact_row_predicate.clone(),
654                Arc::clone(&self.range_read_estimator),
655                self.parquet_metadata_cache.clone(),
656            ),
657            ParquetReaderBackend::DeltaKernel => {
658                delta_kernel_executor(&self.reader_plan, Arc::clone(&self.range_read_estimator))
659            }
660        };
661        let stream = DeltaScanScheduler::new_with_limiter(
662            Arc::clone(&self.reader_plan),
663            Arc::clone(&self.limiter),
664        )
665        .partition_stream(partition, admission, executor)
666        .map_err(datafusion_error)?;
667        let schema = Arc::clone(&self.schema);
668        let projection = self.output_projection.clone();
669        let stream = stream::unfold(
670            (Some(stream), projection),
671            |(stream, projection)| async move {
672                let mut stream = stream?;
673                let result = stream.next().await?;
674                let result = finalize_output_batch(result, projection.as_deref());
675                let stream = result.is_ok().then_some(stream);
676                Some((result, (stream, projection)))
677            },
678        );
679
680        Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
681    }
682
683    fn handle_child_pushdown_result(
684        &self,
685        phase: FilterPushdownPhase,
686        child_pushdown_result: ChildPushdownResult,
687        _config: &ConfigOptions,
688    ) -> DataFusionResult<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
689        let parent_filters = child_pushdown_result
690            .parent_filters
691            .iter()
692            .map(|result| Arc::clone(&result.filter))
693            .collect::<Vec<_>>();
694        let unsupported = || {
695            FilterPushdownPropagation::with_parent_pushdown_result(vec![
696                PushedDown::No;
697                parent_filters.len()
698            ])
699        };
700        if phase != FilterPushdownPhase::Post || parent_filters.is_empty() {
701            return Ok(unsupported());
702        }
703
704        let classification = DynamicFilterClassification::from_filters(
705            &parent_filters,
706            &self.schema,
707            &self.reader_plan.partition_columns,
708        );
709        let accepted_filters = classification
710            .accepted_filters()
711            .cloned()
712            .collect::<Vec<_>>();
713        let accepted = accepted_filters.len();
714        self.metrics
715            .record_dynamic_filters_received(parent_filters.len());
716        self.metrics.record_dynamic_filters_accepted(accepted);
717        self.metrics
718            .record_dynamic_filters_rejected(parent_filters.len().saturating_sub(accepted));
719        if accepted_filters.is_empty() {
720            return Ok(unsupported());
721        }
722
723        let pushed = classification
724            .decisions
725            .iter()
726            .map(|decision| match decision {
727                DynamicFilterDecision::Accepted(_) => PushedDown::Yes,
728                DynamicFilterDecision::Rejected => PushedDown::No,
729            })
730            .collect();
731        Ok(
732            FilterPushdownPropagation::with_parent_pushdown_result(pushed)
733                .with_updated_node(self.with_dynamic_filters(accepted_filters)),
734        )
735    }
736}
737
738fn dynamic_partition_admission_policy(
739    metrics: ScanMetrics,
740    filters: Arc<[RetainedDynamicFilter]>,
741) -> FileAdmissionPolicy<DeltaScanFileTask> {
742    Arc::new(move |task| {
743        if filters.is_empty() {
744            return Ok(FileAdmissionDecision::Admit);
745        }
746
747        let mut unusable_metadata = false;
748        let mut unevaluable_filter = false;
749        for filter in filters.iter() {
750            metrics.record_dynamic_partition_filter_check();
751            match evaluate_dynamic_partition_filter(filter, task) {
752                DynamicPartitionPruningDecision::Prune => {
753                    metrics.record_dynamic_partition_task_pruned();
754                    return Ok(FileAdmissionDecision::Skip);
755                }
756                DynamicPartitionPruningDecision::Keep(reason) => {
757                    unusable_metadata |= is_unusable_metadata(reason);
758                    unevaluable_filter |= is_unevaluable_filter(reason);
759                }
760            }
761        }
762        if unusable_metadata {
763            metrics.record_dynamic_partition_task_kept_unusable_metadata();
764        }
765        if unevaluable_filter {
766            metrics.record_dynamic_partition_task_kept_unevaluable_filter();
767        }
768        metrics.record_dynamic_partition_task_kept();
769        Ok(FileAdmissionDecision::Admit)
770    })
771}
772
773fn is_unusable_metadata(reason: DynamicPartitionKeepReason) -> bool {
774    matches!(
775        reason,
776        DynamicPartitionKeepReason::PartitionMetadataInvalid
777            | DynamicPartitionKeepReason::PartitionValueMissing
778            | DynamicPartitionKeepReason::PartitionValueUnparseable
779    )
780}
781
782fn is_unevaluable_filter(reason: DynamicPartitionKeepReason) -> bool {
783    matches!(
784        reason,
785        DynamicPartitionKeepReason::FilterUnavailable
786            | DynamicPartitionKeepReason::UnsupportedPartitionType
787            | DynamicPartitionKeepReason::EvaluationFailed
788            | DynamicPartitionKeepReason::NonBooleanResult
789    )
790}
791
792fn project_output_batch(
793    batch: RecordBatch,
794    projection: Option<&[usize]>,
795) -> Result<RecordBatch, arrow::error::ArrowError> {
796    match projection {
797        Some(projection) => batch.project(projection),
798        None => Ok(batch),
799    }
800}
801
802fn finalize_output_batch(
803    result: Result<RecordBatch, DeltaReaderError>,
804    projection: Option<&[usize]>,
805) -> DataFusionResult<RecordBatch> {
806    let batch = result.map_err(datafusion_error)?;
807    project_output_batch(batch, projection).map_err(|source| {
808        datafusion_error(DeltaReaderError::DataFusionAdapter {
809            reason: "scan_output_projection_failed",
810            source: Box::new(DataFusionError::from(source)),
811        })
812    })
813}
814
815fn datafusion_error(error: DeltaReaderError) -> DataFusionError {
816    DataFusionError::External(Box::new(error))
817}
818
819fn adapter_error(reason: &'static str) -> DataFusionError {
820    datafusion_error(DeltaReaderError::DataFusionAdapter {
821        reason,
822        source: Box::new(DataFusionError::Execution(reason.to_owned())),
823    })
824}
825
826#[cfg(test)]
827mod tests {
828    use std::{
829        collections::HashSet,
830        error::Error,
831        fs,
832        path::{Path, PathBuf},
833        thread,
834        time::{SystemTime, UNIX_EPOCH},
835    };
836
837    use arrow::array::StringArray;
838    use arrow::{
839        array::Int32Array,
840        datatypes::{DataType, Field, Schema},
841        record_batch::RecordBatch,
842    };
843    use datafusion::physical_plan::filter::FilterExec;
844    use datafusion::{
845        common::config::ConfigOptions,
846        datasource::TableProvider,
847        logical_expr::{Operator, col, lit},
848        physical_expr::expressions::{
849            BinaryExpr, Column, DynamicFilterPhysicalExpr, lit as physical_lit,
850        },
851        physical_plan::{
852            ExecutionPlan,
853            filter_pushdown::{
854                ChildFilterPushdownResult, ChildPushdownResult, FilterPushdownPhase, PushedDown,
855            },
856            union::UnionExec,
857        },
858        prelude::{SessionConfig, SessionContext},
859    };
860    use futures_util::StreamExt;
861    use parquet::arrow::ArrowWriter;
862    use serde_json::{Value, json};
863
864    use super::*;
865    use crate::{
866        DeltaScanExecutionOptions, DeltaTable, DeltaTableBuilder,
867        delta::kernel::kernel_pruning_predicate,
868        reader::datafusion::{
869            DeltaTableProvider, ScanOptions,
870            planning::{FilterCapabilities, plan_datafusion_scan},
871        },
872        reader::planning::{
873            DeltaScanPartitionTargetOptions, build_physical_row_predicate, plan_scan,
874        },
875    };
876
877    type TestResult<T = ()> = Result<T, Box<dyn Error>>;
878
879    struct TestTable(PathBuf);
880
881    impl TestTable {
882        fn empty(name: &str) -> TestResult<Self> {
883            let nanos = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos();
884            let path = Path::new("target")
885                .join("delta-arrow-reader-datafusion-tests")
886                .join(format!("{}-{name}-{nanos}", std::process::id()));
887            fs::create_dir_all(path.join("_delta_log"))?;
888            let table = Self(path);
889            table.write_log(&[protocol(), metadata()])?;
890            Ok(table)
891        }
892
893        fn partitioned(name: &str) -> TestResult<Self> {
894            let table = Self::empty(name)?;
895            let west = table.write_parquet("west.parquet", &[1, 2])?;
896            let east = table.write_parquet("east.parquet", &[3, 4])?;
897            table.write_log(&[
898                protocol(),
899                metadata(),
900                add("west.parquet", west, "west", 2, 1, 2),
901                add("east.parquet", east, "east", 2, 3, 4),
902            ])?;
903            Ok(table)
904        }
905
906        fn late_dynamic(name: &str) -> TestResult<Self> {
907            let table = Self::empty(name)?;
908            let west = table.write_parquet("west.parquet", &[1, 2, 3])?;
909            let east = table.write_parquet("east.parquet", &[4, 5])?;
910            table.write_log(&[
911                protocol(),
912                metadata(),
913                add("west.parquet", west, "west", 3, 1, 3),
914                add("east.parquet", east, "east", 2, 4, 5),
915            ])?;
916            Ok(table)
917        }
918
919        fn missing(name: &str) -> TestResult<Self> {
920            let table = Self::partitioned(name)?;
921            table.write_log(&[
922                protocol(),
923                metadata(),
924                add("missing.parquet", 100, "west", 1, 1, 1),
925            ])?;
926            Ok(table)
927        }
928
929        fn uri(&self) -> String {
930            self.0.to_string_lossy().into_owned()
931        }
932
933        fn write_parquet(&self, name: &str, ids: &[i32]) -> TestResult<u64> {
934            let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
935            let batch = RecordBatch::try_new(
936                Arc::clone(&schema),
937                vec![Arc::new(Int32Array::from(ids.to_vec()))],
938            )?;
939            let path = self.0.join(name);
940            let mut writer = ArrowWriter::try_new(fs::File::create(&path)?, schema, None)?;
941            writer.write(&batch)?;
942            writer.close()?;
943            Ok(fs::metadata(path)?.len())
944        }
945
946        fn write_log(&self, actions: &[Value]) -> TestResult {
947            let contents = actions
948                .iter()
949                .map(Value::to_string)
950                .collect::<Vec<_>>()
951                .join("\n");
952            fs::write(
953                self.0.join("_delta_log/00000000000000000000.json"),
954                format!("{contents}\n"),
955            )?;
956            Ok(())
957        }
958    }
959
960    impl Drop for TestTable {
961        fn drop(&mut self) {
962            let _ = fs::remove_dir_all(&self.0);
963        }
964    }
965
966    fn protocol() -> Value {
967        json!({"protocol": {"minReaderVersion": 1, "minWriterVersion": 2}})
968    }
969
970    fn metadata() -> Value {
971        let schema = json!({
972            "type": "struct",
973            "fields": [
974                {"name": "id", "type": "integer", "nullable": false, "metadata": {}},
975                {"name": "region", "type": "string", "nullable": true, "metadata": {}}
976            ]
977        });
978        json!({
979            "metaData": {
980                "id": "delta-arrow-reader-datafusion-test",
981                "format": {"provider": "parquet", "options": {}},
982                "schemaString": schema.to_string(),
983                "partitionColumns": ["region"],
984                "configuration": {},
985                "createdTime": 1587968585495_i64
986            }
987        })
988    }
989
990    fn add(
991        path: &str,
992        size: u64,
993        region: &str,
994        num_records: u64,
995        min_id: i32,
996        max_id: i32,
997    ) -> Value {
998        let stats = json!({
999            "numRecords": num_records,
1000            "minValues": {"id": min_id},
1001            "maxValues": {"id": max_id},
1002            "nullCount": {"id": 0}
1003        });
1004        json!({
1005            "add": {
1006                "path": path,
1007                "partitionValues": {"region": region},
1008                "size": size,
1009                "modificationTime": 1587968586000_i64,
1010                "dataChange": true,
1011                "stats": stats.to_string()
1012            }
1013        })
1014    }
1015
1016    fn build_plan(
1017        table: &DeltaTable,
1018        projection: Option<&[usize]>,
1019        filters: &[datafusion::logical_expr::Expr],
1020        target_partitions: usize,
1021        execution_options: DeltaScanExecutionOptions,
1022        registration_name: Option<String>,
1023    ) -> Result<Arc<dyn ExecutionPlan>, DeltaReaderError> {
1024        build_plan_with_repartitioning(
1025            table,
1026            projection,
1027            filters,
1028            target_partitions,
1029            execution_options,
1030            registration_name,
1031            IntraFileRepartitioning::default(),
1032        )
1033    }
1034
1035    fn build_plan_with_repartitioning(
1036        table: &DeltaTable,
1037        projection: Option<&[usize]>,
1038        filters: &[datafusion::logical_expr::Expr],
1039        target_partitions: usize,
1040        execution_options: DeltaScanExecutionOptions,
1041        registration_name: Option<String>,
1042        intra_file_repartitioning: IntraFileRepartitioning,
1043    ) -> Result<Arc<dyn ExecutionPlan>, DeltaReaderError> {
1044        let partition_columns = table
1045            .partition_columns()
1046            .iter()
1047            .cloned()
1048            .collect::<HashSet<_>>();
1049        let filter_refs = filters.iter().collect::<Vec<_>>();
1050        let datafusion_plan = plan_datafusion_scan(
1051            &table.schema(),
1052            &partition_columns,
1053            projection,
1054            &filter_refs,
1055            FilterCapabilities {
1056                supports_exact_row_filtering: execution_options.parquet_backend()
1057                    == ParquetReaderBackend::Direct,
1058            },
1059        )?;
1060        let scan_projection = datafusion_plan.projection.scan_projection.clone();
1061        let hidden_columns = datafusion_plan.projection.hidden_columns.clone();
1062        let pruning_predicate = datafusion_plan
1063            .filters
1064            .pruning_predicate
1065            .as_ref()
1066            .and_then(kernel_pruning_predicate);
1067        let exact_row_predicate = match datafusion_plan.filters.exact_row_predicate.as_ref() {
1068            Some(predicate) => Some(kernel_pruning_predicate(predicate).ok_or(
1069                DeltaReaderError::UnsupportedPredicate {
1070                    reason: "exact_row_predicate_not_kernel_safe",
1071                },
1072            )?),
1073            None => None,
1074        };
1075        let exact_row_predicate = build_physical_row_predicate(
1076            table.snapshot(),
1077            scan_projection.as_deref(),
1078            &hidden_columns,
1079            exact_row_predicate,
1080        )?;
1081        let include_stats = datafusion_plan.filters.requires_statistics;
1082        let reader_plan = plan_scan(
1083            table.snapshot(),
1084            scan_projection.as_deref(),
1085            &hidden_columns,
1086            pruning_predicate,
1087            include_stats,
1088            execution_options,
1089            DeltaScanPartitionTargetOptions {
1090                explicit_target_partitions: Some(target_partitions),
1091                datafusion_target_partitions: None,
1092            },
1093        )?;
1094        let metrics = ScanMetrics::new(registration_name, reader_plan.metrics.clone(), true);
1095        Ok(create_datafusion_execution_plan(
1096            reader_plan,
1097            datafusion_plan,
1098            exact_row_predicate,
1099            Arc::default(),
1100            metrics,
1101            intra_file_repartitioning,
1102        ))
1103    }
1104
1105    fn session(batch_size: usize) -> SessionContext {
1106        SessionContext::new_with_config(SessionConfig::new().with_batch_size(batch_size))
1107    }
1108
1109    #[tokio::test]
1110    async fn native_and_datafusion_scans_use_the_recovered_profile() -> TestResult {
1111        use crate::reader::backend::direct_parquet::tests::network_warmup_fixture;
1112        use crate::{DeltaComparison, DeltaPredicate, DeltaScalar};
1113        use futures_util::TryStreamExt;
1114
1115        for datafusion_first in [false, true] {
1116            let fixture = TestTable::partitioned("recovered-profile")?;
1117            let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1118            let original = table.range_read_estimator();
1119            // Inject failed calibration while keeping real Delta query execution.
1120            let (warmup, probe) = network_warmup_fixture().await?;
1121            probe.fail_range(1);
1122            assert!(warmup.run("initialization", &original).await.is_none());
1123            table.range_read_state.defer_network_warmup(warmup).await;
1124            let requests_before_recovery = probe.range_call_count();
1125            let predicate = DeltaPredicate::Compare {
1126                column: "id".into(),
1127                op: DeltaComparison::Gt,
1128                value: DeltaScalar::Int32(1),
1129            };
1130            let filters = [col("id").gt(lit(1_i32))];
1131            let no_matches = [col("id").gt(lit(99_i32))];
1132            let provider = DeltaTableProvider::try_new(table.clone(), ScanOptions::default())?;
1133            let context = SessionContext::new();
1134            for scan in [
1135                table.scan(),
1136                table.scan().with_predicate(predicate.clone()).with_limit(0),
1137                table.scan().with_predicate(DeltaPredicate::Compare {
1138                    column: "id".into(),
1139                    op: DeltaComparison::Gt,
1140                    value: DeltaScalar::Int32(99),
1141                }),
1142            ] {
1143                let _ = scan.build().await?;
1144            }
1145            for (filters, limit) in [
1146                (&[][..], None),
1147                (filters.as_slice(), Some(0)),
1148                (no_matches.as_slice(), None),
1149            ] {
1150                provider
1151                    .scan(&context.state(), None, filters, limit)
1152                    .await?;
1153            }
1154            assert_eq!(probe.range_call_count(), requests_before_recovery);
1155            assert!(table.network_warmup_pending(DeltaScanExecutionOptions::default()));
1156
1157            for use_datafusion in [datafusion_first, !datafusion_first, datafusion_first] {
1158                let mut actual = if use_datafusion {
1159                    let plan = provider
1160                        .scan(&context.state(), None, &filters, None)
1161                        .await?;
1162                    let scan = plan
1163                        .downcast_ref::<DeltaScanExec>()
1164                        .ok_or("expected DeltaScanExec")?;
1165                    assert!(Arc::ptr_eq(
1166                        &scan.range_read_estimator,
1167                        &table.range_read_estimator()
1168                    ));
1169                    ids(&datafusion::physical_plan::collect(plan, context.task_ctx()).await?)
1170                } else {
1171                    let scan = table
1172                        .scan()
1173                        .with_predicate(predicate.clone())
1174                        .build()
1175                        .await?;
1176                    assert!(Arc::ptr_eq(
1177                        &scan.range_read_estimator,
1178                        &table.range_read_estimator()
1179                    ));
1180                    ids(&scan.into_stream().try_collect::<Vec<_>>().await?)
1181                };
1182                actual.sort_unstable();
1183                assert_eq!(actual, [2, 3, 4]);
1184                assert!(!Arc::ptr_eq(&original, &table.range_read_estimator()));
1185                assert_eq!(probe.range_call_count(), requests_before_recovery + 3_180);
1186            }
1187        }
1188        Ok(())
1189    }
1190
1191    #[tokio::test]
1192    async fn table_scopes_the_range_estimator_and_repartitioning_creates_a_metadata_cache()
1193    -> TestResult {
1194        let fixture = TestTable::partitioned("shared-range-read-estimator")?;
1195        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1196        let provider = DeltaTableProvider::try_new(table.clone(), ScanOptions::default())?;
1197        let shared_provider = DeltaTableProvider::try_new(table.clone(), ScanOptions::default())?;
1198        let separate_table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1199        let separate_provider =
1200            DeltaTableProvider::try_new(separate_table, ScanOptions::default())?;
1201        let context = SessionContext::new();
1202
1203        let first = provider.scan(&context.state(), None, &[], None).await?;
1204        let second = provider.scan(&context.state(), None, &[], None).await?;
1205        let shared = shared_provider
1206            .scan(&context.state(), None, &[], None)
1207            .await?;
1208        let separate = separate_provider
1209            .scan(&context.state(), None, &[], None)
1210            .await?;
1211        let first = first
1212            .as_ref()
1213            .downcast_ref::<DeltaScanExec>()
1214            .ok_or("first plan was not DeltaScanExec")?;
1215        let second = second
1216            .as_ref()
1217            .downcast_ref::<DeltaScanExec>()
1218            .ok_or("second plan was not DeltaScanExec")?;
1219        let separate = separate
1220            .as_ref()
1221            .downcast_ref::<DeltaScanExec>()
1222            .ok_or("separate provider plan was not DeltaScanExec")?;
1223        let repartitioned =
1224            first.with_repartitioned_partitions(first.reader_plan.partitions.clone());
1225        let repartitioned = repartitioned
1226            .as_ref()
1227            .downcast_ref::<DeltaScanExec>()
1228            .ok_or("repartitioned plan was not DeltaScanExec")?;
1229        assert!(Arc::ptr_eq(
1230            &first.range_read_estimator,
1231            &second.range_read_estimator
1232        ));
1233        let shared = shared
1234            .as_ref()
1235            .downcast_ref::<DeltaScanExec>()
1236            .ok_or("expected shared scan")?;
1237        assert!(Arc::ptr_eq(
1238            &first.range_read_estimator,
1239            &shared.range_read_estimator
1240        ));
1241        let native = table.scan().build().await?;
1242        let refreshed = table.refresh().await?;
1243        let refreshed_estimator = refreshed.range_read_estimator();
1244        assert!(Arc::ptr_eq(
1245            &first.range_read_estimator,
1246            &native.range_read_estimator
1247        ));
1248        assert!(Arc::ptr_eq(
1249            &first.range_read_estimator,
1250            &refreshed_estimator
1251        ));
1252        assert!(Arc::ptr_eq(
1253            &first.range_read_estimator,
1254            &repartitioned.range_read_estimator
1255        ));
1256        assert!(!Arc::ptr_eq(
1257            &first.range_read_estimator,
1258            &separate.range_read_estimator
1259        ));
1260        assert!(repartitioned.file_tasks_repartitioned);
1261        assert!(repartitioned.parquet_metadata_cache.is_some());
1262        Ok(())
1263    }
1264
1265    fn sized_file_task(path: &str, size: Option<u64>) -> DeltaScanFileTask {
1266        use crate::{
1267            delta::kernel::KernelPhysicalToLogicalTransform,
1268            reader::deletion_vector::DeletionVectorMetadata,
1269        };
1270
1271        DeltaScanFileTask {
1272            path: path.to_owned(),
1273            file_size: size,
1274            parquet_byte_range: None,
1275            modification_time_ms: None,
1276            partition_values: Default::default(),
1277            deletion_vector: DeletionVectorMetadata::default(),
1278            transform: KernelPhysicalToLogicalTransform::default(),
1279        }
1280    }
1281
1282    fn partition(file_tasks: Vec<DeltaScanFileTask>) -> DeltaScanPartition {
1283        DeltaScanPartition { file_tasks }
1284    }
1285
1286    fn partition_estimated_bytes(partition: &DeltaScanPartition) -> Option<u64> {
1287        partition
1288            .file_tasks
1289            .iter()
1290            .map(DeltaScanFileTask::estimated_scan_bytes)
1291            .sum()
1292    }
1293
1294    #[allow(clippy::expect_used)]
1295    fn ids(batches: &[RecordBatch]) -> Vec<i32> {
1296        batches
1297            .iter()
1298            .flat_map(|batch| {
1299                batch
1300                    .column(batch.schema().index_of("id").expect("id column"))
1301                    .as_any()
1302                    .downcast_ref::<Int32Array>()
1303                    .expect("Int32 id")
1304                    .values()
1305                    .iter()
1306                    .copied()
1307                    .collect::<Vec<_>>()
1308            })
1309            .collect()
1310    }
1311
1312    #[test]
1313    fn file_repartitioning_balances_ranges_without_losing_file_identity() -> TestResult {
1314        let input = vec![partition(vec![
1315            sized_file_task("large.parquet", Some(100)),
1316            sized_file_task("small.parquet", Some(20)),
1317        ])];
1318        let partitions = repartition_file_tasks(&input, 4, 1, IntraFileRepartitioning::default())?
1319            .ok_or("not repartitioned")?;
1320
1321        assert_eq!(
1322            partitions
1323                .iter()
1324                .map(partition_estimated_bytes)
1325                .collect::<Vec<_>>(),
1326            vec![Some(30); 4]
1327        );
1328        let mut ranges = std::collections::BTreeMap::<_, Vec<_>>::new();
1329        for task in partitions
1330            .iter()
1331            .flat_map(|partition| &partition.file_tasks)
1332        {
1333            ranges.entry(task.path.as_str()).or_default().push(
1334                task.parquet_byte_range
1335                    .clone()
1336                    .unwrap_or(0..task.file_size.ok_or("missing file size")?),
1337            );
1338            assert_eq!(
1339                task.file_size,
1340                Some(if task.path == "large.parquet" {
1341                    100
1342                } else {
1343                    20
1344                })
1345            );
1346            if task.path == "small.parquet" {
1347                assert!(task.parquet_byte_range.is_none());
1348            }
1349        }
1350        assert_eq!(
1351            ranges.remove("large.parquet"),
1352            Some(vec![0..30, 30..60, 60..90, 90..100])
1353        );
1354        assert_eq!(
1355            ranges.remove("small.parquet"),
1356            Some(std::iter::once(0..20).collect())
1357        );
1358        assert!(ranges.is_empty());
1359
1360        Ok(())
1361    }
1362
1363    #[test]
1364    fn file_repartitioning_never_escapes_an_existing_range() -> TestResult {
1365        let mut task = sized_file_task("partial.parquet", Some(100));
1366        task.parquet_byte_range = Some(20..80);
1367        let partitions = repartition_file_tasks(
1368            &[partition(vec![task])],
1369            3,
1370            1,
1371            IntraFileRepartitioning::default(),
1372        )?
1373        .ok_or("partial range was not repartitioned")?;
1374
1375        assert_eq!(
1376            partitions
1377                .iter()
1378                .flat_map(|partition| &partition.file_tasks)
1379                .map(|task| task.parquet_byte_range.clone())
1380                .collect::<Vec<_>>(),
1381            [Some(20..40), Some(40..60), Some(60..80)]
1382        );
1383        assert!(
1384            partitions
1385                .iter()
1386                .flat_map(|partition| &partition.file_tasks)
1387                .all(|task| task.file_size == Some(100))
1388        );
1389
1390        Ok(())
1391    }
1392
1393    #[test]
1394    fn file_repartitioning_policy_controls_full_whole_file_plans() -> TestResult {
1395        let input = vec![
1396            partition(vec![sized_file_task("huge.parquet", Some(1_000))]),
1397            partition(vec![sized_file_task("small-0.parquet", Some(10))]),
1398            partition(vec![sized_file_task("small-1.parquet", Some(10))]),
1399            partition(vec![sized_file_task("small-2.parquet", Some(10))]),
1400        ];
1401
1402        assert!(
1403            repartition_file_tasks(&input, 4, 1, IntraFileRepartitioning::WhenBelowTarget)?
1404                .is_none()
1405        );
1406        let rebalanced = repartition_file_tasks(&input, 4, 1, IntraFileRepartitioning::Always)?
1407            .ok_or("full plan was not rebalanced")?;
1408        assert_eq!(
1409            rebalanced
1410                .iter()
1411                .map(partition_estimated_bytes)
1412                .collect::<Vec<_>>(),
1413            [Some(258), Some(258), Some(258), Some(256)]
1414        );
1415        assert!(
1416            rebalanced
1417                .iter()
1418                .flat_map(|partition| &partition.file_tasks)
1419                .any(|task| task.parquet_byte_range.is_some())
1420        );
1421        assert!(
1422            repartition_file_tasks(&input, 4, 1_031, IntraFileRepartitioning::Always)?.is_none()
1423        );
1424        assert!(repartition_file_tasks(&[], 4, 1, IntraFileRepartitioning::Always)?.is_none());
1425        assert!(
1426            input
1427                .iter()
1428                .flat_map(|partition| &partition.file_tasks)
1429                .all(|task| task.parquet_byte_range.is_none())
1430        );
1431        Ok(())
1432    }
1433
1434    #[test]
1435    fn file_repartitioning_refuses_unsupported_inputs_and_rejects_invalid_ranges() -> TestResult {
1436        let input = vec![partition(vec![sized_file_task("known.parquet", Some(120))])];
1437
1438        assert!(
1439            repartition_file_tasks(&input, 4, 121, IntraFileRepartitioning::default())?.is_none()
1440        );
1441        assert!(
1442            repartition_file_tasks(
1443                &[partition(vec![sized_file_task("unknown.parquet", None)])],
1444                4,
1445                1,
1446                IntraFileRepartitioning::default(),
1447            )?
1448            .is_none()
1449        );
1450        assert!(
1451            repartition_file_tasks(
1452                &[partition(vec![sized_file_task("empty.parquet", Some(0))])],
1453                4,
1454                1,
1455                IntraFileRepartitioning::default(),
1456            )?
1457            .is_none()
1458        );
1459        assert!(repartition_file_tasks(&input, 0, 1, IntraFileRepartitioning::default()).is_err());
1460
1461        for range in [90..110, std::ops::Range { start: 90, end: 80 }, 90..90] {
1462            let mut invalid = sized_file_task("invalid.parquet", Some(100));
1463            invalid.parquet_byte_range = Some(range);
1464            assert!(
1465                repartition_file_tasks(
1466                    &[partition(vec![invalid])],
1467                    4,
1468                    1,
1469                    IntraFileRepartitioning::default(),
1470                )
1471                .is_err()
1472            );
1473        }
1474
1475        let oversized = u64::try_from(i64::MAX)? + 1;
1476        assert!(
1477            repartition_file_tasks(
1478                &[partition(vec![sized_file_task(
1479                    "oversized.parquet",
1480                    Some(oversized),
1481                )])],
1482                2,
1483                1,
1484                IntraFileRepartitioning::default(),
1485            )
1486            .is_err()
1487        );
1488
1489        Ok(())
1490    }
1491
1492    #[test]
1493    fn task_from_partitioned_file_rejects_malformed_partitioner_output() {
1494        let mut missing_extension = PartitionedFile::new("missing-extension.parquet", 100);
1495        missing_extension.range = Some(FileRange { start: 0, end: 100 });
1496        assert!(task_from_partitioned_file(missing_extension).is_err());
1497
1498        let missing_range = PartitionedFile::new("missing-range.parquet", 100)
1499            .with_extension(sized_file_task("missing-range.parquet", Some(100)));
1500        assert!(task_from_partitioned_file(missing_range).is_err());
1501
1502        for range in [
1503            FileRange {
1504                start: -1,
1505                end: 100,
1506            },
1507            FileRange { start: 50, end: 50 },
1508            FileRange { start: 0, end: 101 },
1509        ] {
1510            let mut invalid = PartitionedFile::new("invalid-range.parquet", 100)
1511                .with_extension(sized_file_task("invalid-range.parquet", Some(100)));
1512            invalid.range = Some(range);
1513            assert!(task_from_partitioned_file(invalid).is_err());
1514        }
1515
1516        let mut missing_size = PartitionedFile::new("missing-size.parquet", 100)
1517            .with_extension(sized_file_task("missing-size.parquet", None));
1518        missing_size.range = Some(FileRange { start: 0, end: 100 });
1519        assert!(task_from_partitioned_file(missing_size).is_err());
1520    }
1521
1522    #[tokio::test]
1523    async fn direct_repartitioning_reads_each_row_once_and_preserves_byte_accounting() -> TestResult
1524    {
1525        let fixture = TestTable::partitioned("file-repartitioning")?;
1526        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1527        let plan = build_plan_with_repartitioning(
1528            &table,
1529            None,
1530            &[],
1531            4,
1532            DeltaScanExecutionOptions::new(),
1533            None,
1534            IntraFileRepartitioning::Always,
1535        )?;
1536        let expected_bytes = collect_scan_metrics(plan.as_ref())[0]
1537            .snapshot()
1538            .reader_metrics
1539            .estimated_input_bytes;
1540        let mut config = ConfigOptions::new();
1541        config.optimizer.repartition_file_min_size = 1;
1542        let repartitioned = plan
1543            .repartitioned(4, &config)?
1544            .ok_or("Direct backend scan was not repartitioned")?;
1545        assert!(repartitioned.repartitioned(4, &config)?.is_none());
1546
1547        assert_eq!(
1548            repartitioned
1549                .properties()
1550                .output_partitioning()
1551                .partition_count(),
1552            4
1553        );
1554        let mut actual_ids = ids(&datafusion::physical_plan::collect(
1555            Arc::clone(&repartitioned),
1556            session(1024).task_ctx(),
1557        )
1558        .await?);
1559        actual_ids.sort_unstable();
1560        assert_eq!(actual_ids, [1, 2, 3, 4]);
1561        let metrics = collect_scan_metrics(repartitioned.as_ref())[0].snapshot();
1562        assert_eq!(metrics.reader_metrics.scan_partitions_planned, 4);
1563        assert_eq!(
1564            metrics.reader_metrics.estimated_parquet_task_bytes_admitted,
1565            expected_bytes
1566        );
1567
1568        let explicit_one =
1569            build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
1570        assert!(explicit_one.repartitioned(4, &config)?.is_none());
1571
1572        {
1573            let kernel = build_plan_with_repartitioning(
1574                &table,
1575                None,
1576                &[],
1577                4,
1578                DeltaScanExecutionOptions::new()
1579                    .with_parquet_backend(ParquetReaderBackend::DeltaKernel),
1580                None,
1581                IntraFileRepartitioning::Always,
1582            )?;
1583            assert!(kernel.repartitioned(4, &config)?.is_none());
1584        }
1585
1586        Ok(())
1587    }
1588
1589    fn dynamic_filter(name: &str, index: usize) -> Arc<DynamicFilterPhysicalExpr> {
1590        Arc::new(DynamicFilterPhysicalExpr::new(
1591            vec![Arc::new(Column::new(name, index))],
1592            physical_lit(true),
1593        ))
1594    }
1595
1596    fn hook_input(
1597        filters: Vec<Arc<dyn datafusion::physical_plan::PhysicalExpr>>,
1598    ) -> ChildPushdownResult {
1599        ChildPushdownResult {
1600            parent_filters: filters
1601                .into_iter()
1602                .map(|filter| ChildFilterPushdownResult {
1603                    filter,
1604                    child_results: Vec::new(),
1605                })
1606                .collect(),
1607            self_filters: Vec::new(),
1608        }
1609    }
1610
1611    #[tokio::test]
1612    async fn properties_projection_partitions_metrics_and_reexecution_match_provider_behavior()
1613    -> TestResult {
1614        let fixture = TestTable::partitioned("properties")?;
1615        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1616        let logical_filter = col("id").gt(lit(1_i32));
1617        let plan = build_plan(
1618            &table,
1619            Some(&[1, 0]),
1620            &[logical_filter],
1621            2,
1622            DeltaScanExecutionOptions::new(),
1623            None,
1624        )?;
1625
1626        assert_eq!(plan.name(), "DeltaScanExec");
1627        assert!(plan.children().is_empty());
1628        assert!(plan.metrics().is_none());
1629        assert_eq!(plan.schema().fields().len(), 2);
1630        assert_eq!(plan.schema().field(0).name(), "region");
1631        assert_eq!(plan.schema().field(1).name(), "id");
1632        assert_eq!(plan.properties().output_partitioning().partition_count(), 2);
1633        assert_eq!(
1634            plan.partition_statistics(None)?,
1635            Arc::new(datafusion::common::Statistics::new_unknown(&plan.schema()))
1636        );
1637        let context = session(1);
1638        let first =
1639            datafusion::physical_plan::collect(Arc::clone(&plan), context.task_ctx()).await?;
1640        assert!(first.iter().all(|batch| batch.num_rows() <= 1));
1641        let mut first_ids = ids(&first);
1642        first_ids.sort_unstable();
1643        assert_eq!(first_ids, [2, 3, 4]);
1644
1645        let second =
1646            datafusion::physical_plan::collect(Arc::clone(&plan), context.task_ctx()).await?;
1647        let mut second_ids = ids(&second);
1648        second_ids.sort_unstable();
1649        assert_eq!(second_ids, first_ids);
1650        let handles = collect_scan_metrics(plan.as_ref());
1651        assert_eq!(handles.len(), 1);
1652        assert_eq!(handles[0].registration_name(), None);
1653        let metrics = handles[0].snapshot();
1654        assert_eq!(metrics.configured_batch_size_rows, Some(1));
1655        assert_eq!(metrics.reader_metrics.scan_partitions_started, 4);
1656        assert_eq!(metrics.reader_metrics.file_tasks_completed, 4);
1657        assert_eq!(metrics.reader_metrics.scheduler_rows_emitted, 6);
1658
1659        let hidden = build_plan(
1660            &table,
1661            Some(&[1]),
1662            &[col("id").gt(lit(1_i32))],
1663            1,
1664            DeltaScanExecutionOptions::new(),
1665            None,
1666        )?;
1667        let hidden_batches = datafusion::physical_plan::collect(
1668            Arc::clone(&hidden),
1669            SessionContext::new().task_ctx(),
1670        )
1671        .await?;
1672        assert_eq!(hidden.schema().fields().len(), 1);
1673        assert_eq!(hidden.schema().field(0).name(), "region");
1674        assert!(hidden_batches.iter().all(|batch| batch.num_columns() == 1));
1675        assert_eq!(
1676            hidden_batches
1677                .iter()
1678                .map(RecordBatch::num_rows)
1679                .sum::<usize>(),
1680            3
1681        );
1682
1683        let partition_filter = build_plan(
1684            &table,
1685            None,
1686            &[col("region").eq(lit("west"))],
1687            2,
1688            DeltaScanExecutionOptions::new(),
1689            None,
1690        )?;
1691        let partition_batches = datafusion::physical_plan::collect(
1692            Arc::clone(&partition_filter),
1693            SessionContext::new().task_ctx(),
1694        )
1695        .await?;
1696        assert_eq!(ids(&partition_batches), [1, 2]);
1697        assert_eq!(
1698            collect_scan_metrics(partition_filter.as_ref())[0]
1699                .snapshot()
1700                .reader_metrics
1701                .file_tasks_started,
1702            1
1703        );
1704
1705        let empty = build_plan(
1706            &table,
1707            Some(&[]),
1708            &[],
1709            1,
1710            DeltaScanExecutionOptions::new(),
1711            None,
1712        )?;
1713        let empty_batches = datafusion::physical_plan::collect(
1714            Arc::clone(&empty),
1715            SessionContext::new().task_ctx(),
1716        )
1717        .await?;
1718        assert!(empty.schema().fields().is_empty());
1719        assert!(empty_batches.iter().all(|batch| batch.num_columns() == 0));
1720        assert_eq!(
1721            empty_batches
1722                .iter()
1723                .map(RecordBatch::num_rows)
1724                .sum::<usize>(),
1725            4
1726        );
1727
1728        let empty_fixture = TestTable::empty("empty-scan")?;
1729        let empty_table = DeltaTableBuilder::new(empty_fixture.uri())
1730            .load_table()
1731            .await?;
1732        let empty_plan = build_plan(
1733            &empty_table,
1734            None,
1735            &[],
1736            1,
1737            DeltaScanExecutionOptions::new(),
1738            None,
1739        )?;
1740        assert_eq!(
1741            empty_plan
1742                .properties()
1743                .output_partitioning()
1744                .partition_count(),
1745            1
1746        );
1747        assert!(
1748            datafusion::physical_plan::collect(empty_plan, SessionContext::new().task_ctx(),)
1749                .await?
1750                .is_empty()
1751        );
1752
1753        let invalid = plan.execute(2, context.task_ctx());
1754        let error = match invalid {
1755            Ok(_) => return Err("out-of-range partition unexpectedly executed".into()),
1756            Err(error) => error,
1757        };
1758        let DataFusionError::External(source) = error else {
1759            return Err("invalid partition did not preserve the reader error".into());
1760        };
1761        let reader = source
1762            .downcast_ref::<DeltaReaderError>()
1763            .ok_or("external error was not DeltaReaderError")?;
1764        assert_eq!(reader.code(), "datafusion_adapter");
1765        Ok(())
1766    }
1767
1768    #[tokio::test]
1769    async fn dynamic_filter_hook_prunes_before_file_start_and_counts_once() -> TestResult {
1770        let fixture = TestTable::partitioned("dynamic")?;
1771        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1772        let plan = build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
1773        let dynamic = dynamic_filter("region", 1);
1774        let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic.clone();
1775        let rejected: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic_filter("id", 0);
1776        let pushed = plan.handle_child_pushdown_result(
1777            FilterPushdownPhase::Post,
1778            hook_input(vec![physical, rejected]),
1779            &ConfigOptions::new(),
1780        )?;
1781        assert!(matches!(
1782            pushed.filters.as_slice(),
1783            [PushedDown::Yes, PushedDown::No]
1784        ));
1785        let updated = pushed.updated_node.ok_or("dynamic plan was not retained")?;
1786        dynamic.update(Arc::new(BinaryExpr::new(
1787            Arc::new(Column::new("region", 1)),
1788            Operator::Eq,
1789            physical_lit("west"),
1790        )))?;
1791
1792        let batches = datafusion::physical_plan::collect(
1793            Arc::clone(&updated),
1794            SessionContext::new().task_ctx(),
1795        )
1796        .await?;
1797        assert_eq!(ids(&batches), [1, 2]);
1798        let metrics = collect_scan_metrics(updated.as_ref())
1799            .pop()
1800            .ok_or("missing dynamic metrics")?
1801            .snapshot();
1802        assert_eq!(metrics.dynamic_filters_received, 2);
1803        assert_eq!(metrics.dynamic_filters_accepted, 1);
1804        assert_eq!(metrics.dynamic_filters_rejected, 1);
1805        assert_eq!(metrics.dynamic_partition_filter_checks, 2);
1806        assert_eq!(metrics.dynamic_partition_tasks_pruned, 1);
1807        assert_eq!(metrics.dynamic_partition_tasks_kept, 1);
1808        assert_eq!(metrics.reader_metrics.file_tasks_started, 1);
1809        assert_eq!(metrics.reader_metrics.file_tasks_completed, 1);
1810        assert_eq!(
1811            collect_scan_metrics(plan.as_ref())[0]
1812                .snapshot()
1813                .dynamic_filters_received,
1814            2
1815        );
1816        Ok(())
1817    }
1818
1819    #[tokio::test]
1820    async fn physical_pushdown_preserves_dynamic_filters_across_plan_rebuild() -> TestResult {
1821        let fixture = TestTable::partitioned("dynamic-plan-rebuild")?;
1822        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1823        let plan = build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
1824        let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> =
1825            dynamic_filter("region", 1);
1826        let pushed = plan.handle_child_pushdown_result(
1827            FilterPushdownPhase::Post,
1828            hook_input(vec![physical]),
1829            &ConfigOptions::new(),
1830        )?;
1831        let updated = pushed.updated_node.ok_or("expected updated scan")?;
1832        let rebuilt = Arc::clone(&updated).with_new_children(Vec::new())?;
1833        let reset = updated.reset_state()?;
1834
1835        for candidate in [&rebuilt, &reset] {
1836            let debug = format!("{candidate:?}");
1837            assert!(debug.contains("dynamic_filter_count: 1"), "{debug}");
1838        }
1839        let display = datafusion::physical_plan::displayable(rebuilt.as_ref())
1840            .one_line()
1841            .to_string();
1842        assert!(display.contains("DeltaScanExec:"), "{display}");
1843        assert!(display.contains("partitions="), "{display}");
1844        assert!(!display.contains("DynamicFilter"), "{display}");
1845        assert!(
1846            Arc::clone(&rebuilt)
1847                .with_new_children(vec![Arc::clone(&rebuilt)])
1848                .is_err()
1849        );
1850        Ok(())
1851    }
1852
1853    #[tokio::test]
1854    async fn late_dynamic_filter_keeps_admitted_file_and_prunes_the_next() -> TestResult {
1855        let fixture = TestTable::late_dynamic("late-dynamic")?;
1856        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1857        let options = DeltaScanExecutionOptions::new()
1858            .with_prefetch_files_per_partition(0)
1859            .with_max_concurrent_file_reads_per_partition(1)?
1860            .with_max_concurrent_file_reads_per_scan(Some(1))?
1861            .with_output_buffer_batches_per_partition(1)?;
1862        let plan = build_plan(&table, None, &[], 1, options, None)?;
1863        let dynamic = dynamic_filter("region", 1);
1864        let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic.clone();
1865        let pushed = plan.handle_child_pushdown_result(
1866            FilterPushdownPhase::Post,
1867            hook_input(vec![physical]),
1868            &ConfigOptions::new(),
1869        )?;
1870        let updated = pushed.updated_node.ok_or("dynamic plan was not retained")?;
1871        let mut stream = updated.execute(0, session(1).task_ctx())?;
1872        let first = stream.next().await.ok_or("missing first batch")??;
1873        assert_eq!(ids(std::slice::from_ref(&first)), [1]);
1874
1875        dynamic.update(Arc::new(BinaryExpr::new(
1876            Arc::new(Column::new("region", 1)),
1877            Operator::Eq,
1878            physical_lit("none"),
1879        )))?;
1880        let mut batches = vec![first];
1881        while let Some(batch) = stream.next().await {
1882            batches.push(batch?);
1883        }
1884
1885        assert_eq!(ids(&batches), [1, 2, 3]);
1886        let metrics = collect_scan_metrics(updated.as_ref())[0].snapshot();
1887        assert_eq!(metrics.dynamic_partition_filter_checks, 2);
1888        assert_eq!(metrics.dynamic_partition_tasks_kept, 1);
1889        assert_eq!(metrics.dynamic_partition_tasks_pruned, 1);
1890        assert_eq!(metrics.reader_metrics.file_tasks_started, 1);
1891        assert_eq!(metrics.reader_metrics.file_tasks_completed, 1);
1892        Ok(())
1893    }
1894
1895    #[tokio::test]
1896    async fn hook_is_post_only_empty_safe_and_collector_is_ordered_and_distinct() -> TestResult {
1897        let fixture = TestTable::partitioned("collector")?;
1898        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1899        let first = build_plan(
1900            &table,
1901            None,
1902            &[],
1903            1,
1904            DeltaScanExecutionOptions::new(),
1905            Some("first".to_owned()),
1906        )?;
1907        let second = build_plan(
1908            &table,
1909            None,
1910            &[],
1911            1,
1912            DeltaScanExecutionOptions::new(),
1913            Some("second".to_owned()),
1914        )?;
1915        let dynamic = dynamic_filter("region", 1);
1916        let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic;
1917        let pre = first.handle_child_pushdown_result(
1918            FilterPushdownPhase::Pre,
1919            hook_input(vec![physical]),
1920            &ConfigOptions::new(),
1921        )?;
1922        assert!(pre.updated_node.is_none());
1923        assert!(matches!(pre.filters.as_slice(), [PushedDown::No]));
1924        let empty = first.handle_child_pushdown_result(
1925            FilterPushdownPhase::Post,
1926            hook_input(Vec::new()),
1927            &ConfigOptions::new(),
1928        )?;
1929        assert!(empty.updated_node.is_none());
1930
1931        let union: Arc<dyn ExecutionPlan> = UnionExec::try_new(vec![
1932            Arc::clone(&first),
1933            Arc::clone(&second),
1934            Arc::clone(&first),
1935        ])?;
1936        let handles = collect_scan_metrics(union.as_ref());
1937        assert_eq!(handles.len(), 2);
1938        assert_eq!(handles[0].registration_name(), Some("first"));
1939        assert_eq!(handles[1].registration_name(), Some("second"));
1940        assert!(!format!("{:?}", handles[0]).contains("first"));
1941        let initial = handles[0].snapshot();
1942        assert_eq!(initial.configured_batch_size_rows, None);
1943        assert_eq!(
1944            [
1945                initial.dynamic_partition_tasks_pruned,
1946                initial.dynamic_partition_tasks_kept,
1947                initial.dynamic_filters_received,
1948                initial.dynamic_filters_accepted,
1949                initial.dynamic_filters_rejected,
1950                initial.dynamic_partition_filter_checks,
1951                initial.dynamic_partition_tasks_kept_unusable_metadata,
1952                initial.dynamic_partition_tasks_kept_unevaluable_filter,
1953            ],
1954            [0; 8]
1955        );
1956
1957        let accepted: Arc<dyn datafusion::physical_plan::PhysicalExpr> =
1958            dynamic_filter("region", 1);
1959        let updated = first
1960            .handle_child_pushdown_result(
1961                FilterPushdownPhase::Post,
1962                hook_input(vec![accepted]),
1963                &ConfigOptions::new(),
1964            )?
1965            .updated_node
1966            .ok_or("expected updated scan")?;
1967        let shared_metrics_union: Arc<dyn ExecutionPlan> =
1968            UnionExec::try_new(vec![updated, Arc::clone(&first), Arc::clone(&second)])?;
1969        let shared_handles = collect_scan_metrics(shared_metrics_union.as_ref());
1970        assert_eq!(shared_handles.len(), 2);
1971        assert_eq!(shared_handles[0].registration_name(), Some("first"));
1972        assert_eq!(shared_handles[1].registration_name(), Some("second"));
1973        assert_eq!(handles[0].identity(), shared_handles[0].identity());
1974        assert_ne!(handles[0].identity(), shared_handles[1].identity());
1975
1976        drop(shared_metrics_union);
1977        drop(union);
1978        drop(first);
1979        drop(second);
1980        assert_eq!(handles[0].snapshot().reader_metrics.file_tasks_started, 0);
1981        assert_eq!(handles[0].registration_name(), Some("first"));
1982        Ok(())
1983    }
1984
1985    #[tokio::test]
1986    async fn dynamic_partition_admission_reason_counts_are_once_per_file_and_saturating()
1987    -> TestResult {
1988        use crate::{
1989            delta::kernel::KernelPhysicalToLogicalTransform,
1990            reader::datafusion::dynamic_filters::DynamicFilterClassification,
1991            reader::deletion_vector::DeletionVectorMetadata,
1992        };
1993
1994        let fixture = TestTable::partitioned("dynamic-counters")?;
1995        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1996        let plan = build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
1997        let metrics = collect_scan_metrics(plan.as_ref())
1998            .pop()
1999            .ok_or("missing metrics")?;
2000        let schema = Arc::new(Schema::new(vec![
2001            Field::new("id", DataType::Int32, false),
2002            Field::new("region", DataType::Utf8, true),
2003        ]));
2004        let retained = |dynamic: Arc<DynamicFilterPhysicalExpr>| -> TestResult<_> {
2005            let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic;
2006            Ok(DynamicFilterClassification::from_filters(
2007                std::slice::from_ref(&physical),
2008                &schema,
2009                &["region".to_owned()],
2010            )
2011            .accepted_filters()
2012            .next()
2013            .cloned()
2014            .ok_or("dynamic filter was not retained")?)
2015        };
2016        let first = retained(dynamic_filter("region", 1))?;
2017        let second = retained(dynamic_filter("region", 1))?;
2018        let missing = DeltaScanFileTask {
2019            path: "missing-partition.parquet".to_owned(),
2020            file_size: None,
2021            parquet_byte_range: None,
2022            modification_time_ms: None,
2023            partition_values: Default::default(),
2024            deletion_vector: DeletionVectorMetadata::default(),
2025            transform: KernelPhysicalToLogicalTransform::default(),
2026        };
2027        assert_eq!(
2028            dynamic_partition_admission_policy(metrics.clone(), Arc::from([first, second]))(
2029                &missing,
2030            )?,
2031            FileAdmissionDecision::Admit
2032        );
2033        let snapshot = metrics.snapshot();
2034        assert_eq!(snapshot.dynamic_partition_filter_checks, 2);
2035        assert_eq!(snapshot.dynamic_partition_tasks_kept, 1);
2036        assert_eq!(snapshot.dynamic_partition_tasks_kept_unusable_metadata, 1);
2037
2038        let rejecting = dynamic_filter("region", 1);
2039        rejecting.update(physical_lit(false))?;
2040        let first = retained(rejecting)?;
2041        let second = retained(dynamic_filter("region", 1))?;
2042        let mut present = missing.clone();
2043        present
2044            .partition_values
2045            .insert("region".to_owned(), "west".to_owned());
2046        assert_eq!(
2047            dynamic_partition_admission_policy(metrics.clone(), Arc::from([first, second]))(
2048                &present,
2049            )?,
2050            FileAdmissionDecision::Skip
2051        );
2052        let snapshot = metrics.snapshot();
2053        assert_eq!(snapshot.dynamic_partition_filter_checks, 3);
2054        assert_eq!(snapshot.dynamic_partition_tasks_pruned, 1);
2055        assert_eq!(snapshot.dynamic_partition_tasks_kept, 1);
2056        assert_eq!(snapshot.dynamic_partition_tasks_kept_unusable_metadata, 1);
2057
2058        let unsupported = dynamic_filter("region", 1);
2059        unsupported.update(physical_lit("not boolean"))?;
2060        let admission = dynamic_partition_admission_policy(
2061            metrics.clone(),
2062            Arc::from([retained(unsupported)?]),
2063        );
2064        assert_eq!(admission(&present)?, FileAdmissionDecision::Admit);
2065        let snapshot = metrics.snapshot();
2066        assert_eq!(snapshot.dynamic_partition_filter_checks, 4);
2067        assert_eq!(snapshot.dynamic_partition_tasks_kept, 2);
2068        assert_eq!(snapshot.dynamic_partition_tasks_kept_unevaluable_filter, 1);
2069
2070        metrics
2071            .inner
2072            .dynamic_filters_received
2073            .store(u64::MAX - 1, Ordering::Relaxed);
2074        metrics.record_dynamic_filters_received(2);
2075        metrics.record_dynamic_filters_received(1);
2076        assert_eq!(metrics.snapshot().dynamic_filters_received, u64::MAX);
2077        Ok(())
2078    }
2079
2080    #[tokio::test]
2081    async fn dynamic_metrics_updates_are_thread_safe() -> TestResult {
2082        const THREADS: usize = 4;
2083        const ITERATIONS: usize = 100;
2084
2085        let fixture = TestTable::partitioned("dynamic-metrics-concurrency")?;
2086        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
2087        let plan = build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
2088        let metrics = collect_scan_metrics(plan.as_ref())
2089            .pop()
2090            .ok_or("missing metrics")?;
2091        let mut handles = Vec::new();
2092
2093        for _ in 0..THREADS {
2094            let metrics = metrics.clone();
2095            handles.push(thread::spawn(move || {
2096                for _ in 0..ITERATIONS {
2097                    metrics.record_dynamic_partition_task_pruned();
2098                    metrics.record_dynamic_partition_task_kept();
2099                    metrics.record_dynamic_filters_received(3);
2100                    metrics.record_dynamic_filters_accepted(1);
2101                    metrics.record_dynamic_filters_rejected(2);
2102                    metrics.record_dynamic_partition_filter_check();
2103                    metrics.record_dynamic_partition_task_kept_unusable_metadata();
2104                    metrics.record_dynamic_partition_task_kept_unevaluable_filter();
2105                }
2106            }));
2107        }
2108        for handle in handles {
2109            handle.join().map_err(|_| "metrics worker panicked")?;
2110        }
2111
2112        let calls = u64::try_from(THREADS * ITERATIONS)?;
2113        let snapshot = metrics.snapshot();
2114        assert_eq!(snapshot.dynamic_partition_tasks_pruned, calls);
2115        assert_eq!(snapshot.dynamic_partition_tasks_kept, calls);
2116        assert_eq!(snapshot.dynamic_filters_received, calls * 3);
2117        assert_eq!(snapshot.dynamic_filters_accepted, calls);
2118        assert_eq!(snapshot.dynamic_filters_rejected, calls * 2);
2119        assert_eq!(snapshot.dynamic_partition_filter_checks, calls);
2120        assert_eq!(
2121            snapshot.dynamic_partition_tasks_kept_unusable_metadata,
2122            calls
2123        );
2124        assert_eq!(
2125            snapshot.dynamic_partition_tasks_kept_unevaluable_filter,
2126            calls
2127        );
2128        Ok(())
2129    }
2130
2131    #[tokio::test]
2132    async fn execution_error_and_stream_drop_preserve_partial_metrics() -> TestResult {
2133        let missing_fixture = TestTable::missing("error")?;
2134        let missing_table = DeltaTableBuilder::new(missing_fixture.uri())
2135            .load_table()
2136            .await?;
2137        let missing_plan = build_plan(
2138            &missing_table,
2139            None,
2140            &[],
2141            1,
2142            DeltaScanExecutionOptions::new(),
2143            None,
2144        )?;
2145        let result = datafusion::physical_plan::collect(
2146            Arc::clone(&missing_plan),
2147            SessionContext::new().task_ctx(),
2148        )
2149        .await;
2150        let error = result.expect_err("missing file must fail");
2151        assert!(matches!(&error, DataFusionError::External(_)));
2152        assert!(!error.to_string().contains("missing.parquet"));
2153        let failed = collect_scan_metrics(missing_plan.as_ref())
2154            .pop()
2155            .ok_or("missing failure metrics")?;
2156        assert_eq!(failed.snapshot().reader_metrics.file_tasks_started, 1);
2157
2158        let fixture = TestTable::partitioned("drop")?;
2159        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
2160        let options = DeltaScanExecutionOptions::new()
2161            .with_prefetch_files_per_partition(0)
2162            .with_max_concurrent_file_reads_per_partition(1)?
2163            .with_max_concurrent_file_reads_per_scan(Some(1))?
2164            .with_output_buffer_batches_per_partition(1)?;
2165        let drop_plan = build_plan(&table, None, &[], 1, options, None)?;
2166        let handle = collect_scan_metrics(drop_plan.as_ref())
2167            .pop()
2168            .ok_or("missing drop metrics")?;
2169        let mut stream = drop_plan.execute(0, SessionContext::new().task_ctx())?;
2170        assert!(stream.next().await.transpose()?.is_some());
2171        drop(stream);
2172        tokio::task::yield_now().await;
2173        let stable = handle.snapshot();
2174        tokio::task::yield_now().await;
2175        assert_eq!(handle.snapshot(), stable);
2176        assert!(stable.reader_metrics.file_tasks_started >= 1);
2177        let retry = datafusion::physical_plan::collect(
2178            Arc::clone(&drop_plan),
2179            SessionContext::new().task_ctx(),
2180        )
2181        .await?;
2182        assert_eq!(ids(&retry), [1, 2, 3, 4]);
2183        Ok(())
2184    }
2185
2186    #[tokio::test]
2187    async fn parquet_backends_produce_the_same_logical_rows() -> TestResult {
2188        let fixture = TestTable::partitioned("backends")?;
2189        let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
2190        let mut outputs = Vec::new();
2191        for backend in [
2192            ParquetReaderBackend::Direct,
2193            ParquetReaderBackend::DeltaKernel,
2194        ] {
2195            let options = DeltaScanExecutionOptions::new().with_parquet_backend(backend);
2196            let plan = build_plan(&table, Some(&[1, 0]), &[], 2, options, None)?;
2197            let mut batches =
2198                datafusion::physical_plan::collect(plan, SessionContext::new().task_ctx()).await?;
2199            batches.sort_by_key(|batch| {
2200                batch
2201                    .column(1)
2202                    .as_any()
2203                    .downcast_ref::<Int32Array>()
2204                    .expect("Int32 id")
2205                    .value(0)
2206            });
2207            outputs.push(
2208                batches
2209                    .iter()
2210                    .flat_map(|batch| {
2211                        let ids = batch
2212                            .column(1)
2213                            .as_any()
2214                            .downcast_ref::<Int32Array>()
2215                            .expect("Int32 id");
2216                        let regions = batch
2217                            .column(0)
2218                            .as_any()
2219                            .downcast_ref::<StringArray>()
2220                            .expect("Utf8 region");
2221                        (0..batch.num_rows())
2222                            .map(|row| (regions.value(row).to_owned(), ids.value(row)))
2223                            .collect::<Vec<_>>()
2224                    })
2225                    .collect::<Vec<_>>(),
2226            );
2227        }
2228        assert_eq!(outputs[0], outputs[1]);
2229
2230        let kernel_options = DeltaScanExecutionOptions::new()
2231            .with_parquet_backend(ParquetReaderBackend::DeltaKernel);
2232        let inexact = build_plan(
2233            &table,
2234            None,
2235            &[col("id").gt(lit(1_i32))],
2236            1,
2237            kernel_options,
2238            None,
2239        )?;
2240        let unfiltered = datafusion::physical_plan::collect(
2241            Arc::clone(&inexact),
2242            SessionContext::new().task_ctx(),
2243        )
2244        .await?;
2245        assert_eq!(ids(&unfiltered), [1, 2, 3, 4]);
2246
2247        let residual: Arc<dyn datafusion::physical_plan::PhysicalExpr> = Arc::new(BinaryExpr::new(
2248            Arc::new(Column::new("id", 0)),
2249            Operator::Gt,
2250            physical_lit(1_i32),
2251        ));
2252        let residual_plan: Arc<dyn ExecutionPlan> =
2253            Arc::new(FilterExec::try_new(residual, inexact)?);
2254        let filtered =
2255            datafusion::physical_plan::collect(residual_plan, SessionContext::new().task_ctx())
2256                .await?;
2257        assert_eq!(ids(&filtered), [2, 3, 4]);
2258        Ok(())
2259    }
2260}