1use 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#[non_exhaustive]
78#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
79pub enum IntraFileRepartitioning {
80 #[default]
82 WhenBelowTarget,
83 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#[non_exhaustive]
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct ScanMetricsSnapshot {
103 pub reader_metrics: DeltaScanMetricsSnapshot,
105 pub uses_arrow_view_types: bool,
107 pub configured_batch_size_rows: Option<u64>,
109 pub dynamic_partition_tasks_pruned: u64,
112 pub dynamic_partition_tasks_kept: u64,
115 pub dynamic_filters_received: u64,
117 pub dynamic_filters_accepted: u64,
119 pub dynamic_filters_rejected: u64,
121 pub dynamic_partition_filter_checks: u64,
123 pub dynamic_partition_tasks_kept_unusable_metadata: u64,
125 pub dynamic_partition_tasks_kept_unevaluable_filter: u64,
127}
128
129#[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 pub fn registration_name(&self) -> Option<&str> {
176 self.inner.registration_name.as_deref()
177 }
178
179 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
278pub 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 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 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 Partitioning::UnknownPartitioning(file_partition_count.max(1)),
421 EmissionType::Incremental,
422 Boundedness::Bounded,
423 )
424 .with_scheduling_type(SchedulingType::Cooperative),
425 )
426}
427
428fn 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 if !policy.allows_repartitioning(partitions.len(), target_partitions) {
450 return Ok(None);
451 }
452
453 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 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 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 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 => delta_kernel_executor(&self.reader_plan),
658 };
659 let stream = DeltaScanScheduler::new_with_limiter(
660 Arc::clone(&self.reader_plan),
661 Arc::clone(&self.limiter),
662 )
663 .partition_stream(partition, admission, executor)
664 .map_err(datafusion_error)?;
665 let schema = Arc::clone(&self.schema);
666 let projection = self.output_projection.clone();
667 let stream = stream::unfold(
668 (Some(stream), projection),
669 |(stream, projection)| async move {
670 let mut stream = stream?;
671 let result = stream.next().await?;
672 let result = finalize_output_batch(result, projection.as_deref());
673 let stream = result.is_ok().then_some(stream);
674 Some((result, (stream, projection)))
675 },
676 );
677
678 Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
679 }
680
681 fn handle_child_pushdown_result(
682 &self,
683 phase: FilterPushdownPhase,
684 child_pushdown_result: ChildPushdownResult,
685 _config: &ConfigOptions,
686 ) -> DataFusionResult<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
687 let parent_filters = child_pushdown_result
688 .parent_filters
689 .iter()
690 .map(|result| Arc::clone(&result.filter))
691 .collect::<Vec<_>>();
692 let unsupported = || {
693 FilterPushdownPropagation::with_parent_pushdown_result(vec![
694 PushedDown::No;
695 parent_filters.len()
696 ])
697 };
698 if phase != FilterPushdownPhase::Post || parent_filters.is_empty() {
699 return Ok(unsupported());
700 }
701
702 let classification = DynamicFilterClassification::from_filters(
703 &parent_filters,
704 &self.schema,
705 &self.reader_plan.partition_columns,
706 );
707 let accepted_filters = classification
708 .accepted_filters()
709 .cloned()
710 .collect::<Vec<_>>();
711 let accepted = accepted_filters.len();
712 self.metrics
713 .record_dynamic_filters_received(parent_filters.len());
714 self.metrics.record_dynamic_filters_accepted(accepted);
715 self.metrics
716 .record_dynamic_filters_rejected(parent_filters.len().saturating_sub(accepted));
717 if accepted_filters.is_empty() {
718 return Ok(unsupported());
719 }
720
721 let pushed = classification
722 .decisions
723 .iter()
724 .map(|decision| match decision {
725 DynamicFilterDecision::Accepted(_) => PushedDown::Yes,
726 DynamicFilterDecision::Rejected => PushedDown::No,
727 })
728 .collect();
729 Ok(
730 FilterPushdownPropagation::with_parent_pushdown_result(pushed)
731 .with_updated_node(self.with_dynamic_filters(accepted_filters)),
732 )
733 }
734}
735
736fn dynamic_partition_admission_policy(
737 metrics: ScanMetrics,
738 filters: Arc<[RetainedDynamicFilter]>,
739) -> FileAdmissionPolicy<DeltaScanFileTask> {
740 Arc::new(move |task| {
741 if filters.is_empty() {
742 return Ok(FileAdmissionDecision::Admit);
743 }
744
745 let mut unusable_metadata = false;
746 let mut unevaluable_filter = false;
747 for filter in filters.iter() {
748 metrics.record_dynamic_partition_filter_check();
749 match evaluate_dynamic_partition_filter(filter, task) {
750 DynamicPartitionPruningDecision::Prune => {
751 metrics.record_dynamic_partition_task_pruned();
752 return Ok(FileAdmissionDecision::Skip);
753 }
754 DynamicPartitionPruningDecision::Keep(reason) => {
755 unusable_metadata |= is_unusable_metadata(reason);
756 unevaluable_filter |= is_unevaluable_filter(reason);
757 }
758 }
759 }
760 if unusable_metadata {
761 metrics.record_dynamic_partition_task_kept_unusable_metadata();
762 }
763 if unevaluable_filter {
764 metrics.record_dynamic_partition_task_kept_unevaluable_filter();
765 }
766 metrics.record_dynamic_partition_task_kept();
767 Ok(FileAdmissionDecision::Admit)
768 })
769}
770
771fn is_unusable_metadata(reason: DynamicPartitionKeepReason) -> bool {
772 matches!(
773 reason,
774 DynamicPartitionKeepReason::PartitionMetadataInvalid
775 | DynamicPartitionKeepReason::PartitionValueMissing
776 | DynamicPartitionKeepReason::PartitionValueUnparseable
777 )
778}
779
780fn is_unevaluable_filter(reason: DynamicPartitionKeepReason) -> bool {
781 matches!(
782 reason,
783 DynamicPartitionKeepReason::FilterUnavailable
784 | DynamicPartitionKeepReason::UnsupportedPartitionType
785 | DynamicPartitionKeepReason::EvaluationFailed
786 | DynamicPartitionKeepReason::NonBooleanResult
787 )
788}
789
790fn project_output_batch(
791 batch: RecordBatch,
792 projection: Option<&[usize]>,
793) -> Result<RecordBatch, arrow::error::ArrowError> {
794 match projection {
795 Some(projection) => batch.project(projection),
796 None => Ok(batch),
797 }
798}
799
800fn finalize_output_batch(
801 result: Result<RecordBatch, DeltaReaderError>,
802 projection: Option<&[usize]>,
803) -> DataFusionResult<RecordBatch> {
804 let batch = result.map_err(datafusion_error)?;
805 project_output_batch(batch, projection).map_err(|source| {
806 datafusion_error(DeltaReaderError::DataFusionAdapter {
807 reason: "scan_output_projection_failed",
808 source: Box::new(DataFusionError::from(source)),
809 })
810 })
811}
812
813fn datafusion_error(error: DeltaReaderError) -> DataFusionError {
814 DataFusionError::External(Box::new(error))
815}
816
817fn adapter_error(reason: &'static str) -> DataFusionError {
818 datafusion_error(DeltaReaderError::DataFusionAdapter {
819 reason,
820 source: Box::new(DataFusionError::Execution(reason.to_owned())),
821 })
822}
823
824#[cfg(test)]
825mod tests {
826 use std::{
827 collections::HashSet,
828 error::Error,
829 fs,
830 path::{Path, PathBuf},
831 thread,
832 time::{SystemTime, UNIX_EPOCH},
833 };
834
835 use arrow::array::StringArray;
836 use arrow::{
837 array::Int32Array,
838 datatypes::{DataType, Field, Schema},
839 record_batch::RecordBatch,
840 };
841 use datafusion::physical_plan::filter::FilterExec;
842 use datafusion::{
843 common::config::ConfigOptions,
844 datasource::TableProvider,
845 logical_expr::{Operator, col, lit},
846 physical_expr::expressions::{
847 BinaryExpr, Column, DynamicFilterPhysicalExpr, lit as physical_lit,
848 },
849 physical_plan::{
850 ExecutionPlan,
851 filter_pushdown::{
852 ChildFilterPushdownResult, ChildPushdownResult, FilterPushdownPhase, PushedDown,
853 },
854 union::UnionExec,
855 },
856 prelude::{SessionConfig, SessionContext},
857 };
858 use futures_util::StreamExt;
859 use parquet::arrow::ArrowWriter;
860 use serde_json::{Value, json};
861
862 use super::*;
863 use crate::{
864 DeltaScanExecutionOptions, DeltaTable, DeltaTableBuilder,
865 delta::kernel::kernel_pruning_predicate,
866 reader::datafusion::{
867 DeltaTableProvider, ScanOptions,
868 planning::{FilterCapabilities, plan_datafusion_scan},
869 },
870 reader::planning::{
871 DeltaScanPartitionTargetOptions, build_physical_row_predicate, plan_scan,
872 },
873 };
874
875 type TestResult<T = ()> = Result<T, Box<dyn Error>>;
876
877 struct TestTable(PathBuf);
878
879 impl TestTable {
880 fn empty(name: &str) -> TestResult<Self> {
881 let nanos = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos();
882 let path = Path::new("target")
883 .join("delta-arrow-reader-datafusion-tests")
884 .join(format!("{}-{name}-{nanos}", std::process::id()));
885 fs::create_dir_all(path.join("_delta_log"))?;
886 let table = Self(path);
887 table.write_log(&[protocol(), metadata()])?;
888 Ok(table)
889 }
890
891 fn partitioned(name: &str) -> TestResult<Self> {
892 let table = Self::empty(name)?;
893 let west = table.write_parquet("west.parquet", &[1, 2])?;
894 let east = table.write_parquet("east.parquet", &[3, 4])?;
895 table.write_log(&[
896 protocol(),
897 metadata(),
898 add("west.parquet", west, "west", 2, 1, 2),
899 add("east.parquet", east, "east", 2, 3, 4),
900 ])?;
901 Ok(table)
902 }
903
904 fn late_dynamic(name: &str) -> TestResult<Self> {
905 let table = Self::empty(name)?;
906 let west = table.write_parquet("west.parquet", &[1, 2, 3])?;
907 let east = table.write_parquet("east.parquet", &[4, 5])?;
908 table.write_log(&[
909 protocol(),
910 metadata(),
911 add("west.parquet", west, "west", 3, 1, 3),
912 add("east.parquet", east, "east", 2, 4, 5),
913 ])?;
914 Ok(table)
915 }
916
917 fn missing(name: &str) -> TestResult<Self> {
918 let table = Self::partitioned(name)?;
919 table.write_log(&[
920 protocol(),
921 metadata(),
922 add("missing.parquet", 100, "west", 1, 1, 1),
923 ])?;
924 Ok(table)
925 }
926
927 fn uri(&self) -> String {
928 self.0.to_string_lossy().into_owned()
929 }
930
931 fn write_parquet(&self, name: &str, ids: &[i32]) -> TestResult<u64> {
932 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
933 let batch = RecordBatch::try_new(
934 Arc::clone(&schema),
935 vec![Arc::new(Int32Array::from(ids.to_vec()))],
936 )?;
937 let path = self.0.join(name);
938 let mut writer = ArrowWriter::try_new(fs::File::create(&path)?, schema, None)?;
939 writer.write(&batch)?;
940 writer.close()?;
941 Ok(fs::metadata(path)?.len())
942 }
943
944 fn write_log(&self, actions: &[Value]) -> TestResult {
945 let contents = actions
946 .iter()
947 .map(Value::to_string)
948 .collect::<Vec<_>>()
949 .join("\n");
950 fs::write(
951 self.0.join("_delta_log/00000000000000000000.json"),
952 format!("{contents}\n"),
953 )?;
954 Ok(())
955 }
956 }
957
958 impl Drop for TestTable {
959 fn drop(&mut self) {
960 let _ = fs::remove_dir_all(&self.0);
961 }
962 }
963
964 fn protocol() -> Value {
965 json!({"protocol": {"minReaderVersion": 1, "minWriterVersion": 2}})
966 }
967
968 fn metadata() -> Value {
969 let schema = json!({
970 "type": "struct",
971 "fields": [
972 {"name": "id", "type": "integer", "nullable": false, "metadata": {}},
973 {"name": "region", "type": "string", "nullable": true, "metadata": {}}
974 ]
975 });
976 json!({
977 "metaData": {
978 "id": "delta-arrow-reader-datafusion-test",
979 "format": {"provider": "parquet", "options": {}},
980 "schemaString": schema.to_string(),
981 "partitionColumns": ["region"],
982 "configuration": {},
983 "createdTime": 1587968585495_i64
984 }
985 })
986 }
987
988 fn add(
989 path: &str,
990 size: u64,
991 region: &str,
992 num_records: u64,
993 min_id: i32,
994 max_id: i32,
995 ) -> Value {
996 let stats = json!({
997 "numRecords": num_records,
998 "minValues": {"id": min_id},
999 "maxValues": {"id": max_id},
1000 "nullCount": {"id": 0}
1001 });
1002 json!({
1003 "add": {
1004 "path": path,
1005 "partitionValues": {"region": region},
1006 "size": size,
1007 "modificationTime": 1587968586000_i64,
1008 "dataChange": true,
1009 "stats": stats.to_string()
1010 }
1011 })
1012 }
1013
1014 fn build_plan(
1015 table: &DeltaTable,
1016 projection: Option<&[usize]>,
1017 filters: &[datafusion::logical_expr::Expr],
1018 target_partitions: usize,
1019 execution_options: DeltaScanExecutionOptions,
1020 registration_name: Option<String>,
1021 ) -> Result<Arc<dyn ExecutionPlan>, DeltaReaderError> {
1022 build_plan_with_repartitioning(
1023 table,
1024 projection,
1025 filters,
1026 target_partitions,
1027 execution_options,
1028 registration_name,
1029 IntraFileRepartitioning::default(),
1030 )
1031 }
1032
1033 fn build_plan_with_repartitioning(
1034 table: &DeltaTable,
1035 projection: Option<&[usize]>,
1036 filters: &[datafusion::logical_expr::Expr],
1037 target_partitions: usize,
1038 execution_options: DeltaScanExecutionOptions,
1039 registration_name: Option<String>,
1040 intra_file_repartitioning: IntraFileRepartitioning,
1041 ) -> Result<Arc<dyn ExecutionPlan>, DeltaReaderError> {
1042 let partition_columns = table
1043 .partition_columns()
1044 .iter()
1045 .cloned()
1046 .collect::<HashSet<_>>();
1047 let filter_refs = filters.iter().collect::<Vec<_>>();
1048 let datafusion_plan = plan_datafusion_scan(
1049 &table.schema(),
1050 &partition_columns,
1051 projection,
1052 &filter_refs,
1053 FilterCapabilities {
1054 supports_exact_row_filtering: execution_options.parquet_backend()
1055 == ParquetReaderBackend::Direct,
1056 },
1057 )?;
1058 let scan_projection = datafusion_plan.projection.scan_projection.clone();
1059 let hidden_columns = datafusion_plan.projection.hidden_columns.clone();
1060 let pruning_predicate = datafusion_plan
1061 .filters
1062 .pruning_predicate
1063 .as_ref()
1064 .and_then(kernel_pruning_predicate);
1065 let exact_row_predicate = match datafusion_plan.filters.exact_row_predicate.as_ref() {
1066 Some(predicate) => Some(kernel_pruning_predicate(predicate).ok_or(
1067 DeltaReaderError::UnsupportedPredicate {
1068 reason: "exact_row_predicate_not_kernel_safe",
1069 },
1070 )?),
1071 None => None,
1072 };
1073 let exact_row_predicate = build_physical_row_predicate(
1074 table.snapshot(),
1075 scan_projection.as_deref(),
1076 &hidden_columns,
1077 exact_row_predicate,
1078 )?;
1079 let include_stats = datafusion_plan.filters.requires_statistics;
1080 let reader_plan = plan_scan(
1081 table.snapshot(),
1082 scan_projection.as_deref(),
1083 &hidden_columns,
1084 pruning_predicate,
1085 include_stats,
1086 execution_options,
1087 DeltaScanPartitionTargetOptions {
1088 explicit_target_partitions: Some(target_partitions),
1089 datafusion_target_partitions: None,
1090 },
1091 )?;
1092 let metrics = ScanMetrics::new(registration_name, reader_plan.metrics.clone(), true);
1093 Ok(create_datafusion_execution_plan(
1094 reader_plan,
1095 datafusion_plan,
1096 exact_row_predicate,
1097 Arc::default(),
1098 metrics,
1099 intra_file_repartitioning,
1100 ))
1101 }
1102
1103 fn session(batch_size: usize) -> SessionContext {
1104 SessionContext::new_with_config(SessionConfig::new().with_batch_size(batch_size))
1105 }
1106
1107 #[tokio::test]
1108 async fn provider_scopes_the_range_estimator_and_repartitioning_creates_a_metadata_cache()
1109 -> TestResult {
1110 let fixture = TestTable::partitioned("shared-range-read-estimator")?;
1111 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1112 let provider = DeltaTableProvider::try_new(table.clone(), ScanOptions::default())?;
1113 let separate_provider = DeltaTableProvider::try_new(table, ScanOptions::default())?;
1114 let context = SessionContext::new();
1115
1116 let first = provider.scan(&context.state(), None, &[], None).await?;
1117 let second = provider.scan(&context.state(), None, &[], None).await?;
1118 let separate = separate_provider
1119 .scan(&context.state(), None, &[], None)
1120 .await?;
1121 let first = first
1122 .as_ref()
1123 .downcast_ref::<DeltaScanExec>()
1124 .ok_or("first plan was not DeltaScanExec")?;
1125 let second = second
1126 .as_ref()
1127 .downcast_ref::<DeltaScanExec>()
1128 .ok_or("second plan was not DeltaScanExec")?;
1129 let separate = separate
1130 .as_ref()
1131 .downcast_ref::<DeltaScanExec>()
1132 .ok_or("separate provider plan was not DeltaScanExec")?;
1133 let repartitioned =
1134 first.with_repartitioned_partitions(first.reader_plan.partitions.clone());
1135 let repartitioned = repartitioned
1136 .as_ref()
1137 .downcast_ref::<DeltaScanExec>()
1138 .ok_or("repartitioned plan was not DeltaScanExec")?;
1139 assert!(Arc::ptr_eq(
1140 &first.range_read_estimator,
1141 &second.range_read_estimator
1142 ));
1143 assert!(Arc::ptr_eq(
1144 &first.range_read_estimator,
1145 &repartitioned.range_read_estimator
1146 ));
1147 assert!(!Arc::ptr_eq(
1148 &first.range_read_estimator,
1149 &separate.range_read_estimator
1150 ));
1151 assert!(repartitioned.file_tasks_repartitioned);
1152 assert!(repartitioned.parquet_metadata_cache.is_some());
1153 Ok(())
1154 }
1155
1156 fn sized_file_task(path: &str, size: Option<u64>) -> DeltaScanFileTask {
1157 use crate::{
1158 delta::kernel::KernelPhysicalToLogicalTransform,
1159 reader::deletion_vector::DeletionVectorMetadata,
1160 };
1161
1162 DeltaScanFileTask {
1163 path: path.to_owned(),
1164 file_size: size,
1165 parquet_byte_range: None,
1166 modification_time_ms: None,
1167 partition_values: Default::default(),
1168 deletion_vector: DeletionVectorMetadata::default(),
1169 transform: KernelPhysicalToLogicalTransform::default(),
1170 }
1171 }
1172
1173 fn partition(file_tasks: Vec<DeltaScanFileTask>) -> DeltaScanPartition {
1174 DeltaScanPartition { file_tasks }
1175 }
1176
1177 fn partition_estimated_bytes(partition: &DeltaScanPartition) -> Option<u64> {
1178 partition
1179 .file_tasks
1180 .iter()
1181 .map(DeltaScanFileTask::estimated_scan_bytes)
1182 .sum()
1183 }
1184
1185 #[allow(clippy::expect_used)]
1186 fn ids(batches: &[RecordBatch]) -> Vec<i32> {
1187 batches
1188 .iter()
1189 .flat_map(|batch| {
1190 batch
1191 .column(batch.schema().index_of("id").expect("id column"))
1192 .as_any()
1193 .downcast_ref::<Int32Array>()
1194 .expect("Int32 id")
1195 .values()
1196 .iter()
1197 .copied()
1198 .collect::<Vec<_>>()
1199 })
1200 .collect()
1201 }
1202
1203 #[test]
1204 fn file_repartitioning_balances_ranges_without_losing_file_identity() -> TestResult {
1205 let input = vec![partition(vec![
1206 sized_file_task("large.parquet", Some(100)),
1207 sized_file_task("small.parquet", Some(20)),
1208 ])];
1209 let partitions = repartition_file_tasks(&input, 4, 1, IntraFileRepartitioning::default())?
1210 .ok_or("not repartitioned")?;
1211
1212 assert_eq!(
1213 partitions
1214 .iter()
1215 .map(partition_estimated_bytes)
1216 .collect::<Vec<_>>(),
1217 vec![Some(30); 4]
1218 );
1219 let mut ranges = std::collections::BTreeMap::<_, Vec<_>>::new();
1220 for task in partitions
1221 .iter()
1222 .flat_map(|partition| &partition.file_tasks)
1223 {
1224 ranges.entry(task.path.as_str()).or_default().push(
1225 task.parquet_byte_range
1226 .clone()
1227 .unwrap_or(0..task.file_size.ok_or("missing file size")?),
1228 );
1229 assert_eq!(
1230 task.file_size,
1231 Some(if task.path == "large.parquet" {
1232 100
1233 } else {
1234 20
1235 })
1236 );
1237 if task.path == "small.parquet" {
1238 assert!(task.parquet_byte_range.is_none());
1239 }
1240 }
1241 assert_eq!(
1242 ranges.remove("large.parquet"),
1243 Some(vec![0..30, 30..60, 60..90, 90..100])
1244 );
1245 assert_eq!(
1246 ranges.remove("small.parquet"),
1247 Some(std::iter::once(0..20).collect())
1248 );
1249 assert!(ranges.is_empty());
1250
1251 Ok(())
1252 }
1253
1254 #[test]
1255 fn file_repartitioning_never_escapes_an_existing_range() -> TestResult {
1256 let mut task = sized_file_task("partial.parquet", Some(100));
1257 task.parquet_byte_range = Some(20..80);
1258 let partitions = repartition_file_tasks(
1259 &[partition(vec![task])],
1260 3,
1261 1,
1262 IntraFileRepartitioning::default(),
1263 )?
1264 .ok_or("partial range was not repartitioned")?;
1265
1266 assert_eq!(
1267 partitions
1268 .iter()
1269 .flat_map(|partition| &partition.file_tasks)
1270 .map(|task| task.parquet_byte_range.clone())
1271 .collect::<Vec<_>>(),
1272 [Some(20..40), Some(40..60), Some(60..80)]
1273 );
1274 assert!(
1275 partitions
1276 .iter()
1277 .flat_map(|partition| &partition.file_tasks)
1278 .all(|task| task.file_size == Some(100))
1279 );
1280
1281 Ok(())
1282 }
1283
1284 #[test]
1285 fn file_repartitioning_policy_controls_full_whole_file_plans() -> TestResult {
1286 let input = vec![
1287 partition(vec![sized_file_task("huge.parquet", Some(1_000))]),
1288 partition(vec![sized_file_task("small-0.parquet", Some(10))]),
1289 partition(vec![sized_file_task("small-1.parquet", Some(10))]),
1290 partition(vec![sized_file_task("small-2.parquet", Some(10))]),
1291 ];
1292
1293 assert!(
1294 repartition_file_tasks(&input, 4, 1, IntraFileRepartitioning::WhenBelowTarget)?
1295 .is_none()
1296 );
1297 let rebalanced = repartition_file_tasks(&input, 4, 1, IntraFileRepartitioning::Always)?
1298 .ok_or("full plan was not rebalanced")?;
1299 assert_eq!(
1300 rebalanced
1301 .iter()
1302 .map(partition_estimated_bytes)
1303 .collect::<Vec<_>>(),
1304 [Some(258), Some(258), Some(258), Some(256)]
1305 );
1306 assert!(
1307 rebalanced
1308 .iter()
1309 .flat_map(|partition| &partition.file_tasks)
1310 .any(|task| task.parquet_byte_range.is_some())
1311 );
1312 assert!(
1313 repartition_file_tasks(&input, 4, 1_031, IntraFileRepartitioning::Always)?.is_none()
1314 );
1315 assert!(repartition_file_tasks(&[], 4, 1, IntraFileRepartitioning::Always)?.is_none());
1316 assert!(
1317 input
1318 .iter()
1319 .flat_map(|partition| &partition.file_tasks)
1320 .all(|task| task.parquet_byte_range.is_none())
1321 );
1322 Ok(())
1323 }
1324
1325 #[test]
1326 fn file_repartitioning_refuses_unsupported_inputs_and_rejects_invalid_ranges() -> TestResult {
1327 let input = vec![partition(vec![sized_file_task("known.parquet", Some(120))])];
1328
1329 assert!(
1330 repartition_file_tasks(&input, 4, 121, IntraFileRepartitioning::default())?.is_none()
1331 );
1332 assert!(
1333 repartition_file_tasks(
1334 &[partition(vec![sized_file_task("unknown.parquet", None)])],
1335 4,
1336 1,
1337 IntraFileRepartitioning::default(),
1338 )?
1339 .is_none()
1340 );
1341 assert!(
1342 repartition_file_tasks(
1343 &[partition(vec![sized_file_task("empty.parquet", Some(0))])],
1344 4,
1345 1,
1346 IntraFileRepartitioning::default(),
1347 )?
1348 .is_none()
1349 );
1350 assert!(repartition_file_tasks(&input, 0, 1, IntraFileRepartitioning::default()).is_err());
1351
1352 for range in [90..110, std::ops::Range { start: 90, end: 80 }, 90..90] {
1353 let mut invalid = sized_file_task("invalid.parquet", Some(100));
1354 invalid.parquet_byte_range = Some(range);
1355 assert!(
1356 repartition_file_tasks(
1357 &[partition(vec![invalid])],
1358 4,
1359 1,
1360 IntraFileRepartitioning::default(),
1361 )
1362 .is_err()
1363 );
1364 }
1365
1366 let oversized = u64::try_from(i64::MAX)? + 1;
1367 assert!(
1368 repartition_file_tasks(
1369 &[partition(vec![sized_file_task(
1370 "oversized.parquet",
1371 Some(oversized),
1372 )])],
1373 2,
1374 1,
1375 IntraFileRepartitioning::default(),
1376 )
1377 .is_err()
1378 );
1379
1380 Ok(())
1381 }
1382
1383 #[test]
1384 fn task_from_partitioned_file_rejects_malformed_partitioner_output() {
1385 let mut missing_extension = PartitionedFile::new("missing-extension.parquet", 100);
1386 missing_extension.range = Some(FileRange { start: 0, end: 100 });
1387 assert!(task_from_partitioned_file(missing_extension).is_err());
1388
1389 let missing_range = PartitionedFile::new("missing-range.parquet", 100)
1390 .with_extension(sized_file_task("missing-range.parquet", Some(100)));
1391 assert!(task_from_partitioned_file(missing_range).is_err());
1392
1393 for range in [
1394 FileRange {
1395 start: -1,
1396 end: 100,
1397 },
1398 FileRange { start: 50, end: 50 },
1399 FileRange { start: 0, end: 101 },
1400 ] {
1401 let mut invalid = PartitionedFile::new("invalid-range.parquet", 100)
1402 .with_extension(sized_file_task("invalid-range.parquet", Some(100)));
1403 invalid.range = Some(range);
1404 assert!(task_from_partitioned_file(invalid).is_err());
1405 }
1406
1407 let mut missing_size = PartitionedFile::new("missing-size.parquet", 100)
1408 .with_extension(sized_file_task("missing-size.parquet", None));
1409 missing_size.range = Some(FileRange { start: 0, end: 100 });
1410 assert!(task_from_partitioned_file(missing_size).is_err());
1411 }
1412
1413 #[tokio::test]
1414 async fn direct_repartitioning_reads_each_row_once_and_preserves_byte_accounting() -> TestResult
1415 {
1416 let fixture = TestTable::partitioned("file-repartitioning")?;
1417 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1418 let plan = build_plan_with_repartitioning(
1419 &table,
1420 None,
1421 &[],
1422 4,
1423 DeltaScanExecutionOptions::new(),
1424 None,
1425 IntraFileRepartitioning::Always,
1426 )?;
1427 let expected_bytes = collect_scan_metrics(plan.as_ref())[0]
1428 .snapshot()
1429 .reader_metrics
1430 .estimated_input_bytes;
1431 let mut config = ConfigOptions::new();
1432 config.optimizer.repartition_file_min_size = 1;
1433 let repartitioned = plan
1434 .repartitioned(4, &config)?
1435 .ok_or("Direct backend scan was not repartitioned")?;
1436 assert!(repartitioned.repartitioned(4, &config)?.is_none());
1437
1438 assert_eq!(
1439 repartitioned
1440 .properties()
1441 .output_partitioning()
1442 .partition_count(),
1443 4
1444 );
1445 let mut actual_ids = ids(&datafusion::physical_plan::collect(
1446 Arc::clone(&repartitioned),
1447 session(1024).task_ctx(),
1448 )
1449 .await?);
1450 actual_ids.sort_unstable();
1451 assert_eq!(actual_ids, [1, 2, 3, 4]);
1452 let metrics = collect_scan_metrics(repartitioned.as_ref())[0].snapshot();
1453 assert_eq!(metrics.reader_metrics.scan_partitions_planned, 4);
1454 assert_eq!(
1455 metrics.reader_metrics.estimated_parquet_task_bytes_admitted,
1456 expected_bytes
1457 );
1458
1459 let explicit_one =
1460 build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
1461 assert!(explicit_one.repartitioned(4, &config)?.is_none());
1462
1463 {
1464 let kernel = build_plan_with_repartitioning(
1465 &table,
1466 None,
1467 &[],
1468 4,
1469 DeltaScanExecutionOptions::new()
1470 .with_parquet_backend(ParquetReaderBackend::DeltaKernel),
1471 None,
1472 IntraFileRepartitioning::Always,
1473 )?;
1474 assert!(kernel.repartitioned(4, &config)?.is_none());
1475 }
1476
1477 Ok(())
1478 }
1479
1480 fn dynamic_filter(name: &str, index: usize) -> Arc<DynamicFilterPhysicalExpr> {
1481 Arc::new(DynamicFilterPhysicalExpr::new(
1482 vec![Arc::new(Column::new(name, index))],
1483 physical_lit(true),
1484 ))
1485 }
1486
1487 fn hook_input(
1488 filters: Vec<Arc<dyn datafusion::physical_plan::PhysicalExpr>>,
1489 ) -> ChildPushdownResult {
1490 ChildPushdownResult {
1491 parent_filters: filters
1492 .into_iter()
1493 .map(|filter| ChildFilterPushdownResult {
1494 filter,
1495 child_results: Vec::new(),
1496 })
1497 .collect(),
1498 self_filters: Vec::new(),
1499 }
1500 }
1501
1502 #[tokio::test]
1503 async fn properties_projection_partitions_metrics_and_reexecution_match_provider_behavior()
1504 -> TestResult {
1505 let fixture = TestTable::partitioned("properties")?;
1506 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1507 let logical_filter = col("id").gt(lit(1_i32));
1508 let plan = build_plan(
1509 &table,
1510 Some(&[1, 0]),
1511 &[logical_filter],
1512 2,
1513 DeltaScanExecutionOptions::new(),
1514 None,
1515 )?;
1516
1517 assert_eq!(plan.name(), "DeltaScanExec");
1518 assert!(plan.children().is_empty());
1519 assert!(plan.metrics().is_none());
1520 assert_eq!(plan.schema().fields().len(), 2);
1521 assert_eq!(plan.schema().field(0).name(), "region");
1522 assert_eq!(plan.schema().field(1).name(), "id");
1523 assert_eq!(plan.properties().output_partitioning().partition_count(), 2);
1524 assert_eq!(
1525 plan.partition_statistics(None)?,
1526 Arc::new(datafusion::common::Statistics::new_unknown(&plan.schema()))
1527 );
1528 let context = session(1);
1529 let first =
1530 datafusion::physical_plan::collect(Arc::clone(&plan), context.task_ctx()).await?;
1531 assert!(first.iter().all(|batch| batch.num_rows() <= 1));
1532 let mut first_ids = ids(&first);
1533 first_ids.sort_unstable();
1534 assert_eq!(first_ids, [2, 3, 4]);
1535
1536 let second =
1537 datafusion::physical_plan::collect(Arc::clone(&plan), context.task_ctx()).await?;
1538 let mut second_ids = ids(&second);
1539 second_ids.sort_unstable();
1540 assert_eq!(second_ids, first_ids);
1541 let handles = collect_scan_metrics(plan.as_ref());
1542 assert_eq!(handles.len(), 1);
1543 assert_eq!(handles[0].registration_name(), None);
1544 let metrics = handles[0].snapshot();
1545 assert_eq!(metrics.configured_batch_size_rows, Some(1));
1546 assert_eq!(metrics.reader_metrics.scan_partitions_started, 4);
1547 assert_eq!(metrics.reader_metrics.file_tasks_completed, 4);
1548 assert_eq!(metrics.reader_metrics.scheduler_rows_emitted, 6);
1549
1550 let hidden = build_plan(
1551 &table,
1552 Some(&[1]),
1553 &[col("id").gt(lit(1_i32))],
1554 1,
1555 DeltaScanExecutionOptions::new(),
1556 None,
1557 )?;
1558 let hidden_batches = datafusion::physical_plan::collect(
1559 Arc::clone(&hidden),
1560 SessionContext::new().task_ctx(),
1561 )
1562 .await?;
1563 assert_eq!(hidden.schema().fields().len(), 1);
1564 assert_eq!(hidden.schema().field(0).name(), "region");
1565 assert!(hidden_batches.iter().all(|batch| batch.num_columns() == 1));
1566 assert_eq!(
1567 hidden_batches
1568 .iter()
1569 .map(RecordBatch::num_rows)
1570 .sum::<usize>(),
1571 3
1572 );
1573
1574 let partition_filter = build_plan(
1575 &table,
1576 None,
1577 &[col("region").eq(lit("west"))],
1578 2,
1579 DeltaScanExecutionOptions::new(),
1580 None,
1581 )?;
1582 let partition_batches = datafusion::physical_plan::collect(
1583 Arc::clone(&partition_filter),
1584 SessionContext::new().task_ctx(),
1585 )
1586 .await?;
1587 assert_eq!(ids(&partition_batches), [1, 2]);
1588 assert_eq!(
1589 collect_scan_metrics(partition_filter.as_ref())[0]
1590 .snapshot()
1591 .reader_metrics
1592 .file_tasks_started,
1593 1
1594 );
1595
1596 let empty = build_plan(
1597 &table,
1598 Some(&[]),
1599 &[],
1600 1,
1601 DeltaScanExecutionOptions::new(),
1602 None,
1603 )?;
1604 let empty_batches = datafusion::physical_plan::collect(
1605 Arc::clone(&empty),
1606 SessionContext::new().task_ctx(),
1607 )
1608 .await?;
1609 assert!(empty.schema().fields().is_empty());
1610 assert!(empty_batches.iter().all(|batch| batch.num_columns() == 0));
1611 assert_eq!(
1612 empty_batches
1613 .iter()
1614 .map(RecordBatch::num_rows)
1615 .sum::<usize>(),
1616 4
1617 );
1618
1619 let empty_fixture = TestTable::empty("empty-scan")?;
1620 let empty_table = DeltaTableBuilder::new(empty_fixture.uri())
1621 .load_table()
1622 .await?;
1623 let empty_plan = build_plan(
1624 &empty_table,
1625 None,
1626 &[],
1627 1,
1628 DeltaScanExecutionOptions::new(),
1629 None,
1630 )?;
1631 assert_eq!(
1632 empty_plan
1633 .properties()
1634 .output_partitioning()
1635 .partition_count(),
1636 1
1637 );
1638 assert!(
1639 datafusion::physical_plan::collect(empty_plan, SessionContext::new().task_ctx(),)
1640 .await?
1641 .is_empty()
1642 );
1643
1644 let invalid = plan.execute(2, context.task_ctx());
1645 let error = match invalid {
1646 Ok(_) => return Err("out-of-range partition unexpectedly executed".into()),
1647 Err(error) => error,
1648 };
1649 let DataFusionError::External(source) = error else {
1650 return Err("invalid partition did not preserve the reader error".into());
1651 };
1652 let reader = source
1653 .downcast_ref::<DeltaReaderError>()
1654 .ok_or("external error was not DeltaReaderError")?;
1655 assert_eq!(reader.code(), "datafusion_adapter");
1656 Ok(())
1657 }
1658
1659 #[tokio::test]
1660 async fn dynamic_filter_hook_prunes_before_file_start_and_counts_once() -> TestResult {
1661 let fixture = TestTable::partitioned("dynamic")?;
1662 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1663 let plan = build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
1664 let dynamic = dynamic_filter("region", 1);
1665 let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic.clone();
1666 let rejected: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic_filter("id", 0);
1667 let pushed = plan.handle_child_pushdown_result(
1668 FilterPushdownPhase::Post,
1669 hook_input(vec![physical, rejected]),
1670 &ConfigOptions::new(),
1671 )?;
1672 assert!(matches!(
1673 pushed.filters.as_slice(),
1674 [PushedDown::Yes, PushedDown::No]
1675 ));
1676 let updated = pushed.updated_node.ok_or("dynamic plan was not retained")?;
1677 dynamic.update(Arc::new(BinaryExpr::new(
1678 Arc::new(Column::new("region", 1)),
1679 Operator::Eq,
1680 physical_lit("west"),
1681 )))?;
1682
1683 let batches = datafusion::physical_plan::collect(
1684 Arc::clone(&updated),
1685 SessionContext::new().task_ctx(),
1686 )
1687 .await?;
1688 assert_eq!(ids(&batches), [1, 2]);
1689 let metrics = collect_scan_metrics(updated.as_ref())
1690 .pop()
1691 .ok_or("missing dynamic metrics")?
1692 .snapshot();
1693 assert_eq!(metrics.dynamic_filters_received, 2);
1694 assert_eq!(metrics.dynamic_filters_accepted, 1);
1695 assert_eq!(metrics.dynamic_filters_rejected, 1);
1696 assert_eq!(metrics.dynamic_partition_filter_checks, 2);
1697 assert_eq!(metrics.dynamic_partition_tasks_pruned, 1);
1698 assert_eq!(metrics.dynamic_partition_tasks_kept, 1);
1699 assert_eq!(metrics.reader_metrics.file_tasks_started, 1);
1700 assert_eq!(metrics.reader_metrics.file_tasks_completed, 1);
1701 assert_eq!(
1702 collect_scan_metrics(plan.as_ref())[0]
1703 .snapshot()
1704 .dynamic_filters_received,
1705 2
1706 );
1707 Ok(())
1708 }
1709
1710 #[tokio::test]
1711 async fn physical_pushdown_preserves_dynamic_filters_across_plan_rebuild() -> TestResult {
1712 let fixture = TestTable::partitioned("dynamic-plan-rebuild")?;
1713 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1714 let plan = build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
1715 let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> =
1716 dynamic_filter("region", 1);
1717 let pushed = plan.handle_child_pushdown_result(
1718 FilterPushdownPhase::Post,
1719 hook_input(vec![physical]),
1720 &ConfigOptions::new(),
1721 )?;
1722 let updated = pushed.updated_node.ok_or("expected updated scan")?;
1723 let rebuilt = Arc::clone(&updated).with_new_children(Vec::new())?;
1724 let reset = updated.reset_state()?;
1725
1726 for candidate in [&rebuilt, &reset] {
1727 let debug = format!("{candidate:?}");
1728 assert!(debug.contains("dynamic_filter_count: 1"), "{debug}");
1729 }
1730 let display = datafusion::physical_plan::displayable(rebuilt.as_ref())
1731 .one_line()
1732 .to_string();
1733 assert!(display.contains("DeltaScanExec:"), "{display}");
1734 assert!(display.contains("partitions="), "{display}");
1735 assert!(!display.contains("DynamicFilter"), "{display}");
1736 assert!(
1737 Arc::clone(&rebuilt)
1738 .with_new_children(vec![Arc::clone(&rebuilt)])
1739 .is_err()
1740 );
1741 Ok(())
1742 }
1743
1744 #[tokio::test]
1745 async fn late_dynamic_filter_keeps_admitted_file_and_prunes_the_next() -> TestResult {
1746 let fixture = TestTable::late_dynamic("late-dynamic")?;
1747 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1748 let options = DeltaScanExecutionOptions::new()
1749 .with_prefetch_files_per_partition(0)
1750 .with_max_concurrent_file_reads_per_partition(1)?
1751 .with_max_concurrent_file_reads_per_scan(Some(1))?
1752 .with_output_buffer_batches_per_partition(1)?;
1753 let plan = build_plan(&table, None, &[], 1, options, None)?;
1754 let dynamic = dynamic_filter("region", 1);
1755 let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic.clone();
1756 let pushed = plan.handle_child_pushdown_result(
1757 FilterPushdownPhase::Post,
1758 hook_input(vec![physical]),
1759 &ConfigOptions::new(),
1760 )?;
1761 let updated = pushed.updated_node.ok_or("dynamic plan was not retained")?;
1762 let mut stream = updated.execute(0, session(1).task_ctx())?;
1763 let first = stream.next().await.ok_or("missing first batch")??;
1764 assert_eq!(ids(std::slice::from_ref(&first)), [1]);
1765
1766 dynamic.update(Arc::new(BinaryExpr::new(
1767 Arc::new(Column::new("region", 1)),
1768 Operator::Eq,
1769 physical_lit("none"),
1770 )))?;
1771 let mut batches = vec![first];
1772 while let Some(batch) = stream.next().await {
1773 batches.push(batch?);
1774 }
1775
1776 assert_eq!(ids(&batches), [1, 2, 3]);
1777 let metrics = collect_scan_metrics(updated.as_ref())[0].snapshot();
1778 assert_eq!(metrics.dynamic_partition_filter_checks, 2);
1779 assert_eq!(metrics.dynamic_partition_tasks_kept, 1);
1780 assert_eq!(metrics.dynamic_partition_tasks_pruned, 1);
1781 assert_eq!(metrics.reader_metrics.file_tasks_started, 1);
1782 assert_eq!(metrics.reader_metrics.file_tasks_completed, 1);
1783 Ok(())
1784 }
1785
1786 #[tokio::test]
1787 async fn hook_is_post_only_empty_safe_and_collector_is_ordered_and_distinct() -> TestResult {
1788 let fixture = TestTable::partitioned("collector")?;
1789 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1790 let first = build_plan(
1791 &table,
1792 None,
1793 &[],
1794 1,
1795 DeltaScanExecutionOptions::new(),
1796 Some("first".to_owned()),
1797 )?;
1798 let second = build_plan(
1799 &table,
1800 None,
1801 &[],
1802 1,
1803 DeltaScanExecutionOptions::new(),
1804 Some("second".to_owned()),
1805 )?;
1806 let dynamic = dynamic_filter("region", 1);
1807 let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic;
1808 let pre = first.handle_child_pushdown_result(
1809 FilterPushdownPhase::Pre,
1810 hook_input(vec![physical]),
1811 &ConfigOptions::new(),
1812 )?;
1813 assert!(pre.updated_node.is_none());
1814 assert!(matches!(pre.filters.as_slice(), [PushedDown::No]));
1815 let empty = first.handle_child_pushdown_result(
1816 FilterPushdownPhase::Post,
1817 hook_input(Vec::new()),
1818 &ConfigOptions::new(),
1819 )?;
1820 assert!(empty.updated_node.is_none());
1821
1822 let union: Arc<dyn ExecutionPlan> = UnionExec::try_new(vec![
1823 Arc::clone(&first),
1824 Arc::clone(&second),
1825 Arc::clone(&first),
1826 ])?;
1827 let handles = collect_scan_metrics(union.as_ref());
1828 assert_eq!(handles.len(), 2);
1829 assert_eq!(handles[0].registration_name(), Some("first"));
1830 assert_eq!(handles[1].registration_name(), Some("second"));
1831 assert!(!format!("{:?}", handles[0]).contains("first"));
1832 let initial = handles[0].snapshot();
1833 assert_eq!(initial.configured_batch_size_rows, None);
1834 assert_eq!(
1835 [
1836 initial.dynamic_partition_tasks_pruned,
1837 initial.dynamic_partition_tasks_kept,
1838 initial.dynamic_filters_received,
1839 initial.dynamic_filters_accepted,
1840 initial.dynamic_filters_rejected,
1841 initial.dynamic_partition_filter_checks,
1842 initial.dynamic_partition_tasks_kept_unusable_metadata,
1843 initial.dynamic_partition_tasks_kept_unevaluable_filter,
1844 ],
1845 [0; 8]
1846 );
1847
1848 let accepted: Arc<dyn datafusion::physical_plan::PhysicalExpr> =
1849 dynamic_filter("region", 1);
1850 let updated = first
1851 .handle_child_pushdown_result(
1852 FilterPushdownPhase::Post,
1853 hook_input(vec![accepted]),
1854 &ConfigOptions::new(),
1855 )?
1856 .updated_node
1857 .ok_or("expected updated scan")?;
1858 let shared_metrics_union: Arc<dyn ExecutionPlan> =
1859 UnionExec::try_new(vec![updated, Arc::clone(&first), Arc::clone(&second)])?;
1860 let shared_handles = collect_scan_metrics(shared_metrics_union.as_ref());
1861 assert_eq!(shared_handles.len(), 2);
1862 assert_eq!(shared_handles[0].registration_name(), Some("first"));
1863 assert_eq!(shared_handles[1].registration_name(), Some("second"));
1864 assert_eq!(handles[0].identity(), shared_handles[0].identity());
1865 assert_ne!(handles[0].identity(), shared_handles[1].identity());
1866
1867 drop(shared_metrics_union);
1868 drop(union);
1869 drop(first);
1870 drop(second);
1871 assert_eq!(handles[0].snapshot().reader_metrics.file_tasks_started, 0);
1872 assert_eq!(handles[0].registration_name(), Some("first"));
1873 Ok(())
1874 }
1875
1876 #[tokio::test]
1877 async fn dynamic_partition_admission_reason_counts_are_once_per_file_and_saturating()
1878 -> TestResult {
1879 use crate::{
1880 delta::kernel::KernelPhysicalToLogicalTransform,
1881 reader::datafusion::dynamic_filters::DynamicFilterClassification,
1882 reader::deletion_vector::DeletionVectorMetadata,
1883 };
1884
1885 let fixture = TestTable::partitioned("dynamic-counters")?;
1886 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1887 let plan = build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
1888 let metrics = collect_scan_metrics(plan.as_ref())
1889 .pop()
1890 .ok_or("missing metrics")?;
1891 let schema = Arc::new(Schema::new(vec![
1892 Field::new("id", DataType::Int32, false),
1893 Field::new("region", DataType::Utf8, true),
1894 ]));
1895 let retained = |dynamic: Arc<DynamicFilterPhysicalExpr>| -> TestResult<_> {
1896 let physical: Arc<dyn datafusion::physical_plan::PhysicalExpr> = dynamic;
1897 Ok(DynamicFilterClassification::from_filters(
1898 std::slice::from_ref(&physical),
1899 &schema,
1900 &["region".to_owned()],
1901 )
1902 .accepted_filters()
1903 .next()
1904 .cloned()
1905 .ok_or("dynamic filter was not retained")?)
1906 };
1907 let first = retained(dynamic_filter("region", 1))?;
1908 let second = retained(dynamic_filter("region", 1))?;
1909 let missing = DeltaScanFileTask {
1910 path: "missing-partition.parquet".to_owned(),
1911 file_size: None,
1912 parquet_byte_range: None,
1913 modification_time_ms: None,
1914 partition_values: Default::default(),
1915 deletion_vector: DeletionVectorMetadata::default(),
1916 transform: KernelPhysicalToLogicalTransform::default(),
1917 };
1918 assert_eq!(
1919 dynamic_partition_admission_policy(metrics.clone(), Arc::from([first, second]))(
1920 &missing,
1921 )?,
1922 FileAdmissionDecision::Admit
1923 );
1924 let snapshot = metrics.snapshot();
1925 assert_eq!(snapshot.dynamic_partition_filter_checks, 2);
1926 assert_eq!(snapshot.dynamic_partition_tasks_kept, 1);
1927 assert_eq!(snapshot.dynamic_partition_tasks_kept_unusable_metadata, 1);
1928
1929 let rejecting = dynamic_filter("region", 1);
1930 rejecting.update(physical_lit(false))?;
1931 let first = retained(rejecting)?;
1932 let second = retained(dynamic_filter("region", 1))?;
1933 let mut present = missing.clone();
1934 present
1935 .partition_values
1936 .insert("region".to_owned(), "west".to_owned());
1937 assert_eq!(
1938 dynamic_partition_admission_policy(metrics.clone(), Arc::from([first, second]))(
1939 &present,
1940 )?,
1941 FileAdmissionDecision::Skip
1942 );
1943 let snapshot = metrics.snapshot();
1944 assert_eq!(snapshot.dynamic_partition_filter_checks, 3);
1945 assert_eq!(snapshot.dynamic_partition_tasks_pruned, 1);
1946 assert_eq!(snapshot.dynamic_partition_tasks_kept, 1);
1947 assert_eq!(snapshot.dynamic_partition_tasks_kept_unusable_metadata, 1);
1948
1949 let unsupported = dynamic_filter("region", 1);
1950 unsupported.update(physical_lit("not boolean"))?;
1951 let admission = dynamic_partition_admission_policy(
1952 metrics.clone(),
1953 Arc::from([retained(unsupported)?]),
1954 );
1955 assert_eq!(admission(&present)?, FileAdmissionDecision::Admit);
1956 let snapshot = metrics.snapshot();
1957 assert_eq!(snapshot.dynamic_partition_filter_checks, 4);
1958 assert_eq!(snapshot.dynamic_partition_tasks_kept, 2);
1959 assert_eq!(snapshot.dynamic_partition_tasks_kept_unevaluable_filter, 1);
1960
1961 metrics
1962 .inner
1963 .dynamic_filters_received
1964 .store(u64::MAX - 1, Ordering::Relaxed);
1965 metrics.record_dynamic_filters_received(2);
1966 metrics.record_dynamic_filters_received(1);
1967 assert_eq!(metrics.snapshot().dynamic_filters_received, u64::MAX);
1968 Ok(())
1969 }
1970
1971 #[tokio::test]
1972 async fn dynamic_metrics_updates_are_thread_safe() -> TestResult {
1973 const THREADS: usize = 4;
1974 const ITERATIONS: usize = 100;
1975
1976 let fixture = TestTable::partitioned("dynamic-metrics-concurrency")?;
1977 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
1978 let plan = build_plan(&table, None, &[], 1, DeltaScanExecutionOptions::new(), None)?;
1979 let metrics = collect_scan_metrics(plan.as_ref())
1980 .pop()
1981 .ok_or("missing metrics")?;
1982 let mut handles = Vec::new();
1983
1984 for _ in 0..THREADS {
1985 let metrics = metrics.clone();
1986 handles.push(thread::spawn(move || {
1987 for _ in 0..ITERATIONS {
1988 metrics.record_dynamic_partition_task_pruned();
1989 metrics.record_dynamic_partition_task_kept();
1990 metrics.record_dynamic_filters_received(3);
1991 metrics.record_dynamic_filters_accepted(1);
1992 metrics.record_dynamic_filters_rejected(2);
1993 metrics.record_dynamic_partition_filter_check();
1994 metrics.record_dynamic_partition_task_kept_unusable_metadata();
1995 metrics.record_dynamic_partition_task_kept_unevaluable_filter();
1996 }
1997 }));
1998 }
1999 for handle in handles {
2000 handle.join().map_err(|_| "metrics worker panicked")?;
2001 }
2002
2003 let calls = u64::try_from(THREADS * ITERATIONS)?;
2004 let snapshot = metrics.snapshot();
2005 assert_eq!(snapshot.dynamic_partition_tasks_pruned, calls);
2006 assert_eq!(snapshot.dynamic_partition_tasks_kept, calls);
2007 assert_eq!(snapshot.dynamic_filters_received, calls * 3);
2008 assert_eq!(snapshot.dynamic_filters_accepted, calls);
2009 assert_eq!(snapshot.dynamic_filters_rejected, calls * 2);
2010 assert_eq!(snapshot.dynamic_partition_filter_checks, calls);
2011 assert_eq!(
2012 snapshot.dynamic_partition_tasks_kept_unusable_metadata,
2013 calls
2014 );
2015 assert_eq!(
2016 snapshot.dynamic_partition_tasks_kept_unevaluable_filter,
2017 calls
2018 );
2019 Ok(())
2020 }
2021
2022 #[tokio::test]
2023 async fn execution_error_and_stream_drop_preserve_partial_metrics() -> TestResult {
2024 let missing_fixture = TestTable::missing("error")?;
2025 let missing_table = DeltaTableBuilder::new(missing_fixture.uri())
2026 .load_table()
2027 .await?;
2028 let missing_plan = build_plan(
2029 &missing_table,
2030 None,
2031 &[],
2032 1,
2033 DeltaScanExecutionOptions::new(),
2034 None,
2035 )?;
2036 let result = datafusion::physical_plan::collect(
2037 Arc::clone(&missing_plan),
2038 SessionContext::new().task_ctx(),
2039 )
2040 .await;
2041 let error = result.expect_err("missing file must fail");
2042 assert!(matches!(&error, DataFusionError::External(_)));
2043 assert!(!error.to_string().contains("missing.parquet"));
2044 let failed = collect_scan_metrics(missing_plan.as_ref())
2045 .pop()
2046 .ok_or("missing failure metrics")?;
2047 assert_eq!(failed.snapshot().reader_metrics.file_tasks_started, 1);
2048
2049 let fixture = TestTable::partitioned("drop")?;
2050 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
2051 let options = DeltaScanExecutionOptions::new()
2052 .with_prefetch_files_per_partition(0)
2053 .with_max_concurrent_file_reads_per_partition(1)?
2054 .with_max_concurrent_file_reads_per_scan(Some(1))?
2055 .with_output_buffer_batches_per_partition(1)?;
2056 let drop_plan = build_plan(&table, None, &[], 1, options, None)?;
2057 let handle = collect_scan_metrics(drop_plan.as_ref())
2058 .pop()
2059 .ok_or("missing drop metrics")?;
2060 let mut stream = drop_plan.execute(0, SessionContext::new().task_ctx())?;
2061 assert!(stream.next().await.transpose()?.is_some());
2062 drop(stream);
2063 tokio::task::yield_now().await;
2064 let stable = handle.snapshot();
2065 tokio::task::yield_now().await;
2066 assert_eq!(handle.snapshot(), stable);
2067 assert!(stable.reader_metrics.file_tasks_started >= 1);
2068 let retry = datafusion::physical_plan::collect(
2069 Arc::clone(&drop_plan),
2070 SessionContext::new().task_ctx(),
2071 )
2072 .await?;
2073 assert_eq!(ids(&retry), [1, 2, 3, 4]);
2074 Ok(())
2075 }
2076
2077 #[tokio::test]
2078 async fn parquet_backends_produce_the_same_logical_rows() -> TestResult {
2079 let fixture = TestTable::partitioned("backends")?;
2080 let table = DeltaTableBuilder::new(fixture.uri()).load_table().await?;
2081 let mut outputs = Vec::new();
2082 for backend in [
2083 ParquetReaderBackend::Direct,
2084 ParquetReaderBackend::DeltaKernel,
2085 ] {
2086 let options = DeltaScanExecutionOptions::new().with_parquet_backend(backend);
2087 let plan = build_plan(&table, Some(&[1, 0]), &[], 2, options, None)?;
2088 let mut batches =
2089 datafusion::physical_plan::collect(plan, SessionContext::new().task_ctx()).await?;
2090 batches.sort_by_key(|batch| {
2091 batch
2092 .column(1)
2093 .as_any()
2094 .downcast_ref::<Int32Array>()
2095 .expect("Int32 id")
2096 .value(0)
2097 });
2098 outputs.push(
2099 batches
2100 .iter()
2101 .flat_map(|batch| {
2102 let ids = batch
2103 .column(1)
2104 .as_any()
2105 .downcast_ref::<Int32Array>()
2106 .expect("Int32 id");
2107 let regions = batch
2108 .column(0)
2109 .as_any()
2110 .downcast_ref::<StringArray>()
2111 .expect("Utf8 region");
2112 (0..batch.num_rows())
2113 .map(|row| (regions.value(row).to_owned(), ids.value(row)))
2114 .collect::<Vec<_>>()
2115 })
2116 .collect::<Vec<_>>(),
2117 );
2118 }
2119 assert_eq!(outputs[0], outputs[1]);
2120
2121 let kernel_options = DeltaScanExecutionOptions::new()
2122 .with_parquet_backend(ParquetReaderBackend::DeltaKernel);
2123 let inexact = build_plan(
2124 &table,
2125 None,
2126 &[col("id").gt(lit(1_i32))],
2127 1,
2128 kernel_options,
2129 None,
2130 )?;
2131 let unfiltered = datafusion::physical_plan::collect(
2132 Arc::clone(&inexact),
2133 SessionContext::new().task_ctx(),
2134 )
2135 .await?;
2136 assert_eq!(ids(&unfiltered), [1, 2, 3, 4]);
2137
2138 let residual: Arc<dyn datafusion::physical_plan::PhysicalExpr> = Arc::new(BinaryExpr::new(
2139 Arc::new(Column::new("id", 0)),
2140 Operator::Gt,
2141 physical_lit(1_i32),
2142 ));
2143 let residual_plan: Arc<dyn ExecutionPlan> =
2144 Arc::new(FilterExec::try_new(residual, inexact)?);
2145 let filtered =
2146 datafusion::physical_plan::collect(residual_plan, SessionContext::new().task_ctx())
2147 .await?;
2148 assert_eq!(ids(&filtered), [2, 3, 4]);
2149 Ok(())
2150 }
2151}