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 => {
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 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}