1pub(crate) mod sort_pushdown;
22
23#[cfg(feature = "proto")]
28mod proto;
29
30use crate::file_groups::FileGroup;
31use crate::{
32 PartitionedFile, display::FileGroupsDisplay, file::FileSource,
33 file_compression_type::FileCompressionType, file_stream::FileStreamBuilder,
34 file_stream::work_source::SharedWorkSource, source::DataSource,
35 statistics::MinMaxStatistics,
36};
37use arrow::datatypes::Fields;
38use arrow::datatypes::{DataType, Schema, SchemaRef};
39use datafusion_common::config::ConfigOptions;
40use datafusion_common::tree_node::TreeNodeRecursion;
41use datafusion_common::{
42 Constraints, Result, ScalarValue, Statistics, internal_datafusion_err, internal_err,
43};
44use datafusion_execution::{
45 SendableRecordBatchStream, TaskContext, object_store::ObjectStoreUrl,
46};
47use datafusion_expr::Operator;
48
49use crate::source::OpenArgs;
50use datafusion_common::stats::Precision;
51use datafusion_physical_expr::expressions::{BinaryExpr, Column};
52use datafusion_physical_expr::projection::{ProjectionExprs, ProjectionMapping};
53use datafusion_physical_expr::utils::reassign_expr_columns;
54use datafusion_physical_expr::{EquivalenceProperties, Partitioning, split_conjunction};
55use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
56use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, is_volatile};
57use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
58use datafusion_physical_plan::SortOrderPushdownResult;
59use datafusion_physical_plan::coop::cooperative;
60use datafusion_physical_plan::execution_plan::SchedulingType;
61use datafusion_physical_plan::{
62 DisplayAs, DisplayFormatType,
63 display::{ProjectSchemaDisplay, display_orderings},
64 filter_pushdown::FilterPushdownPropagation,
65 metrics::ExecutionPlanMetricsSet,
66};
67use log::{debug, warn};
68use std::any::Any;
69use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc};
70
71#[derive(Clone)]
152pub struct FileScanConfig {
153 pub object_store_url: ObjectStoreUrl,
165 pub file_groups: Vec<FileGroup>,
181 pub constraints: Constraints,
183 pub limit: Option<usize>,
186 pub preserve_order: bool,
191 pub output_ordering: Vec<LexOrdering>,
199 pub file_compression_type: FileCompressionType,
201 pub file_source: Arc<dyn FileSource>,
203 pub batch_size: Option<usize>,
206 pub expr_adapter_factory: Option<Arc<dyn PhysicalExprAdapterFactory>>,
209 pub(crate) statistics: Statistics,
219 pub output_partitioning: Option<Partitioning>,
226}
227
228#[derive(Clone)]
285pub struct FileScanConfigBuilder {
286 object_store_url: ObjectStoreUrl,
287 file_source: Arc<dyn FileSource>,
288 limit: Option<usize>,
289 preserve_order: bool,
290 constraints: Option<Constraints>,
291 file_groups: Vec<FileGroup>,
292 statistics: Option<Statistics>,
293 output_ordering: Vec<LexOrdering>,
294 output_partitioning: Option<Partitioning>,
295 file_compression_type: Option<FileCompressionType>,
296 batch_size: Option<usize>,
297 expr_adapter_factory: Option<Arc<dyn PhysicalExprAdapterFactory>>,
298}
299
300impl FileScanConfigBuilder {
301 pub fn new(
308 object_store_url: ObjectStoreUrl,
309 file_source: Arc<dyn FileSource>,
310 ) -> Self {
311 Self {
312 object_store_url,
313 file_source,
314 file_groups: vec![],
315 statistics: None,
316 output_ordering: vec![],
317 output_partitioning: None,
318 file_compression_type: None,
319 limit: None,
320 preserve_order: false,
321 constraints: None,
322 batch_size: None,
323 expr_adapter_factory: None,
324 }
325 }
326
327 pub fn with_limit(mut self, limit: Option<usize>) -> Self {
331 self.limit = limit;
332 self
333 }
334
335 pub fn with_preserve_order(mut self, order_sensitive: bool) -> Self {
342 self.preserve_order = order_sensitive;
343 self
344 }
345
346 pub fn with_source(mut self, file_source: Arc<dyn FileSource>) -> Self {
351 self.file_source = file_source;
352 self
353 }
354
355 pub fn table_schema(&self) -> &SchemaRef {
357 self.file_source.table_schema().table_schema()
358 }
359
360 #[deprecated(since = "51.0.0", note = "Use with_projection_indices instead")]
366 pub fn with_projection(self, indices: Option<Vec<usize>>) -> Self {
367 match self.clone().with_projection_indices(indices) {
368 Ok(builder) => builder,
369 Err(e) => {
370 warn!(
371 "Failed to push down projection in FileScanConfigBuilder::with_projection: {e}"
372 );
373 self
374 }
375 }
376 }
377
378 pub fn with_projection_indices(
387 mut self,
388 indices: Option<Vec<usize>>,
389 ) -> Result<Self> {
390 let projection_exprs = indices.map(|indices| {
391 ProjectionExprs::from_indices(
392 &indices,
393 self.file_source.table_schema().table_schema(),
394 )
395 });
396 let Some(projection_exprs) = projection_exprs else {
397 return Ok(self);
398 };
399 let new_source = self
400 .file_source
401 .try_pushdown_projection(&projection_exprs)
402 .map_err(|e| {
403 internal_datafusion_err!(
404 "Failed to push down projection in FileScanConfigBuilder::build: {e}"
405 )
406 })?;
407 if let Some(new_source) = new_source {
408 self.file_source = new_source;
409 } else {
410 internal_err!(
411 "FileSource {} does not support projection pushdown",
412 self.file_source.file_type()
413 )?;
414 }
415 Ok(self)
416 }
417
418 pub fn with_constraints(mut self, constraints: Constraints) -> Self {
420 self.constraints = Some(constraints);
421 self
422 }
423
424 pub fn with_statistics(mut self, statistics: Statistics) -> Self {
437 self.statistics = Some(statistics);
438 self
439 }
440
441 pub fn with_file_groups(mut self, file_groups: Vec<FileGroup>) -> Self {
451 self.file_groups = file_groups;
452 self
453 }
454
455 pub fn with_file_group(mut self, file_group: FileGroup) -> Self {
459 self.file_groups.push(file_group);
460 self
461 }
462
463 pub fn with_file(self, partitioned_file: PartitionedFile) -> Self {
467 self.with_file_group(FileGroup::new(vec![partitioned_file]))
468 }
469
470 pub fn with_output_ordering(mut self, output_ordering: Vec<LexOrdering>) -> Self {
479 self.output_ordering = output_ordering;
480 self
481 }
482
483 pub fn with_output_partitioning(
485 mut self,
486 output_partitioning: Option<Partitioning>,
487 ) -> Self {
488 self.output_partitioning = output_partitioning;
489 self
490 }
491
492 pub fn with_file_compression_type(
494 mut self,
495 file_compression_type: FileCompressionType,
496 ) -> Self {
497 self.file_compression_type = Some(file_compression_type);
498 self
499 }
500
501 pub fn with_batch_size(mut self, batch_size: Option<usize>) -> Self {
503 self.batch_size = batch_size;
504 self
505 }
506
507 pub fn with_expr_adapter(
514 mut self,
515 expr_adapter: Option<Arc<dyn PhysicalExprAdapterFactory>>,
516 ) -> Self {
517 self.expr_adapter_factory = expr_adapter;
518 self
519 }
520
521 pub fn build(self) -> FileScanConfig {
529 let Self {
530 object_store_url,
531 file_source,
532 limit,
533 preserve_order,
534 constraints,
535 file_groups,
536 statistics,
537 output_ordering,
538 output_partitioning,
539 file_compression_type,
540 batch_size,
541 expr_adapter_factory: expr_adapter,
542 } = self;
543
544 let constraints = constraints.unwrap_or_default();
545 let statistics = statistics.unwrap_or_else(|| {
546 Statistics::new_unknown(file_source.table_schema().table_schema())
547 });
548 let file_compression_type =
549 file_compression_type.unwrap_or(FileCompressionType::UNCOMPRESSED);
550
551 let preserve_order = preserve_order || !output_ordering.is_empty();
553
554 FileScanConfig {
555 object_store_url,
556 file_source,
557 limit,
558 preserve_order,
559 constraints,
560 file_groups,
561 output_ordering,
562 file_compression_type,
563 batch_size,
564 expr_adapter_factory: expr_adapter,
565 statistics,
566 output_partitioning,
567 }
568 }
569}
570
571impl From<FileScanConfig> for FileScanConfigBuilder {
572 fn from(config: FileScanConfig) -> Self {
573 Self {
574 object_store_url: config.object_store_url,
575 file_source: Arc::<dyn FileSource>::clone(&config.file_source),
576 file_groups: config.file_groups,
577 statistics: Some(config.statistics),
578 output_ordering: config.output_ordering,
579 output_partitioning: config.output_partitioning,
580 file_compression_type: Some(config.file_compression_type),
581 limit: config.limit,
582 preserve_order: config.preserve_order,
583 constraints: Some(config.constraints),
584 batch_size: config.batch_size,
585 expr_adapter_factory: config.expr_adapter_factory,
586 }
587 }
588}
589
590pub fn output_partitioning_from_partition_fields(
595 schema: &Schema,
596 partition_cols: &Fields,
597 partition_count: usize,
598) -> Option<Partitioning> {
599 if partition_cols.is_empty() {
600 return None;
601 }
602
603 let mut exprs: Vec<Arc<dyn PhysicalExpr>> = Vec::with_capacity(partition_cols.len());
604 for partition_col in partition_cols {
605 let name = partition_col.name();
606 let idx = schema
607 .fields()
608 .iter()
609 .position(|field| field.name() == name)?;
610 exprs.push(Arc::new(Column::new(name, idx)));
611 }
612
613 Some(Partitioning::Hash(exprs, partition_count))
614}
615
616fn project_output_partitioning(
617 partitioning: &Partitioning,
618 mapping: &ProjectionMapping,
619 input_schema: &SchemaRef,
620 partition_count: usize,
621) -> Partitioning {
622 let input_eq_properties = EquivalenceProperties::new(Arc::clone(input_schema));
623 match partitioning {
624 Partitioning::Hash(exprs, _) => {
625 let projected_exprs = input_eq_properties
626 .project_expressions(exprs, mapping)
627 .collect::<Option<Vec<_>>>();
628 projected_exprs
629 .map(|exprs| Partitioning::Hash(exprs, partition_count))
630 .unwrap_or_else(|| Partitioning::UnknownPartitioning(partition_count))
631 }
632 Partitioning::Range(_)
633 | Partitioning::RoundRobinBatch(_)
634 | Partitioning::UnknownPartitioning(_) => {
635 partitioning.project(mapping, &input_eq_properties)
636 }
637 }
638}
639
640fn would_duplicate_costly_exprs(
663 inner: &ProjectionExprs,
664 outer: &ProjectionExprs,
665) -> bool {
666 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
667
668 let inner_exprs = inner.as_ref();
669
670 let mut ref_counts = vec![0usize; inner_exprs.len()];
671 for proj_expr in outer.as_ref() {
672 proj_expr
673 .expr
674 .apply(|e| {
675 if let Some(col) = e.as_ref().downcast_ref::<Column>()
676 && let Some(count) = ref_counts.get_mut(col.index())
677 {
678 *count += 1;
679 }
680 Ok(TreeNodeRecursion::Continue)
681 })
682 .expect("infallible closure should not fail");
683 }
684
685 ref_counts.iter().enumerate().any(|(idx, &count)| {
686 let expr = &inner_exprs[idx].expr;
687 count > 1 && (is_volatile(expr) || !expr.placement().should_push_to_leaves())
688 })
689}
690
691impl DataSource for FileScanConfig {
692 fn open(
693 &self,
694 partition: usize,
695 context: Arc<TaskContext>,
696 ) -> Result<SendableRecordBatchStream> {
697 self.open_with_args(OpenArgs::new(partition, context))
698 }
699
700 fn open_with_args(&self, args: OpenArgs) -> Result<SendableRecordBatchStream> {
701 let OpenArgs {
702 partition,
703 context,
704 sibling_state,
705 } = args;
706 let object_store = context.runtime_env().object_store(&self.object_store_url)?;
707 let batch_size = self
708 .batch_size
709 .unwrap_or_else(|| context.session_config().batch_size());
710
711 let source = self.file_source.with_batch_size(batch_size);
712
713 let morselizer = source.create_morselizer(object_store, self, partition)?;
714
715 let shared_work_source = sibling_state
719 .as_ref()
720 .and_then(|state| state.downcast_ref::<SharedWorkSource>())
721 .cloned();
722
723 let stream = FileStreamBuilder::new(self)
724 .with_partition(partition)
725 .with_shared_work_source(shared_work_source)
726 .with_morselizer(morselizer)
727 .with_metrics(source.metrics())
728 .build()?;
729 Ok(Box::pin(cooperative(stream)))
730 }
731
732 fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult {
733 match t {
734 DisplayFormatType::Default | DisplayFormatType::Verbose => {
735 let schema = self.projected_schema().map_err(|_| std::fmt::Error {})?;
736 let orderings =
737 sort_pushdown::get_projected_output_ordering(self, &schema);
738
739 write!(f, "file_groups=")?;
740 FileGroupsDisplay(&self.file_groups).fmt_as(t, f)?;
741
742 if !schema.fields().is_empty() {
743 if let Some(projection) = self.file_source.projection() {
744 let expr: Vec<String> = projection
747 .as_ref()
748 .iter()
749 .map(|proj_expr| {
750 if let Some(column) =
751 proj_expr.expr.downcast_ref::<Column>()
752 {
753 if column.name() == proj_expr.alias {
754 column.name().to_string()
755 } else {
756 format!(
757 "{} as {}",
758 proj_expr.expr, proj_expr.alias
759 )
760 }
761 } else {
762 format!("{} as {}", proj_expr.expr, proj_expr.alias)
763 }
764 })
765 .collect();
766 write!(f, ", projection=[{}]", expr.join(", "))?;
767 } else {
768 write!(f, ", projection={}", ProjectSchemaDisplay(&schema))?;
769 }
770 }
771
772 if let Some(limit) = self.limit {
773 write!(f, ", limit={limit}")?;
774 }
775
776 display_orderings(f, &orderings)?;
777
778 if self.output_partitioning.is_some() {
779 write!(f, ", output_partitioning={}", self.output_partitioning())?;
780 }
781
782 if !self.constraints.is_empty() {
783 write!(f, ", {}", self.constraints)?;
784 }
785
786 self.fmt_file_source(t, f)
787 }
788 DisplayFormatType::TreeRender => {
789 writeln!(f, "format={}", self.file_source.file_type())?;
790 self.file_source.fmt_extra(t, f)?;
791 let num_files = self.file_groups.iter().map(|fg| fg.len()).sum::<usize>();
792 writeln!(f, "files={num_files}")?;
793 Ok(())
794 }
795 }
796 }
797
798 fn repartitioned(
800 &self,
801 target_partitions: usize,
802 repartition_file_min_size: usize,
803 output_ordering: Option<LexOrdering>,
804 ) -> Result<Option<Arc<dyn DataSource>>> {
805 if self.output_partitioning.is_some() {
808 return Ok(None);
809 }
810
811 let source = self.file_source.repartitioned(
812 target_partitions,
813 repartition_file_min_size,
814 output_ordering,
815 self,
816 )?;
817
818 Ok(source.map(|s| Arc::new(s) as _))
819 }
820
821 fn output_partitioning(&self) -> Partitioning {
841 let Some(output_partitioning) = self.output_partitioning.clone() else {
842 return Partitioning::UnknownPartitioning(self.file_groups.len());
843 };
844 if output_partitioning.partition_count() != self.file_groups.len() {
845 warn!(
846 "Declared output partitioning has {} partitions, but file scan has {} file groups. Falling back to UnknownPartitioning.",
847 output_partitioning.partition_count(),
848 self.file_groups.len()
849 );
850 return Partitioning::UnknownPartitioning(self.file_groups.len());
851 }
852
853 if let Some(projection) = self.file_source.projection() {
854 let schema = self.file_source.table_schema().table_schema();
855 return match projection.projection_mapping(schema) {
856 Ok(mapping) => project_output_partitioning(
857 &output_partitioning,
858 &mapping,
859 schema,
860 self.file_groups.len(),
861 ),
862 Err(e) => {
863 debug!(
864 "Could not project output partitioning, falling back to UnknownPartitioning: {e}"
865 );
866 Partitioning::UnknownPartitioning(self.file_groups.len())
867 }
868 };
869 }
870
871 output_partitioning
872 }
873
874 fn eq_properties(&self) -> EquivalenceProperties {
878 let schema = self.file_source.table_schema().table_schema();
879 let mut eq_properties = EquivalenceProperties::new_with_orderings(
880 Arc::clone(schema),
881 self.validated_output_ordering(),
882 )
883 .with_constraints(self.constraints.clone());
884
885 if let Some(filter) = self.file_source.filter() {
886 match Self::add_filter_equivalence_info(&filter, &mut eq_properties, schema) {
889 Ok(()) => {}
890 Err(e) => {
891 warn!("Failed to add filter equivalence info: {e}");
892 #[cfg(debug_assertions)]
893 panic!("Failed to add filter equivalence info: {e}");
894 }
895 }
896 }
897
898 if let Some(projection) = self.file_source.projection() {
899 match (
900 projection.project_schema(schema),
901 projection.projection_mapping(schema),
902 ) {
903 (Ok(output_schema), Ok(mapping)) => {
904 eq_properties =
905 eq_properties.project(&mapping, Arc::new(output_schema));
906 }
907 (Err(e), _) | (_, Err(e)) => {
908 warn!("Failed to project equivalence properties: {e}");
909 #[cfg(debug_assertions)]
910 panic!("Failed to project equivalence properties: {e}");
911 }
912 }
913 }
914
915 eq_properties
916 }
917
918 fn scheduling_type(&self) -> SchedulingType {
919 SchedulingType::Cooperative
920 }
921
922 fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
923 if let Some(partition) = partition {
924 if let Some(file_group) = self.file_groups.get(partition)
927 && let Some(stat) = file_group.file_statistics(None)
928 {
929 let output_schema = self.projected_schema()?;
931 return if let Some(projection) = self.file_source.projection() {
932 Ok(Arc::new(
933 projection.project_statistics(stat.clone(), &output_schema)?,
934 ))
935 } else {
936 Ok(Arc::new(stat.clone()))
937 };
938 }
939 Ok(Arc::new(Statistics::new_unknown(
941 self.projected_schema()?.as_ref(),
942 )))
943 } else {
944 let statistics = self.statistics();
946 let projection = self.file_source.projection();
947 let output_schema = self.projected_schema()?;
948 if let Some(projection) = &projection {
949 Ok(Arc::new(
950 projection.project_statistics(statistics.clone(), &output_schema)?,
951 ))
952 } else {
953 Ok(Arc::new(statistics))
954 }
955 }
956 }
957
958 fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn DataSource>> {
959 let source = FileScanConfigBuilder::from(self.clone())
960 .with_limit(limit)
961 .build();
962 Some(Arc::new(source))
963 }
964
965 fn fetch(&self) -> Option<usize> {
966 self.limit
967 }
968
969 fn metrics(&self) -> ExecutionPlanMetricsSet {
970 self.file_source.metrics().clone()
971 }
972
973 fn try_swapping_with_projection(
974 &self,
975 projection: &ProjectionExprs,
976 ) -> Result<Option<Arc<dyn DataSource>>> {
977 if let Some(inner) = self.file_source.projection()
984 && would_duplicate_costly_exprs(inner, projection)
985 {
986 return Ok(None);
987 }
988 match self.file_source.try_pushdown_projection(projection)? {
989 Some(new_source) => {
990 let mut new_file_scan_config = self.clone();
991 new_file_scan_config.file_source = new_source;
992 Ok(Some(Arc::new(new_file_scan_config) as Arc<dyn DataSource>))
993 }
994 None => Ok(None),
995 }
996 }
997
998 fn try_pushdown_filters(
999 &self,
1000 filters: Vec<Arc<dyn PhysicalExpr>>,
1001 config: &ConfigOptions,
1002 ) -> Result<FilterPushdownPropagation<Arc<dyn DataSource>>> {
1003 let table_schema = self.file_source.table_schema().table_schema();
1012 let filters_to_remap = if let Some(projection) = self.file_source.projection() {
1013 filters
1014 .into_iter()
1015 .map(|filter| projection.unproject_expr(&filter))
1016 .collect::<Result<Vec<_>>>()?
1017 } else {
1018 filters
1019 };
1020 let remapped_filters = filters_to_remap
1022 .into_iter()
1023 .map(|filter| reassign_expr_columns(filter, table_schema))
1024 .collect::<Result<Vec<_>>>()?;
1025
1026 let result = self
1027 .file_source
1028 .try_pushdown_filters(remapped_filters, config)?;
1029 match result.updated_node {
1030 Some(new_file_source) => {
1031 let mut new_file_scan_config = self.clone();
1032 new_file_scan_config.file_source = new_file_source;
1033 Ok(FilterPushdownPropagation {
1034 filters: result.filters,
1035 updated_node: Some(Arc::new(new_file_scan_config) as _),
1036 })
1037 }
1038 None => {
1039 Ok(FilterPushdownPropagation {
1041 filters: result.filters,
1042 updated_node: None,
1043 })
1044 }
1045 }
1046 }
1047
1048 fn try_pushdown_sort(
1089 &self,
1090 order: &[PhysicalSortExpr],
1091 ) -> Result<SortOrderPushdownResult<Arc<dyn DataSource>>> {
1092 let pushdown_result = self
1093 .file_source
1094 .try_pushdown_sort(order, &self.eq_properties())?;
1095
1096 match pushdown_result {
1097 SortOrderPushdownResult::Exact { inner } => {
1098 let config = self.rebuild_with_source(inner, true, order)?;
1099 if config.output_ordering.is_empty() {
1103 Ok(SortOrderPushdownResult::Inexact {
1104 inner: Arc::new(config),
1105 })
1106 } else {
1107 Ok(SortOrderPushdownResult::Exact {
1108 inner: Arc::new(config),
1109 })
1110 }
1111 }
1112 SortOrderPushdownResult::Inexact { inner } => {
1113 let mut config = self.rebuild_with_source(inner, false, order)?;
1114 if config.output_ordering.is_empty() {
1123 return Ok(SortOrderPushdownResult::Inexact {
1124 inner: Arc::new(config),
1125 });
1126 }
1127 config.file_source = Arc::clone(&self.file_source);
1147 Ok(SortOrderPushdownResult::Exact {
1148 inner: Arc::new(config),
1149 })
1150 }
1151 SortOrderPushdownResult::Unsupported => {
1152 self.try_sort_file_groups_by_statistics(order)
1153 }
1154 }
1155 }
1156
1157 fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn DataSource>> {
1158 if self.preserve_order == preserve_order {
1159 return Some(Arc::new(self.clone()));
1160 }
1161
1162 let new_config = FileScanConfig {
1163 preserve_order,
1164 ..self.clone()
1165 };
1166 Some(Arc::new(new_config))
1167 }
1168
1169 fn apply_expressions(
1170 &self,
1171 f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
1172 ) -> Result<TreeNodeRecursion> {
1173 self.file_source.apply_expressions(f)
1175 }
1176
1177 fn create_sibling_state(
1186 &self,
1187 config: &ConfigOptions,
1188 ) -> Option<Arc<dyn Any + Send + Sync>> {
1189 if self.preserve_order
1190 || self.output_partitioning.is_some()
1191 || !config.execution.enable_file_stream_work_stealing
1192 {
1193 return None;
1194 }
1195
1196 Some(Arc::new(SharedWorkSource::from_config(self)) as Arc<dyn Any + Send + Sync>)
1197 }
1198
1199 #[cfg(feature = "proto")]
1204 fn try_to_proto(
1205 &self,
1206 ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
1207 ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
1208 self.file_source().try_to_proto(self, ctx)
1209 }
1210}
1211
1212impl FileScanConfig {
1213 fn validated_output_ordering(&self) -> Vec<LexOrdering> {
1243 let schema = self.file_source.table_schema().table_schema();
1244 sort_pushdown::validate_orderings(
1245 &self.output_ordering,
1246 schema,
1247 &self.file_groups,
1248 None,
1249 )
1250 }
1251
1252 pub fn file_schema(&self) -> &SchemaRef {
1254 self.file_source.table_schema().file_schema()
1255 }
1256
1257 pub fn table_partition_cols(&self) -> &Fields {
1259 self.file_source.table_schema().table_partition_cols()
1260 }
1261
1262 pub fn statistics(&self) -> Statistics {
1268 let filter_may_change_row_count = self.file_source.filter().is_some()
1269 && self.statistics.num_rows != Precision::Exact(0);
1270 if filter_may_change_row_count {
1271 self.statistics.clone().to_inexact()
1272 } else {
1273 self.statistics.clone()
1274 }
1275 }
1276
1277 pub fn projected_schema(&self) -> Result<Arc<Schema>> {
1278 let schema = self.file_source.table_schema().table_schema();
1279 match self.file_source.projection() {
1280 Some(proj) => Ok(Arc::new(proj.project_schema(schema)?)),
1281 None => Ok(Arc::clone(schema)),
1282 }
1283 }
1284
1285 fn add_filter_equivalence_info(
1286 filter: &Arc<dyn PhysicalExpr>,
1287 eq_properties: &mut EquivalenceProperties,
1288 schema: &Schema,
1289 ) -> Result<()> {
1290 let equal_pairs = split_conjunction(filter).into_iter().filter_map(|expr| {
1292 reassign_expr_columns(Arc::clone(expr), schema)
1295 .ok()
1296 .and_then(|expr| match expr.downcast_ref::<BinaryExpr>() {
1297 Some(expr) if expr.op() == &Operator::Eq => {
1298 Some((Arc::clone(expr.left()), Arc::clone(expr.right())))
1299 }
1300 _ => None,
1301 })
1302 });
1303
1304 for (lhs, rhs) in equal_pairs {
1305 eq_properties.add_equal_conditions(lhs, rhs)?
1306 }
1307
1308 Ok(())
1309 }
1310
1311 #[deprecated(
1320 since = "52.0.0",
1321 note = "newlines_in_values has moved to CsvSource. Access it via CsvSource::csv_options().newlines_in_values instead. It will be removed in 58.0.0 or 6 months after 52.0.0 is released, whichever comes first."
1322 )]
1323 pub fn newlines_in_values(&self) -> bool {
1324 false
1325 }
1326
1327 #[deprecated(
1328 since = "52.0.0",
1329 note = "This method is no longer used, use eq_properties instead. It will be removed in 58.0.0 or 6 months after 52.0.0 is released, whichever comes first."
1330 )]
1331 pub fn projected_constraints(&self) -> Constraints {
1332 let props = self.eq_properties();
1333 props.constraints().clone()
1334 }
1335
1336 #[deprecated(
1337 since = "52.0.0",
1338 note = "This method is no longer used, use eq_properties instead. It will be removed in 58.0.0 or 6 months after 52.0.0 is released, whichever comes first."
1339 )]
1340 pub fn file_column_projection_indices(&self) -> Option<Vec<usize>> {
1341 #[expect(deprecated)]
1342 self.file_source.projection().as_ref().map(|p| {
1343 p.ordered_column_indices()
1344 .into_iter()
1345 .filter(|&i| i < self.file_schema().fields().len())
1346 .collect::<Vec<_>>()
1347 })
1348 }
1349
1350 pub fn split_groups_by_statistics_with_target_partitions(
1372 table_schema: &SchemaRef,
1373 file_groups: &[FileGroup],
1374 sort_order: &LexOrdering,
1375 target_partitions: usize,
1376 ) -> Result<Vec<FileGroup>> {
1377 if target_partitions == 0 {
1378 return Err(internal_datafusion_err!(
1379 "target_partitions must be greater than 0"
1380 ));
1381 }
1382
1383 let flattened_files = file_groups
1384 .iter()
1385 .flat_map(FileGroup::iter)
1386 .collect::<Vec<_>>();
1387
1388 if flattened_files.is_empty() {
1389 return Ok(vec![]);
1390 }
1391
1392 let statistics = MinMaxStatistics::new_from_files(
1393 sort_order,
1394 table_schema,
1395 None,
1396 flattened_files.iter().copied(),
1397 )?;
1398
1399 let indices_sorted_by_min = statistics.min_values_sorted();
1400
1401 let mut file_groups_indices: Vec<Vec<usize>> = vec![vec![]; target_partitions];
1403
1404 for (idx, min) in indices_sorted_by_min {
1405 if let Some((_, group)) = file_groups_indices
1406 .iter_mut()
1407 .enumerate()
1408 .filter(|(_, group)| {
1409 group.is_empty()
1410 || min
1411 > statistics
1412 .max(*group.last().expect("groups should not be empty"))
1413 })
1414 .min_by_key(|(_, group)| group.len())
1415 {
1416 group.push(idx);
1417 } else {
1418 file_groups_indices.push(vec![idx]);
1420 }
1421 }
1422
1423 file_groups_indices.retain(|group| !group.is_empty());
1425
1426 Ok(file_groups_indices
1428 .into_iter()
1429 .map(|file_group_indices| {
1430 FileGroup::new(
1431 file_group_indices
1432 .into_iter()
1433 .map(|idx| flattened_files[idx].clone())
1434 .collect(),
1435 )
1436 })
1437 .collect())
1438 }
1439
1440 pub fn split_groups_by_statistics(
1444 table_schema: &SchemaRef,
1445 file_groups: &[FileGroup],
1446 sort_order: &LexOrdering,
1447 ) -> Result<Vec<FileGroup>> {
1448 let flattened_files = file_groups
1449 .iter()
1450 .flat_map(FileGroup::iter)
1451 .collect::<Vec<_>>();
1452 if flattened_files.is_empty() {
1464 return Ok(vec![]);
1465 }
1466
1467 let statistics = MinMaxStatistics::new_from_files(
1468 sort_order,
1469 table_schema,
1470 None,
1471 flattened_files.iter().copied(),
1472 )
1473 .map_err(|e| {
1474 e.context("construct min/max statistics for split_groups_by_statistics")
1475 })?;
1476
1477 let indices_sorted_by_min = statistics.min_values_sorted();
1478 let mut file_groups_indices: Vec<Vec<usize>> = vec![];
1479
1480 for (idx, min) in indices_sorted_by_min {
1481 let file_group_to_insert = file_groups_indices.iter_mut().find(|group| {
1482 min > statistics.max(
1485 *group
1486 .last()
1487 .expect("groups should be nonempty at construction"),
1488 )
1489 });
1490 match file_group_to_insert {
1491 Some(group) => group.push(idx),
1492 None => file_groups_indices.push(vec![idx]),
1493 }
1494 }
1495
1496 Ok(file_groups_indices
1498 .into_iter()
1499 .map(|file_group_indices| {
1500 file_group_indices
1501 .into_iter()
1502 .map(|idx| flattened_files[idx].clone())
1503 .collect()
1504 })
1505 .collect())
1506 }
1507
1508 fn fmt_file_source(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult {
1510 write!(f, ", file_type={}", self.file_source.file_type())?;
1511 self.file_source.fmt_extra(t, f)
1512 }
1513
1514 pub fn file_source(&self) -> &Arc<dyn FileSource> {
1516 &self.file_source
1517 }
1518
1519 }
1523
1524impl Debug for FileScanConfig {
1525 fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
1526 write!(f, "FileScanConfig {{")?;
1527 write!(f, "object_store_url={:?}, ", self.object_store_url)?;
1528
1529 write!(f, "statistics={:?}, ", self.statistics())?;
1530
1531 DisplayAs::fmt_as(self, DisplayFormatType::Verbose, f)?;
1532 write!(f, "}}")
1533 }
1534}
1535
1536impl DisplayAs for FileScanConfig {
1537 fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult {
1538 let schema = self.projected_schema().map_err(|_| std::fmt::Error {})?;
1539 let orderings = sort_pushdown::get_projected_output_ordering(self, &schema);
1540
1541 write!(f, "file_groups=")?;
1542 FileGroupsDisplay(&self.file_groups).fmt_as(t, f)?;
1543
1544 if !schema.fields().is_empty() {
1545 write!(f, ", projection={}", ProjectSchemaDisplay(&schema))?;
1546 }
1547
1548 if let Some(limit) = self.limit {
1549 write!(f, ", limit={limit}")?;
1550 }
1551
1552 display_orderings(f, &orderings)?;
1553
1554 if !self.constraints.is_empty() {
1555 write!(f, ", {}", self.constraints)?;
1556 }
1557
1558 Ok(())
1559 }
1560}
1561
1562pub fn wrap_partition_type_in_dict(val_type: DataType) -> DataType {
1573 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(val_type))
1574}
1575
1576pub fn wrap_partition_value_in_dict(val: ScalarValue) -> ScalarValue {
1580 ScalarValue::Dictionary(Box::new(DataType::UInt16), Box::new(val))
1581}
1582
1583#[cfg(test)]
1584mod tests {
1585 use std::collections::HashMap;
1586
1587 use super::*;
1588 use crate::source::DataSourceExec;
1589 use crate::test_util::col;
1590 use crate::{TableSchema, TableSchemaBuilder};
1591 use crate::{
1592 generate_test_files, test_util::MockSource, tests::aggr_test_schema,
1593 verify_sort_integrity,
1594 };
1595
1596 use arrow::array::{Int32Array, RecordBatch};
1597 use arrow::datatypes::Field;
1598 use datafusion_common::ColumnStatistics;
1599 use datafusion_common::stats::Precision;
1600 use datafusion_common::tree_node::TreeNodeRecursion;
1601 use datafusion_common::{Result, assert_batches_eq, internal_err};
1602 use datafusion_execution::TaskContext;
1603 use datafusion_expr::SortExpr;
1604 use datafusion_physical_expr::PhysicalExpr;
1605
1606 #[cfg(feature = "proto")]
1607 use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
1608 use datafusion_physical_expr::create_physical_sort_expr;
1609 use datafusion_physical_expr::expressions::Literal;
1610 use datafusion_physical_expr::projection::ProjectionExpr;
1611 use datafusion_physical_expr::projection::ProjectionExprs;
1612 use datafusion_physical_plan::ExecutionPlan;
1613 use datafusion_physical_plan::execution_plan::collect;
1614 #[cfg(feature = "proto")]
1615 use datafusion_physical_plan::proto::{ExecutionPlanEncode, ExecutionPlanEncodeCtx};
1616 #[cfg(feature = "proto")]
1617 use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode};
1618 use futures::FutureExt as _;
1619 use futures::StreamExt as _;
1620 use futures::stream;
1621 use object_store::ObjectStore;
1622 use std::fmt::Debug;
1623
1624 #[derive(Clone)]
1625 struct InexactSortPushdownSource {
1626 metrics: ExecutionPlanMetricsSet,
1627 table_schema: TableSchema,
1628 }
1629
1630 impl InexactSortPushdownSource {
1631 fn new(table_schema: TableSchema) -> Self {
1632 Self {
1633 metrics: ExecutionPlanMetricsSet::new(),
1634 table_schema,
1635 }
1636 }
1637 }
1638
1639 impl FileSource for InexactSortPushdownSource {
1640 fn create_file_opener(
1641 &self,
1642 _object_store: Arc<dyn ObjectStore>,
1643 _base_config: &FileScanConfig,
1644 _partition: usize,
1645 ) -> Result<Arc<dyn crate::file_stream::FileOpener>> {
1646 unimplemented!()
1647 }
1648
1649 fn table_schema(&self) -> &TableSchema {
1650 &self.table_schema
1651 }
1652
1653 fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> {
1654 Arc::new(self.clone())
1655 }
1656
1657 fn metrics(&self) -> &ExecutionPlanMetricsSet {
1658 &self.metrics
1659 }
1660
1661 fn file_type(&self) -> &str {
1662 "mock"
1663 }
1664
1665 fn try_pushdown_sort(
1666 &self,
1667 _order: &[PhysicalSortExpr],
1668 _eq_properties: &EquivalenceProperties,
1669 ) -> Result<SortOrderPushdownResult<Arc<dyn FileSource>>> {
1670 Ok(SortOrderPushdownResult::Inexact {
1671 inner: Arc::new(self.clone()) as Arc<dyn FileSource>,
1672 })
1673 }
1674
1675 fn apply_expressions(
1676 &self,
1677 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
1678 ) -> Result<TreeNodeRecursion> {
1679 Ok(TreeNodeRecursion::Continue)
1680 }
1681 }
1682
1683 #[cfg(feature = "proto")]
1684 #[derive(Clone)]
1685 struct ProtoHookSource {
1686 metrics: ExecutionPlanMetricsSet,
1687 table_schema: TableSchema,
1688 }
1689
1690 #[cfg(feature = "proto")]
1691 impl ProtoHookSource {
1692 fn new(table_schema: TableSchema) -> Self {
1693 Self {
1694 metrics: ExecutionPlanMetricsSet::new(),
1695 table_schema,
1696 }
1697 }
1698 }
1699
1700 #[cfg(feature = "proto")]
1701 impl FileSource for ProtoHookSource {
1702 fn create_file_opener(
1703 &self,
1704 _object_store: Arc<dyn ObjectStore>,
1705 _base_config: &FileScanConfig,
1706 _partition: usize,
1707 ) -> Result<Arc<dyn crate::file_stream::FileOpener>> {
1708 internal_err!("not needed for proto delegation test")
1709 }
1710
1711 fn table_schema(&self) -> &TableSchema {
1712 &self.table_schema
1713 }
1714
1715 fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> {
1716 Arc::new(self.clone())
1717 }
1718
1719 fn metrics(&self) -> &ExecutionPlanMetricsSet {
1720 &self.metrics
1721 }
1722
1723 fn file_type(&self) -> &str {
1724 "proto-hook-test"
1725 }
1726
1727 fn apply_expressions(
1728 &self,
1729 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
1730 ) -> Result<TreeNodeRecursion> {
1731 Ok(TreeNodeRecursion::Continue)
1732 }
1733
1734 fn try_to_proto(
1735 &self,
1736 _base: &FileScanConfig,
1737 _ctx: &ExecutionPlanEncodeCtx<'_>,
1738 ) -> Result<Option<PhysicalPlanNode>> {
1739 Ok(Some(PhysicalPlanNode::default()))
1740 }
1741 }
1742
1743 #[cfg(feature = "proto")]
1744 struct UnusedPlanEncoder;
1745
1746 #[cfg(feature = "proto")]
1747 impl ExecutionPlanEncode for UnusedPlanEncoder {
1748 fn encode_plan(
1749 &self,
1750 _plan: &Arc<dyn ExecutionPlan>,
1751 ) -> Result<PhysicalPlanNode> {
1752 internal_err!("not needed for proto delegation test")
1753 }
1754
1755 fn encode_expr(&self, _expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode> {
1756 internal_err!("not needed for proto delegation test")
1757 }
1758
1759 fn encode_udf(&self, _udf: &ScalarUDF) -> Result<Option<Vec<u8>>> {
1760 internal_err!("not needed for proto delegation test")
1761 }
1762
1763 fn encode_udaf(&self, _udaf: &AggregateUDF) -> Result<Option<Vec<u8>>> {
1764 internal_err!("not needed for proto delegation test")
1765 }
1766
1767 fn encode_udwf(&self, _udwf: &WindowUDF) -> Result<Option<Vec<u8>>> {
1768 internal_err!("not needed for proto delegation test")
1769 }
1770 }
1771
1772 #[cfg(feature = "proto")]
1773 #[test]
1774 fn data_source_exec_delegates_proto_to_file_source() -> Result<()> {
1775 let schema = Arc::new(Schema::new(vec![Field::new(
1776 "value",
1777 DataType::Int32,
1778 false,
1779 )]));
1780 let source = Arc::new(ProtoHookSource::new(TableSchema::from(&schema)));
1781 let config =
1782 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source)
1783 .build();
1784 let exec = DataSourceExec::from_data_source(config);
1785 let encoder = UnusedPlanEncoder;
1786 let ctx = ExecutionPlanEncodeCtx::new(&encoder);
1787
1788 assert_eq!(exec.try_to_proto(&ctx)?, Some(PhysicalPlanNode::default()));
1789 Ok(())
1790 }
1791
1792 #[test]
1793 fn physical_plan_config_no_projection_tab_cols_as_field() {
1794 let file_schema = aggr_test_schema();
1795
1796 let table_partition_col =
1798 Field::new("date", wrap_partition_type_in_dict(DataType::Utf8), true)
1799 .with_metadata(HashMap::from_iter(vec![(
1800 "key_whatever".to_owned(),
1801 "value_whatever".to_owned(),
1802 )]));
1803
1804 let conf = config_for_projection(
1805 Arc::clone(&file_schema),
1806 None,
1807 Statistics::new_unknown(&file_schema),
1808 vec![table_partition_col.clone()],
1809 );
1810
1811 let proj_schema = conf.projected_schema().unwrap();
1813 assert_eq!(proj_schema.fields().len(), file_schema.fields().len() + 1);
1814 assert_eq!(
1815 *proj_schema.field(file_schema.fields().len()),
1816 table_partition_col,
1817 "partition columns are the last columns and ust have all values defined in created field"
1818 );
1819 }
1820
1821 #[test]
1822 fn test_split_groups_by_statistics() -> Result<()> {
1823 use chrono::TimeZone;
1824 use datafusion_common::DFSchema;
1825 use datafusion_expr::execution_props::ExecutionProps;
1826 use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
1827 use object_store::{ObjectMeta, path::Path};
1828
1829 struct File {
1830 name: &'static str,
1831 date: &'static str,
1832 statistics: Vec<Option<(Option<f64>, Option<f64>)>>,
1833 }
1834 impl File {
1835 fn new(
1836 name: &'static str,
1837 date: &'static str,
1838 statistics: Vec<Option<(f64, f64)>>,
1839 ) -> Self {
1840 Self::new_nullable(
1841 name,
1842 date,
1843 statistics
1844 .into_iter()
1845 .map(|opt| opt.map(|(min, max)| (Some(min), Some(max))))
1846 .collect(),
1847 )
1848 }
1849
1850 fn new_nullable(
1851 name: &'static str,
1852 date: &'static str,
1853 statistics: Vec<Option<(Option<f64>, Option<f64>)>>,
1854 ) -> Self {
1855 Self {
1856 name,
1857 date,
1858 statistics,
1859 }
1860 }
1861 }
1862
1863 struct TestCase {
1864 name: &'static str,
1865 file_schema: Schema,
1866 files: Vec<File>,
1867 sort: Vec<SortExpr>,
1868 expected_result: Result<Vec<Vec<&'static str>>, &'static str>,
1869 }
1870
1871 use datafusion_expr::col;
1872 let cases = vec![
1873 TestCase {
1874 name: "test sort",
1875 file_schema: Schema::new(vec![Field::new(
1876 "value".to_string(),
1877 DataType::Float64,
1878 false,
1879 )]),
1880 files: vec![
1881 File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1882 File::new("1", "2023-01-01", vec![Some((0.50, 1.00))]),
1883 File::new("2", "2023-01-02", vec![Some((0.00, 1.00))]),
1884 ],
1885 sort: vec![col("value").sort(true, false)],
1886 expected_result: Ok(vec![vec!["0", "1"], vec!["2"]]),
1887 },
1888 TestCase {
1891 name: "test sort with files ordered differently",
1892 file_schema: Schema::new(vec![Field::new(
1893 "value".to_string(),
1894 DataType::Float64,
1895 false,
1896 )]),
1897 files: vec![
1898 File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1899 File::new("2", "2023-01-02", vec![Some((0.00, 1.00))]),
1900 File::new("1", "2023-01-01", vec![Some((0.50, 1.00))]),
1901 ],
1902 sort: vec![col("value").sort(true, false)],
1903 expected_result: Ok(vec![vec!["0", "1"], vec!["2"]]),
1904 },
1905 TestCase {
1906 name: "reverse sort",
1907 file_schema: Schema::new(vec![Field::new(
1908 "value".to_string(),
1909 DataType::Float64,
1910 false,
1911 )]),
1912 files: vec![
1913 File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1914 File::new("1", "2023-01-01", vec![Some((0.50, 1.00))]),
1915 File::new("2", "2023-01-02", vec![Some((0.00, 1.00))]),
1916 ],
1917 sort: vec![col("value").sort(false, true)],
1918 expected_result: Ok(vec![vec!["1", "0"], vec!["2"]]),
1919 },
1920 TestCase {
1921 name: "nullable sort columns, nulls last",
1922 file_schema: Schema::new(vec![Field::new(
1923 "value".to_string(),
1924 DataType::Float64,
1925 true,
1926 )]),
1927 files: vec![
1928 File::new_nullable(
1929 "0",
1930 "2023-01-01",
1931 vec![Some((Some(0.00), Some(0.49)))],
1932 ),
1933 File::new_nullable("1", "2023-01-01", vec![Some((Some(0.50), None))]),
1934 File::new_nullable("2", "2023-01-02", vec![Some((Some(0.00), None))]),
1935 ],
1936 sort: vec![col("value").sort(true, false)],
1937 expected_result: Ok(vec![vec!["0", "1"], vec!["2"]]),
1938 },
1939 TestCase {
1940 name: "nullable sort columns, nulls first",
1941 file_schema: Schema::new(vec![Field::new(
1942 "value".to_string(),
1943 DataType::Float64,
1944 true,
1945 )]),
1946 files: vec![
1947 File::new_nullable("0", "2023-01-01", vec![Some((None, Some(0.49)))]),
1948 File::new_nullable(
1949 "1",
1950 "2023-01-01",
1951 vec![Some((Some(0.50), Some(1.00)))],
1952 ),
1953 File::new_nullable("2", "2023-01-02", vec![Some((None, Some(1.00)))]),
1954 ],
1955 sort: vec![col("value").sort(true, true)],
1956 expected_result: Ok(vec![vec!["0", "1"], vec!["2"]]),
1957 },
1958 TestCase {
1959 name: "all three non-overlapping",
1960 file_schema: Schema::new(vec![Field::new(
1961 "value".to_string(),
1962 DataType::Float64,
1963 false,
1964 )]),
1965 files: vec![
1966 File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1967 File::new("1", "2023-01-01", vec![Some((0.50, 0.99))]),
1968 File::new("2", "2023-01-02", vec![Some((1.00, 1.49))]),
1969 ],
1970 sort: vec![col("value").sort(true, false)],
1971 expected_result: Ok(vec![vec!["0", "1", "2"]]),
1972 },
1973 TestCase {
1974 name: "all three overlapping",
1975 file_schema: Schema::new(vec![Field::new(
1976 "value".to_string(),
1977 DataType::Float64,
1978 false,
1979 )]),
1980 files: vec![
1981 File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
1982 File::new("1", "2023-01-01", vec![Some((0.00, 0.49))]),
1983 File::new("2", "2023-01-02", vec![Some((0.00, 0.49))]),
1984 ],
1985 sort: vec![col("value").sort(true, false)],
1986 expected_result: Ok(vec![vec!["0"], vec!["1"], vec!["2"]]),
1987 },
1988 TestCase {
1989 name: "empty input",
1990 file_schema: Schema::new(vec![Field::new(
1991 "value".to_string(),
1992 DataType::Float64,
1993 false,
1994 )]),
1995 files: vec![],
1996 sort: vec![col("value").sort(true, false)],
1997 expected_result: Ok(vec![]),
1998 },
1999 TestCase {
2000 name: "one file missing statistics",
2001 file_schema: Schema::new(vec![Field::new(
2002 "value".to_string(),
2003 DataType::Float64,
2004 false,
2005 )]),
2006 files: vec![
2007 File::new("0", "2023-01-01", vec![Some((0.00, 0.49))]),
2008 File::new("1", "2023-01-01", vec![Some((0.00, 0.49))]),
2009 File::new("2", "2023-01-02", vec![None]),
2010 ],
2011 sort: vec![col("value").sort(true, false)],
2012 expected_result: Err(
2013 "construct min/max statistics for split_groups_by_statistics\ncaused by\ncollect min/max values\ncaused by\nget min/max for column: 'value'\ncaused by\nError during planning: statistics not found",
2014 ),
2015 },
2016 ];
2017
2018 for case in cases {
2019 let table_schema = Arc::new(Schema::new(
2020 case.file_schema
2021 .fields()
2022 .clone()
2023 .into_iter()
2024 .cloned()
2025 .chain(Some(Arc::new(Field::new(
2026 "date".to_string(),
2027 DataType::Utf8,
2028 false,
2029 ))))
2030 .collect::<Vec<_>>(),
2031 ));
2032 let Some(sort_order) = LexOrdering::new(
2033 case.sort
2034 .into_iter()
2035 .map(|expr| {
2036 create_physical_sort_expr(
2037 &expr,
2038 &DFSchema::try_from(Arc::clone(&table_schema))?,
2039 &ExecutionProps::default(),
2040 &PhysicalPlanningContext::default(),
2041 )
2042 })
2043 .collect::<Result<Vec<_>>>()?,
2044 ) else {
2045 return internal_err!("This test should always use an ordering");
2046 };
2047
2048 let partitioned_files = FileGroup::new(
2049 case.files.into_iter().map(From::from).collect::<Vec<_>>(),
2050 );
2051 let result = FileScanConfig::split_groups_by_statistics(
2052 &table_schema,
2053 std::slice::from_ref(&partitioned_files),
2054 &sort_order,
2055 );
2056 let results_by_name = result
2057 .as_ref()
2058 .map(|file_groups| {
2059 file_groups
2060 .iter()
2061 .map(|file_group| {
2062 file_group
2063 .iter()
2064 .map(|file| {
2065 partitioned_files
2066 .iter()
2067 .find_map(|f| {
2068 if f.object_meta == file.object_meta {
2069 Some(
2070 f.object_meta
2071 .location
2072 .as_ref()
2073 .rsplit('/')
2074 .next()
2075 .unwrap()
2076 .trim_end_matches(".parquet"),
2077 )
2078 } else {
2079 None
2080 }
2081 })
2082 .unwrap()
2083 })
2084 .collect::<Vec<_>>()
2085 })
2086 .collect::<Vec<_>>()
2087 })
2088 .map_err(|e| e.strip_backtrace().leak() as &'static str);
2089
2090 assert_eq!(results_by_name, case.expected_result, "{}", case.name);
2091 }
2092
2093 return Ok(());
2094
2095 impl From<File> for PartitionedFile {
2096 fn from(file: File) -> Self {
2097 let object_meta = ObjectMeta {
2098 location: Path::from(format!(
2099 "data/date={}/{}.parquet",
2100 file.date, file.name
2101 )),
2102 last_modified: chrono::Utc.timestamp_nanos(0),
2103 size: 0,
2104 e_tag: None,
2105 version: None,
2106 };
2107 let statistics = Arc::new(Statistics {
2108 num_rows: Precision::Absent,
2109 total_byte_size: Precision::Absent,
2110 column_statistics: file
2111 .statistics
2112 .into_iter()
2113 .map(|stats| {
2114 stats
2115 .map(|(min, max)| ColumnStatistics {
2116 min_value: Precision::Exact(ScalarValue::Float64(
2117 min,
2118 )),
2119 max_value: Precision::Exact(ScalarValue::Float64(
2120 max,
2121 )),
2122 ..Default::default()
2123 })
2124 .unwrap_or_default()
2125 })
2126 .collect::<Vec<_>>(),
2127 });
2128 PartitionedFile::new_from_meta(object_meta)
2129 .with_partition_values(vec![ScalarValue::from(file.date)])
2130 .with_statistics(statistics)
2131 }
2132 }
2133 }
2134
2135 fn config_for_projection(
2137 file_schema: SchemaRef,
2138 projection: Option<Vec<usize>>,
2139 statistics: Statistics,
2140 table_partition_cols: Vec<Field>,
2141 ) -> FileScanConfig {
2142 let table_schema = TableSchema::builder(file_schema)
2143 .with_table_partition_cols(
2144 table_partition_cols
2145 .into_iter()
2146 .map(Arc::new)
2147 .collect::<Fields>(),
2148 )
2149 .build();
2150 FileScanConfigBuilder::new(
2151 ObjectStoreUrl::parse("test:///").unwrap(),
2152 Arc::new(MockSource::new(table_schema.clone())),
2153 )
2154 .with_projection_indices(projection)
2155 .unwrap()
2156 .with_statistics(statistics)
2157 .build()
2158 }
2159
2160 #[test]
2161 fn test_file_scan_config_builder() {
2162 let file_schema = aggr_test_schema();
2163 let object_store_url = ObjectStoreUrl::parse("test:///").unwrap();
2164
2165 let table_schema = TableSchemaBuilder::from(&file_schema)
2166 .with_table_partition_cols(vec![Arc::new(Field::new(
2167 "date",
2168 wrap_partition_type_in_dict(DataType::Utf8),
2169 false,
2170 ))])
2171 .build();
2172
2173 let file_source: Arc<dyn FileSource> =
2174 Arc::new(MockSource::new(table_schema.clone()));
2175
2176 let builder = FileScanConfigBuilder::new(
2178 object_store_url.clone(),
2179 Arc::clone(&file_source),
2180 );
2181
2182 let config = builder
2184 .with_limit(Some(1000))
2185 .with_projection_indices(Some(vec![0, 1]))
2186 .unwrap()
2187 .with_statistics(Statistics::new_unknown(&file_schema))
2188 .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new(
2189 "test.parquet".to_string(),
2190 1024,
2191 )])])
2192 .with_output_ordering(vec![
2193 [PhysicalSortExpr::new_default(Arc::new(Column::new(
2194 "date", 0,
2195 )))]
2196 .into(),
2197 ])
2198 .with_file_compression_type(FileCompressionType::UNCOMPRESSED)
2199 .build();
2200
2201 assert_eq!(config.object_store_url, object_store_url);
2203 assert_eq!(*config.file_schema(), file_schema);
2204 assert_eq!(config.limit, Some(1000));
2205 assert_eq!(
2206 config
2207 .file_source
2208 .projection()
2209 .as_ref()
2210 .map(|p| p.column_indices()),
2211 Some(vec![0, 1])
2212 );
2213 assert_eq!(config.table_partition_cols().len(), 1);
2214 assert_eq!(config.table_partition_cols()[0].name(), "date");
2215 assert_eq!(config.file_groups.len(), 1);
2216 assert_eq!(config.file_groups[0].len(), 1);
2217 assert_eq!(
2218 config.file_groups[0][0].object_meta.location.as_ref(),
2219 "test.parquet"
2220 );
2221 assert_eq!(
2222 config.file_compression_type,
2223 FileCompressionType::UNCOMPRESSED
2224 );
2225 assert_eq!(config.output_ordering.len(), 1);
2226 }
2227
2228 #[test]
2229 fn equivalence_properties_after_schema_change() {
2230 let file_schema = aggr_test_schema();
2231 let object_store_url = ObjectStoreUrl::parse("test:///").unwrap();
2232
2233 let table_schema = TableSchema::from(&file_schema);
2234
2235 let file_source: Arc<dyn FileSource> = Arc::new(
2237 MockSource::new(table_schema.clone()).with_filter(Arc::new(BinaryExpr::new(
2238 col("c2", &file_schema).unwrap(),
2239 Operator::Eq,
2240 Arc::new(Literal::new(ScalarValue::Int32(Some(10)))),
2241 ))),
2242 );
2243
2244 let config = FileScanConfigBuilder::new(
2245 object_store_url.clone(),
2246 Arc::clone(&file_source),
2247 )
2248 .with_projection_indices(Some(vec![0, 1, 2]))
2249 .unwrap()
2250 .build();
2251
2252 let exprs = ProjectionExprs::new(vec![ProjectionExpr::new(
2255 col("c1", &file_schema).unwrap(),
2256 "c1",
2257 )]);
2258 let data_source = config
2259 .try_swapping_with_projection(&exprs)
2260 .unwrap()
2261 .unwrap();
2262
2263 let eq_properties = data_source.eq_properties();
2266 let eq_group = eq_properties.eq_group();
2267
2268 for class in eq_group.iter() {
2269 for expr in class.iter() {
2270 if let Some(col) = expr.downcast_ref::<Column>() {
2271 assert_ne!(
2272 col.name(),
2273 "c2",
2274 "c2 should not be present in any equivalence class"
2275 );
2276 }
2277 }
2278 }
2279 }
2280
2281 #[test]
2282 fn test_file_scan_config_builder_defaults() {
2283 let file_schema = aggr_test_schema();
2284 let object_store_url = ObjectStoreUrl::parse("test:///").unwrap();
2285
2286 let table_schema = TableSchema::from(&file_schema);
2287
2288 let file_source: Arc<dyn FileSource> =
2289 Arc::new(MockSource::new(table_schema.clone()));
2290
2291 let config = FileScanConfigBuilder::new(
2293 object_store_url.clone(),
2294 Arc::clone(&file_source),
2295 )
2296 .build();
2297
2298 assert_eq!(config.object_store_url, object_store_url);
2300 assert_eq!(*config.file_schema(), file_schema);
2301 assert_eq!(config.limit, None);
2302 let expected_projection: Vec<usize> = (0..file_schema.fields().len()).collect();
2305 assert_eq!(
2306 config
2307 .file_source
2308 .projection()
2309 .as_ref()
2310 .map(|p| p.column_indices()),
2311 Some(expected_projection)
2312 );
2313 assert!(config.table_partition_cols().is_empty());
2314 assert!(config.file_groups.is_empty());
2315 assert_eq!(
2316 config.file_compression_type,
2317 FileCompressionType::UNCOMPRESSED
2318 );
2319 assert!(config.output_ordering.is_empty());
2320 assert!(config.constraints.is_empty());
2321
2322 assert_eq!(config.statistics().num_rows, Precision::Absent);
2324 assert_eq!(config.statistics().total_byte_size, Precision::Absent);
2325 assert_eq!(
2326 config.statistics().column_statistics.len(),
2327 file_schema.fields().len()
2328 );
2329 for stat in config.statistics().column_statistics {
2330 assert_eq!(stat.distinct_count, Precision::Absent);
2331 assert_eq!(stat.min_value, Precision::Absent);
2332 assert_eq!(stat.max_value, Precision::Absent);
2333 assert_eq!(stat.null_count, Precision::Absent);
2334 }
2335 }
2336
2337 #[test]
2338 fn test_file_scan_config_builder_new_from() {
2339 let schema = aggr_test_schema();
2340 let object_store_url = ObjectStoreUrl::parse("test:///").unwrap();
2341 let partition_cols = vec![Field::new(
2342 "date",
2343 wrap_partition_type_in_dict(DataType::Utf8),
2344 false,
2345 )];
2346 let file = PartitionedFile::new("test_file.parquet", 100);
2347
2348 let table_schema = TableSchemaBuilder::from(&schema)
2349 .with_table_partition_cols(
2350 partition_cols
2351 .iter()
2352 .map(|f| Arc::new(f.clone()))
2353 .collect::<Fields>(),
2354 )
2355 .build();
2356
2357 let file_source: Arc<dyn FileSource> =
2358 Arc::new(MockSource::new(table_schema.clone()));
2359
2360 let original_config = FileScanConfigBuilder::new(
2362 object_store_url.clone(),
2363 Arc::clone(&file_source),
2364 )
2365 .with_projection_indices(Some(vec![0, 2]))
2366 .unwrap()
2367 .with_limit(Some(10))
2368 .with_file(file.clone())
2369 .with_constraints(Constraints::default())
2370 .build();
2371
2372 let new_builder = FileScanConfigBuilder::from(original_config);
2374
2375 let new_config = new_builder.build();
2377
2378 let partition_cols = partition_cols.into_iter().map(Arc::new).collect::<Vec<_>>();
2380 assert_eq!(new_config.object_store_url, object_store_url);
2381 assert_eq!(*new_config.file_schema(), schema);
2382 assert_eq!(
2383 new_config
2384 .file_source
2385 .projection()
2386 .as_ref()
2387 .map(|p| p.column_indices()),
2388 Some(vec![0, 2])
2389 );
2390 assert_eq!(new_config.limit, Some(10));
2391 assert_eq!(
2392 *new_config.table_partition_cols(),
2393 Fields::from(partition_cols)
2394 );
2395 assert_eq!(new_config.file_groups.len(), 1);
2396 assert_eq!(new_config.file_groups[0].len(), 1);
2397 assert_eq!(
2398 new_config.file_groups[0][0].object_meta.location.as_ref(),
2399 "test_file.parquet"
2400 );
2401 assert_eq!(new_config.constraints, Constraints::default());
2402 }
2403
2404 #[test]
2405 fn test_split_groups_by_statistics_with_target_partitions() -> Result<()> {
2406 use datafusion_common::DFSchema;
2407 use datafusion_expr::{
2408 col, execution_props::ExecutionProps,
2409 physical_planning_context::PhysicalPlanningContext,
2410 };
2411
2412 let schema = Arc::new(Schema::new(vec![Field::new(
2413 "value",
2414 DataType::Float64,
2415 false,
2416 )]));
2417
2418 let exec_props = ExecutionProps::new();
2420 let df_schema = DFSchema::try_from_qualified_schema("test", schema.as_ref())?;
2421 let sort_expr = [col("value").sort(true, false)];
2422 let sort_ordering = sort_expr
2423 .map(|expr| {
2424 create_physical_sort_expr(
2425 &expr,
2426 &df_schema,
2427 &exec_props,
2428 &PhysicalPlanningContext::default(),
2429 )
2430 .unwrap()
2431 })
2432 .into();
2433
2434 struct TestCase {
2436 name: String,
2437 file_count: usize,
2438 overlap_factor: f64,
2439 target_partitions: usize,
2440 expected_partition_count: usize,
2441 }
2442
2443 let test_cases = vec![
2444 TestCase {
2446 name: "no_overlap_10_files_4_partitions".to_string(),
2447 file_count: 10,
2448 overlap_factor: 0.0,
2449 target_partitions: 4,
2450 expected_partition_count: 4,
2451 },
2452 TestCase {
2453 name: "medium_overlap_20_files_5_partitions".to_string(),
2454 file_count: 20,
2455 overlap_factor: 0.5,
2456 target_partitions: 5,
2457 expected_partition_count: 5,
2458 },
2459 TestCase {
2460 name: "high_overlap_30_files_3_partitions".to_string(),
2461 file_count: 30,
2462 overlap_factor: 0.8,
2463 target_partitions: 3,
2464 expected_partition_count: 7,
2465 },
2466 TestCase {
2468 name: "fewer_files_than_partitions".to_string(),
2469 file_count: 3,
2470 overlap_factor: 0.0,
2471 target_partitions: 10,
2472 expected_partition_count: 3, },
2474 TestCase {
2475 name: "single_file".to_string(),
2476 file_count: 1,
2477 overlap_factor: 0.0,
2478 target_partitions: 5,
2479 expected_partition_count: 1, },
2481 TestCase {
2482 name: "empty_files".to_string(),
2483 file_count: 0,
2484 overlap_factor: 0.0,
2485 target_partitions: 3,
2486 expected_partition_count: 0, },
2488 ];
2489
2490 for case in test_cases {
2491 println!("Running test case: {}", case.name);
2492
2493 let file_groups = generate_test_files(case.file_count, case.overlap_factor);
2495
2496 let result =
2498 FileScanConfig::split_groups_by_statistics_with_target_partitions(
2499 &schema,
2500 &file_groups,
2501 &sort_ordering,
2502 case.target_partitions,
2503 )?;
2504
2505 println!(
2507 "Created {} partitions (target was {})",
2508 result.len(),
2509 case.target_partitions
2510 );
2511
2512 assert_eq!(
2514 result.len(),
2515 case.expected_partition_count,
2516 "Case '{}': Unexpected partition count",
2517 case.name
2518 );
2519
2520 assert!(
2522 verify_sort_integrity(&result),
2523 "Case '{}': Files within partitions are not properly ordered",
2524 case.name
2525 );
2526
2527 if case.file_count > 1 && case.expected_partition_count > 1 {
2529 let group_sizes: Vec<usize> = result.iter().map(FileGroup::len).collect();
2530 let max_size = *group_sizes.iter().max().unwrap();
2531 let min_size = *group_sizes.iter().min().unwrap();
2532
2533 let avg_files_per_partition =
2535 case.file_count as f64 / case.expected_partition_count as f64;
2536 assert!(
2537 (max_size as f64) < 2.0 * avg_files_per_partition,
2538 "Case '{}': Unbalanced distribution. Max partition size {} exceeds twice the average {}",
2539 case.name,
2540 max_size,
2541 avg_files_per_partition
2542 );
2543
2544 println!("Distribution - min files: {min_size}, max files: {max_size}");
2545 }
2546 }
2547
2548 let empty_groups: Vec<FileGroup> = vec![];
2550 let err = FileScanConfig::split_groups_by_statistics_with_target_partitions(
2551 &schema,
2552 &empty_groups,
2553 &sort_ordering,
2554 0,
2555 )
2556 .unwrap_err();
2557
2558 assert!(
2559 err.to_string()
2560 .contains("target_partitions must be greater than 0"),
2561 "Expected error for zero target partitions"
2562 );
2563
2564 Ok(())
2565 }
2566
2567 #[test]
2568 fn test_partition_statistics_projection() {
2569 use crate::source::DataSourceExec;
2575 use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext};
2576
2577 let schema = Arc::new(Schema::new(vec![
2578 Field::new("col0", DataType::Int32, false),
2579 Field::new("col1", DataType::Int32, false),
2580 Field::new("col2", DataType::Int32, false),
2581 Field::new("col3", DataType::Int32, false),
2582 ]));
2583
2584 let file_group_stats = Statistics {
2586 num_rows: Precision::Exact(100),
2587 total_byte_size: Precision::Exact(1024),
2588 column_statistics: vec![
2589 ColumnStatistics {
2590 null_count: Precision::Exact(0),
2591 ..ColumnStatistics::new_unknown()
2592 },
2593 ColumnStatistics {
2594 null_count: Precision::Exact(5),
2595 ..ColumnStatistics::new_unknown()
2596 },
2597 ColumnStatistics {
2598 null_count: Precision::Exact(10),
2599 ..ColumnStatistics::new_unknown()
2600 },
2601 ColumnStatistics {
2602 null_count: Precision::Exact(15),
2603 ..ColumnStatistics::new_unknown()
2604 },
2605 ],
2606 };
2607
2608 let file_group = FileGroup::new(vec![PartitionedFile::new("test.parquet", 1024)])
2610 .with_statistics(Arc::new(file_group_stats));
2611
2612 let table_schema = TableSchema::from(&schema);
2613
2614 let config = FileScanConfigBuilder::new(
2616 ObjectStoreUrl::parse("test:///").unwrap(),
2617 Arc::new(MockSource::new(table_schema.clone())),
2618 )
2619 .with_projection_indices(Some(vec![0, 2]))
2620 .unwrap() .with_file_groups(vec![file_group])
2622 .build();
2623
2624 let exec = DataSourceExec::from_data_source(config);
2626
2627 let partition_stats = StatisticsContext::new()
2629 .compute(
2630 exec.as_ref(),
2631 &StatisticsArgs::new().with_partition(Some(0)),
2632 )
2633 .unwrap();
2634
2635 assert_eq!(
2637 partition_stats.column_statistics.len(),
2638 2,
2639 "Expected 2 column statistics (projected), but got {}",
2640 partition_stats.column_statistics.len()
2641 );
2642
2643 assert_eq!(
2645 partition_stats.column_statistics[0].null_count,
2646 Precision::Exact(0),
2647 "First projected column should be col0 with 0 nulls"
2648 );
2649 assert_eq!(
2650 partition_stats.column_statistics[1].null_count,
2651 Precision::Exact(10),
2652 "Second projected column should be col2 with 10 nulls"
2653 );
2654
2655 assert_eq!(partition_stats.num_rows, Precision::Exact(100));
2657 assert_eq!(partition_stats.total_byte_size, Precision::Exact(800));
2658 }
2659
2660 #[test]
2661 fn test_statistics_with_filter() {
2662 assert_num_rows_with_filter(Precision::Absent, Precision::Absent);
2663 assert_num_rows_with_filter(Precision::Exact(100), Precision::Inexact(100));
2664 assert_num_rows_with_filter(Precision::Inexact(100), Precision::Inexact(100));
2665 assert_num_rows_with_filter(Precision::Exact(0), Precision::Exact(0));
2666
2667 fn assert_num_rows_with_filter(
2670 input_num_rows: Precision<usize>,
2671 expected_num_rows: Precision<usize>,
2672 ) {
2673 let schema = Arc::new(Schema::new(vec![Field::new(
2674 "col0",
2675 DataType::Int32,
2676 false,
2677 )]));
2678
2679 let stats =
2680 Statistics::new_unknown(schema.as_ref()).with_num_rows(input_num_rows);
2681 let file_group =
2682 FileGroup::new(vec![PartitionedFile::new("test.parquet", 1024)]);
2683
2684 let table_schema = TableSchema::from(&schema);
2685 let config = FileScanConfigBuilder::new(
2686 ObjectStoreUrl::parse("test:///").unwrap(),
2687 Arc::new(MockSource::new(table_schema.clone()).with_filter(Arc::new(
2688 Literal::new(ScalarValue::Boolean(Some(true))),
2689 ))),
2690 )
2691 .with_file_groups(vec![file_group])
2692 .with_statistics(stats)
2693 .build();
2694
2695 assert_eq!(config.statistics().num_rows, expected_num_rows,);
2696 }
2697 }
2698
2699 #[tokio::test]
2706 async fn reset_state_recreates_shared_work_source() -> Result<()> {
2707 let schema = Arc::new(Schema::new(vec![Field::new(
2708 "value",
2709 DataType::Int32,
2710 false,
2711 )]));
2712 let file_source = Arc::new(
2713 MockSource::new(Arc::clone(&schema))
2714 .with_file_opener(Arc::new(ResetStateTestFileOpener { schema })),
2715 );
2716
2717 let config =
2718 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
2719 .with_file_group(FileGroup::new(vec![
2720 PartitionedFile::new("file1.parquet", 100),
2721 PartitionedFile::new("file2.parquet", 100),
2722 ]))
2723 .build();
2724
2725 let exec: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(config);
2726 let task_ctx = Arc::new(TaskContext::default());
2727
2728 let first_run = collect(Arc::clone(&exec), Arc::clone(&task_ctx)).await?;
2731 let reset_exec = exec.reset_state()?;
2732 let second_run = collect(reset_exec, task_ctx).await?;
2733
2734 let expected = [
2735 "+-------+",
2736 "| value |",
2737 "+-------+",
2738 "| 1 |",
2739 "| 2 |",
2740 "+-------+",
2741 ];
2742 assert_batches_eq!(expected, &first_run);
2743 assert_batches_eq!(expected, &second_run);
2744
2745 Ok(())
2746 }
2747
2748 #[derive(Debug)]
2751 struct ResetStateTestFileOpener {
2752 schema: SchemaRef,
2753 }
2754
2755 impl crate::file_stream::FileOpener for ResetStateTestFileOpener {
2756 fn open(
2757 &self,
2758 file: PartitionedFile,
2759 ) -> Result<crate::file_stream::FileOpenFuture> {
2760 let value = file
2761 .object_meta
2762 .location
2763 .as_ref()
2764 .trim_start_matches("file")
2765 .trim_end_matches(".parquet")
2766 .parse::<i32>()
2767 .expect("invalid test file name");
2768 let schema = Arc::clone(&self.schema);
2769 Ok(async move {
2770 let batch = RecordBatch::try_new(
2771 schema,
2772 vec![Arc::new(Int32Array::from(vec![value]))],
2773 )
2774 .expect("test batch should be valid");
2775 Ok(stream::iter(vec![Ok(batch)]).boxed())
2776 }
2777 .boxed())
2778 }
2779 }
2780
2781 #[test]
2782 fn test_output_partitioning_not_partitioned_by_file_group() {
2783 let file_schema = aggr_test_schema();
2784 let partition_col =
2785 Field::new("date", wrap_partition_type_in_dict(DataType::Utf8), false);
2786
2787 let config = config_for_projection(
2788 Arc::clone(&file_schema),
2789 None,
2790 Statistics::new_unknown(&file_schema),
2791 vec![partition_col],
2792 );
2793
2794 let partitioning = config.output_partitioning();
2796 assert!(matches!(partitioning, Partitioning::UnknownPartitioning(_)));
2797 }
2798
2799 #[test]
2800 fn test_declared_output_partitioning_projects_with_scan() {
2801 let file_schema = aggr_test_schema();
2802 let output_partitioning =
2803 Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4);
2804
2805 let mut config = config_for_projection(
2806 Arc::clone(&file_schema),
2807 Some(vec![1, 2]),
2808 Statistics::new_unknown(&file_schema),
2809 vec![],
2810 );
2811 config.file_groups = vec![
2812 FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]),
2813 FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]),
2814 FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]),
2815 FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]),
2816 ];
2817 config.output_partitioning = Some(output_partitioning);
2818
2819 match config.output_partitioning() {
2820 Partitioning::Hash(exprs, num_partitions) => {
2821 assert_eq!(num_partitions, 4);
2822 assert_eq!(exprs.len(), 1);
2823 let column = exprs[0].downcast_ref::<Column>().unwrap();
2824 assert_eq!(column.name(), "c2");
2825 assert_eq!(column.index(), 0);
2826 }
2827 _ => panic!("Expected Hash partitioning"),
2828 }
2829
2830 let mut config = config_for_projection(
2831 Arc::clone(&file_schema),
2832 Some(vec![2]),
2833 Statistics::new_unknown(&file_schema),
2834 vec![],
2835 );
2836 config.file_groups = vec![
2837 FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]),
2838 FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]),
2839 FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]),
2840 FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]),
2841 ];
2842 config.output_partitioning =
2843 Some(Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4));
2844
2845 assert!(matches!(
2846 config.output_partitioning(),
2847 Partitioning::UnknownPartitioning(4)
2848 ));
2849 }
2850
2851 #[test]
2852 fn test_output_partitioning_no_partition_columns() {
2853 let file_schema = aggr_test_schema();
2854 let config = config_for_projection(
2855 Arc::clone(&file_schema),
2856 None,
2857 Statistics::new_unknown(&file_schema),
2858 vec![], );
2860
2861 let partitioning = config.output_partitioning();
2862 assert!(matches!(partitioning, Partitioning::UnknownPartitioning(_)));
2863 }
2864
2865 #[test]
2866 fn test_output_partitioning_with_partition_columns() {
2867 let file_schema = aggr_test_schema();
2868
2869 let single_partition_col = vec![Field::new(
2871 "date",
2872 wrap_partition_type_in_dict(DataType::Utf8),
2873 false,
2874 )];
2875
2876 let mut config = config_for_projection(
2877 Arc::clone(&file_schema),
2878 None,
2879 Statistics::new_unknown(&file_schema),
2880 single_partition_col,
2881 );
2882 config.file_groups = vec![
2883 FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]),
2884 FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]),
2885 FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]),
2886 ];
2887 config.output_partitioning = output_partitioning_from_partition_fields(
2888 config.file_source.table_schema().table_schema(),
2889 config.table_partition_cols(),
2890 config.file_groups.len(),
2891 );
2892
2893 let partitioning = config.output_partitioning();
2894 match partitioning {
2895 Partitioning::Hash(exprs, num_partitions) => {
2896 assert_eq!(num_partitions, 3);
2897 assert_eq!(exprs.len(), 1);
2898 assert_eq!(exprs[0].downcast_ref::<Column>().unwrap().name(), "date");
2899 }
2900 _ => panic!("Expected Hash partitioning"),
2901 }
2902
2903 let multiple_partition_cols = vec![
2905 Field::new("year", wrap_partition_type_in_dict(DataType::Utf8), false),
2906 Field::new("month", wrap_partition_type_in_dict(DataType::Utf8), false),
2907 ];
2908
2909 config = config_for_projection(
2910 Arc::clone(&file_schema),
2911 None,
2912 Statistics::new_unknown(&file_schema),
2913 multiple_partition_cols,
2914 );
2915 config.file_groups = vec![
2916 FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]),
2917 FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]),
2918 ];
2919 config.output_partitioning = output_partitioning_from_partition_fields(
2920 config.file_source.table_schema().table_schema(),
2921 config.table_partition_cols(),
2922 config.file_groups.len(),
2923 );
2924
2925 let partitioning = config.output_partitioning();
2926 match partitioning {
2927 Partitioning::Hash(exprs, num_partitions) => {
2928 assert_eq!(num_partitions, 2);
2929 assert_eq!(exprs.len(), 2);
2930 let col_names: Vec<_> = exprs
2931 .iter()
2932 .map(|e| e.downcast_ref::<Column>().unwrap().name())
2933 .collect();
2934 assert_eq!(col_names, vec!["year", "month"]);
2935 }
2936 _ => panic!("Expected Hash partitioning"),
2937 }
2938 }
2939
2940 #[test]
2941 fn try_pushdown_sort_reverses_file_groups_only_when_requested_is_reverse()
2942 -> Result<()> {
2943 let file_schema =
2944 Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
2945
2946 let table_schema = TableSchema::from(&file_schema);
2947 let file_source = Arc::new(InexactSortPushdownSource::new(table_schema));
2948
2949 let file_groups = vec![FileGroup::new(vec![
2950 PartitionedFile::new("file1", 1),
2951 PartitionedFile::new("file2", 1),
2952 ])];
2953
2954 let sort_expr_asc = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
2955 let config =
2956 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
2957 .with_file_groups(file_groups)
2958 .with_output_ordering(vec![
2959 LexOrdering::new(vec![sort_expr_asc.clone()]).unwrap(),
2960 ])
2961 .build();
2962
2963 let requested_asc = vec![sort_expr_asc.clone()];
2964 let result = config.try_pushdown_sort(&requested_asc)?;
2965 let SortOrderPushdownResult::Inexact { inner } = result else {
2966 panic!("Expected Inexact result");
2967 };
2968 let pushed_config = inner
2969 .downcast_ref::<FileScanConfig>()
2970 .expect("Expected FileScanConfig");
2971 let pushed_files = pushed_config.file_groups[0].files();
2972 assert_eq!(pushed_files[0].object_meta.location.as_ref(), "file1");
2973 assert_eq!(pushed_files[1].object_meta.location.as_ref(), "file2");
2974
2975 let requested_desc = vec![sort_expr_asc.reverse()];
2976 let result = config.try_pushdown_sort(&requested_desc)?;
2977 let SortOrderPushdownResult::Inexact { inner } = result else {
2978 panic!("Expected Inexact result");
2979 };
2980 let pushed_config = inner
2981 .downcast_ref::<FileScanConfig>()
2982 .expect("Expected FileScanConfig");
2983 let pushed_files = pushed_config.file_groups[0].files();
2984 assert_eq!(pushed_files[0].object_meta.location.as_ref(), "file2");
2985 assert_eq!(pushed_files[1].object_meta.location.as_ref(), "file1");
2986
2987 Ok(())
2988 }
2989
2990 fn make_file_with_stats(name: &str, min: f64, max: f64) -> PartitionedFile {
2991 PartitionedFile::new(name.to_string(), 1024).with_statistics(Arc::new(
2992 Statistics {
2993 num_rows: Precision::Exact(100),
2994 total_byte_size: Precision::Exact(1024),
2995 column_statistics: vec![ColumnStatistics {
2996 null_count: Precision::Exact(0),
2997 min_value: Precision::Exact(ScalarValue::Float64(Some(min))),
2998 max_value: Precision::Exact(ScalarValue::Float64(Some(max))),
2999 ..Default::default()
3000 }],
3001 },
3002 ))
3003 }
3004
3005 #[derive(Clone)]
3006 struct ExactSortPushdownSource {
3007 metrics: ExecutionPlanMetricsSet,
3008 table_schema: TableSchema,
3009 }
3010
3011 impl ExactSortPushdownSource {
3012 fn new(table_schema: TableSchema) -> Self {
3013 Self {
3014 metrics: ExecutionPlanMetricsSet::new(),
3015 table_schema,
3016 }
3017 }
3018 }
3019
3020 impl FileSource for ExactSortPushdownSource {
3021 fn create_file_opener(
3022 &self,
3023 _object_store: Arc<dyn ObjectStore>,
3024 _base_config: &FileScanConfig,
3025 _partition: usize,
3026 ) -> Result<Arc<dyn crate::file_stream::FileOpener>> {
3027 unimplemented!()
3028 }
3029
3030 fn table_schema(&self) -> &TableSchema {
3031 &self.table_schema
3032 }
3033
3034 fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> {
3035 Arc::new(self.clone())
3036 }
3037
3038 fn metrics(&self) -> &ExecutionPlanMetricsSet {
3039 &self.metrics
3040 }
3041
3042 fn file_type(&self) -> &str {
3043 "mock_exact"
3044 }
3045
3046 fn try_pushdown_sort(
3047 &self,
3048 _order: &[PhysicalSortExpr],
3049 _eq_properties: &EquivalenceProperties,
3050 ) -> Result<SortOrderPushdownResult<Arc<dyn FileSource>>> {
3051 Ok(SortOrderPushdownResult::Exact {
3052 inner: Arc::new(self.clone()) as Arc<dyn FileSource>,
3053 })
3054 }
3055
3056 fn apply_expressions(
3057 &self,
3058 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
3059 ) -> Result<TreeNodeRecursion> {
3060 Ok(TreeNodeRecursion::Continue)
3061 }
3062 }
3063
3064 #[test]
3065 fn sort_pushdown_unsupported_source_files_get_sorted() -> Result<()> {
3066 let file_schema =
3067 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3068 let table_schema = TableSchema::from(&file_schema);
3069 let file_source = Arc::new(MockSource::new(table_schema));
3070
3071 let file_groups = vec![FileGroup::new(vec![
3072 make_file_with_stats("file3", 20.0, 30.0),
3073 make_file_with_stats("file1", 0.0, 9.0),
3074 make_file_with_stats("file2", 10.0, 19.0),
3075 ])];
3076
3077 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3078 let config =
3079 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3080 .with_file_groups(file_groups)
3081 .build();
3082
3083 let result = config.try_pushdown_sort(&[sort_expr])?;
3084 let SortOrderPushdownResult::Inexact { inner } = result else {
3085 panic!("Expected Inexact result, got {result:?}");
3086 };
3087 let pushed_config = inner
3088 .downcast_ref::<FileScanConfig>()
3089 .expect("Expected FileScanConfig");
3090 let files = pushed_config.file_groups[0].files();
3091 assert_eq!(files[0].object_meta.location.as_ref(), "file1");
3092 assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3093 assert_eq!(files[2].object_meta.location.as_ref(), "file3");
3094 assert!(pushed_config.output_ordering.is_empty());
3095 Ok(())
3096 }
3097
3098 #[test]
3099 fn sort_pushdown_unsupported_source_already_sorted() -> Result<()> {
3100 let file_schema =
3101 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3102 let table_schema = TableSchema::from(&file_schema);
3103 let file_source = Arc::new(MockSource::new(table_schema));
3104
3105 let file_groups = vec![FileGroup::new(vec![
3106 make_file_with_stats("file1", 0.0, 9.0),
3107 make_file_with_stats("file2", 10.0, 19.0),
3108 make_file_with_stats("file3", 20.0, 30.0),
3109 ])];
3110
3111 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3112 let config =
3113 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3114 .with_file_groups(file_groups)
3115 .build();
3116
3117 let result = config.try_pushdown_sort(&[sort_expr])?;
3118 assert!(matches!(result, SortOrderPushdownResult::Unsupported));
3119 Ok(())
3120 }
3121
3122 #[test]
3123 fn sort_pushdown_unsupported_source_descending_sort() -> Result<()> {
3124 let file_schema =
3125 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3126 let table_schema = TableSchema::from(&file_schema);
3127 let file_source = Arc::new(MockSource::new(table_schema));
3128
3129 let file_groups = vec![FileGroup::new(vec![
3130 make_file_with_stats("file1", 0.0, 9.0),
3131 make_file_with_stats("file3", 20.0, 30.0),
3132 make_file_with_stats("file2", 10.0, 19.0),
3133 ])];
3134
3135 let sort_expr = PhysicalSortExpr::new(
3136 Arc::new(Column::new("a", 0)),
3137 arrow::compute::SortOptions {
3138 descending: true,
3139 nulls_first: true,
3140 },
3141 );
3142 let config =
3143 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3144 .with_file_groups(file_groups)
3145 .build();
3146
3147 let result = config.try_pushdown_sort(&[sort_expr])?;
3148 let SortOrderPushdownResult::Inexact { inner } = result else {
3149 panic!("Expected Inexact result");
3150 };
3151 let pushed_config = inner
3152 .downcast_ref::<FileScanConfig>()
3153 .expect("Expected FileScanConfig");
3154 let files = pushed_config.file_groups[0].files();
3155 assert_eq!(files[0].object_meta.location.as_ref(), "file3");
3156 assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3157 assert_eq!(files[2].object_meta.location.as_ref(), "file1");
3158 Ok(())
3159 }
3160
3161 #[test]
3162 fn sort_pushdown_exact_source_non_overlapping_returns_exact() -> Result<()> {
3163 let file_schema =
3164 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3165 let table_schema = TableSchema::from(&file_schema);
3166 let file_source = Arc::new(ExactSortPushdownSource::new(table_schema));
3167
3168 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3169
3170 let file_groups = vec![FileGroup::new(vec![
3171 make_file_with_stats("file1", 0.0, 9.0),
3172 make_file_with_stats("file2", 10.0, 19.0),
3173 make_file_with_stats("file3", 20.0, 30.0),
3174 ])];
3175
3176 let config =
3177 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3178 .with_file_groups(file_groups)
3179 .with_output_ordering(vec![
3180 LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3181 ])
3182 .build();
3183
3184 let result = config.try_pushdown_sort(&[sort_expr])?;
3185 let SortOrderPushdownResult::Exact { inner } = result else {
3186 panic!("Expected Exact result, got {result:?}");
3187 };
3188 let pushed_config = inner
3189 .downcast_ref::<FileScanConfig>()
3190 .expect("Expected FileScanConfig");
3191 assert!(!pushed_config.output_ordering.is_empty());
3192 Ok(())
3193 }
3194
3195 #[test]
3196 fn sort_pushdown_exact_source_overlapping_downgraded_to_inexact() -> Result<()> {
3197 let file_schema =
3198 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3199 let table_schema = TableSchema::from(&file_schema);
3200 let file_source = Arc::new(ExactSortPushdownSource::new(table_schema));
3201
3202 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3203
3204 let file_groups = vec![FileGroup::new(vec![
3205 make_file_with_stats("file1", 0.0, 15.0),
3206 make_file_with_stats("file2", 10.0, 25.0),
3207 make_file_with_stats("file3", 20.0, 30.0),
3208 ])];
3209
3210 let config =
3211 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3212 .with_file_groups(file_groups)
3213 .with_output_ordering(vec![
3214 LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3215 ])
3216 .build();
3217
3218 let result = config.try_pushdown_sort(&[sort_expr])?;
3219 let SortOrderPushdownResult::Inexact { inner } = result else {
3220 panic!("Expected Inexact (downgraded), got {result:?}");
3221 };
3222 let pushed_config = inner
3223 .downcast_ref::<FileScanConfig>()
3224 .expect("Expected FileScanConfig");
3225 assert!(pushed_config.output_ordering.is_empty());
3226 Ok(())
3227 }
3228
3229 #[test]
3230 fn sort_pushdown_exact_source_out_of_order_returns_exact() -> Result<()> {
3231 let file_schema =
3232 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3233 let table_schema = TableSchema::from(&file_schema);
3234 let file_source = Arc::new(ExactSortPushdownSource::new(table_schema));
3235
3236 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3237
3238 let file_groups = vec![FileGroup::new(vec![
3239 make_file_with_stats("file3", 20.0, 30.0),
3240 make_file_with_stats("file1", 0.0, 9.0),
3241 make_file_with_stats("file2", 10.0, 19.0),
3242 ])];
3243
3244 let config =
3245 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3246 .with_file_groups(file_groups)
3247 .with_output_ordering(vec![
3248 LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3249 ])
3250 .build();
3251
3252 let result = config.try_pushdown_sort(&[sort_expr])?;
3253 let SortOrderPushdownResult::Exact { inner } = result else {
3254 panic!("Expected Exact result, got {result:?}");
3255 };
3256 let pushed_config = inner
3257 .downcast_ref::<FileScanConfig>()
3258 .expect("Expected FileScanConfig");
3259 let files = pushed_config.file_groups[0].files();
3260 assert_eq!(files[0].object_meta.location.as_ref(), "file1");
3261 assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3262 assert_eq!(files[2].object_meta.location.as_ref(), "file3");
3263 assert!(!pushed_config.output_ordering.is_empty());
3264 Ok(())
3265 }
3266
3267 #[test]
3268 fn sort_pushdown_unsupported_source_single_file_groups() -> Result<()> {
3269 let file_schema =
3270 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3271 let table_schema = TableSchema::from(&file_schema);
3272 let file_source = Arc::new(MockSource::new(table_schema));
3273
3274 let file_groups = vec![
3275 FileGroup::new(vec![make_file_with_stats("file1", 0.0, 9.0)]),
3276 FileGroup::new(vec![make_file_with_stats("file2", 10.0, 19.0)]),
3277 ];
3278
3279 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3280 let config =
3281 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3282 .with_file_groups(file_groups)
3283 .build();
3284
3285 let result = config.try_pushdown_sort(&[sort_expr])?;
3286 assert!(
3287 matches!(result, SortOrderPushdownResult::Unsupported),
3288 "Expected Unsupported for single-file groups"
3289 );
3290 Ok(())
3291 }
3292
3293 #[test]
3294 fn sort_pushdown_unsupported_source_multiple_groups() -> Result<()> {
3295 let file_schema =
3296 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3297 let table_schema = TableSchema::from(&file_schema);
3298 let file_source = Arc::new(MockSource::new(table_schema));
3299
3300 let file_groups = vec![
3301 FileGroup::new(vec![
3302 make_file_with_stats("file_b", 10.0, 19.0),
3303 make_file_with_stats("file_a", 0.0, 9.0),
3304 ]),
3305 FileGroup::new(vec![
3306 make_file_with_stats("file_d", 30.0, 39.0),
3307 make_file_with_stats("file_c", 20.0, 29.0),
3308 ]),
3309 ];
3310
3311 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3312 let config =
3313 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3314 .with_file_groups(file_groups)
3315 .build();
3316
3317 let result = config.try_pushdown_sort(&[sort_expr])?;
3318 let SortOrderPushdownResult::Inexact { inner } = result else {
3319 panic!("Expected Inexact result");
3320 };
3321 let pushed_config = inner
3322 .downcast_ref::<FileScanConfig>()
3323 .expect("Expected FileScanConfig");
3324 let files0 = pushed_config.file_groups[0].files();
3325 assert_eq!(files0[0].object_meta.location.as_ref(), "file_a");
3326 assert_eq!(files0[1].object_meta.location.as_ref(), "file_b");
3327 let files1 = pushed_config.file_groups[1].files();
3328 assert_eq!(files1[0].object_meta.location.as_ref(), "file_c");
3329 assert_eq!(files1[1].object_meta.location.as_ref(), "file_d");
3330 Ok(())
3331 }
3332
3333 #[test]
3334 fn sort_pushdown_unsupported_source_partial_statistics() -> Result<()> {
3335 let file_schema =
3336 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3337 let table_schema = TableSchema::from(&file_schema);
3338 let file_source = Arc::new(MockSource::new(table_schema));
3339
3340 let file_groups = vec![
3341 FileGroup::new(vec![
3342 make_file_with_stats("file_b", 10.0, 19.0),
3343 make_file_with_stats("file_a", 0.0, 9.0),
3344 ]),
3345 FileGroup::new(vec![
3346 PartitionedFile::new("file_d".to_string(), 1024),
3347 PartitionedFile::new("file_c".to_string(), 1024),
3348 ]),
3349 ];
3350
3351 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3352 let config =
3353 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3354 .with_file_groups(file_groups)
3355 .build();
3356
3357 let result = config.try_pushdown_sort(&[sort_expr])?;
3358 let SortOrderPushdownResult::Inexact { inner } = result else {
3359 panic!("Expected Inexact result");
3360 };
3361 let pushed_config = inner
3362 .downcast_ref::<FileScanConfig>()
3363 .expect("Expected FileScanConfig");
3364 let files0 = pushed_config.file_groups[0].files();
3365 assert_eq!(files0[0].object_meta.location.as_ref(), "file_a");
3366 assert_eq!(files0[1].object_meta.location.as_ref(), "file_b");
3367 let files1 = pushed_config.file_groups[1].files();
3368 assert_eq!(files1[0].object_meta.location.as_ref(), "file_d");
3369 assert_eq!(files1[1].object_meta.location.as_ref(), "file_c");
3370 Ok(())
3371 }
3372
3373 #[test]
3374 fn sort_pushdown_inexact_source_with_statistics_sorting() -> Result<()> {
3375 let file_schema =
3376 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3377 let table_schema = TableSchema::from(&file_schema);
3378 let file_source = Arc::new(InexactSortPushdownSource::new(table_schema));
3379
3380 let file_groups = vec![FileGroup::new(vec![
3381 make_file_with_stats("file2", 10.0, 19.0),
3382 make_file_with_stats("file1", 0.0, 9.0),
3383 ])];
3384
3385 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3386 let config =
3387 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3388 .with_file_groups(file_groups)
3389 .build();
3390
3391 let result = config.try_pushdown_sort(&[sort_expr])?;
3392 let SortOrderPushdownResult::Inexact { inner } = result else {
3393 panic!("Expected Inexact result");
3394 };
3395 let pushed_config = inner
3396 .downcast_ref::<FileScanConfig>()
3397 .expect("Expected FileScanConfig");
3398 let files = pushed_config.file_groups[0].files();
3399 assert_eq!(files[0].object_meta.location.as_ref(), "file1");
3400 assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3401 assert!(pushed_config.output_ordering.is_empty());
3402 Ok(())
3403 }
3404
3405 #[test]
3406 fn sort_pushdown_exact_multi_group_preserves_parallelism() -> Result<()> {
3407 let file_schema =
3413 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3414 let table_schema = TableSchema::from(&file_schema);
3415 let file_source = Arc::new(ExactSortPushdownSource::new(table_schema));
3416
3417 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3418
3419 let file_groups = vec![
3423 FileGroup::new(vec![
3424 make_file_with_stats("file_01", 0.0, 9.0),
3425 make_file_with_stats("file_03", 20.0, 29.0),
3426 ]),
3427 FileGroup::new(vec![
3428 make_file_with_stats("file_02", 10.0, 19.0),
3429 make_file_with_stats("file_04", 30.0, 39.0),
3430 ]),
3431 ];
3432
3433 let config =
3434 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3435 .with_file_groups(file_groups)
3436 .with_output_ordering(vec![
3437 LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3438 ])
3439 .build();
3440
3441 let result = config.try_pushdown_sort(&[sort_expr])?;
3442 let SortOrderPushdownResult::Exact { inner } = result else {
3443 panic!("Expected Exact result, got {result:?}");
3444 };
3445 let pushed_config = inner
3446 .downcast_ref::<FileScanConfig>()
3447 .expect("Expected FileScanConfig");
3448
3449 assert_eq!(pushed_config.file_groups.len(), 2);
3451
3452 let files0 = pushed_config.file_groups[0].files();
3455 assert_eq!(files0[0].object_meta.location.as_ref(), "file_01");
3456 assert_eq!(files0[1].object_meta.location.as_ref(), "file_03");
3457 let files1 = pushed_config.file_groups[1].files();
3458 assert_eq!(files1[0].object_meta.location.as_ref(), "file_02");
3459 assert_eq!(files1[1].object_meta.location.as_ref(), "file_04");
3460
3461 assert!(!pushed_config.output_ordering.is_empty());
3463 Ok(())
3464 }
3465
3466 #[test]
3467 fn sort_pushdown_reverse_preserves_file_order_with_stats() -> Result<()> {
3468 let file_schema =
3471 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)]));
3472 let table_schema = TableSchema::from(&file_schema);
3473 let file_source = Arc::new(InexactSortPushdownSource::new(table_schema));
3474
3475 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3476
3477 let file_groups = vec![FileGroup::new(vec![
3479 make_file_with_stats("file1", 0.0, 9.0),
3480 make_file_with_stats("file2", 10.0, 19.0),
3481 make_file_with_stats("file3", 20.0, 30.0),
3482 ])];
3483
3484 let config =
3485 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3486 .with_file_groups(file_groups)
3487 .with_output_ordering(vec![
3488 LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3489 ])
3490 .build();
3491
3492 let result = config.try_pushdown_sort(&[sort_expr.reverse()])?;
3494 let SortOrderPushdownResult::Inexact { inner } = result else {
3495 panic!("Expected Inexact for reverse scan, got {result:?}");
3496 };
3497 let pushed_config = inner
3498 .downcast_ref::<FileScanConfig>()
3499 .expect("Expected FileScanConfig");
3500
3501 let files = pushed_config.file_groups[0].files();
3503 assert_eq!(files[0].object_meta.location.as_ref(), "file3");
3504 assert_eq!(files[1].object_meta.location.as_ref(), "file2");
3505 assert_eq!(files[2].object_meta.location.as_ref(), "file1");
3506
3507 assert!(pushed_config.output_ordering.is_empty());
3509 Ok(())
3510 }
3511
3512 fn make_file_with_null_stats(
3514 name: &str,
3515 min: f64,
3516 max: f64,
3517 null_count: usize,
3518 ) -> PartitionedFile {
3519 PartitionedFile::new(name.to_string(), 1024).with_statistics(Arc::new(
3520 Statistics {
3521 num_rows: Precision::Exact(100),
3522 total_byte_size: Precision::Exact(1024),
3523 column_statistics: vec![ColumnStatistics {
3524 null_count: Precision::Exact(null_count),
3525 min_value: Precision::Exact(ScalarValue::Float64(Some(min))),
3526 max_value: Precision::Exact(ScalarValue::Float64(Some(max))),
3527 ..Default::default()
3528 }],
3529 },
3530 ))
3531 }
3532
3533 #[test]
3534 fn sort_pushdown_unsupported_with_nulls_does_not_upgrade_to_exact() -> Result<()> {
3535 let file_schema =
3538 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)]));
3539 let table_schema = TableSchema::from(&file_schema);
3540 let file_source = Arc::new(MockSource::new(table_schema));
3541
3542 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3543
3544 let file_groups = vec![FileGroup::new(vec![
3546 make_file_with_null_stats("b_no_nulls", 10.0, 19.0, 0),
3547 make_file_with_null_stats("a_with_nulls", 0.0, 9.0, 5), ])];
3549
3550 let config =
3551 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3552 .with_file_groups(file_groups)
3553 .with_output_ordering(vec![
3554 LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3555 ])
3556 .build();
3557
3558 let result = config.try_pushdown_sort(&[sort_expr])?;
3559 assert!(
3561 matches!(result, SortOrderPushdownResult::Inexact { .. }),
3562 "Expected Inexact due to NULLs, got {result:?}"
3563 );
3564 Ok(())
3565 }
3566
3567 #[test]
3568 fn sort_pushdown_unsupported_no_nulls_upgrades_to_exact() -> Result<()> {
3569 let file_schema =
3571 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)]));
3572 let table_schema = TableSchema::from(&file_schema);
3573 let file_source = Arc::new(MockSource::new(table_schema));
3574
3575 let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
3576
3577 let file_groups = vec![FileGroup::new(vec![
3578 make_file_with_null_stats("b_high", 10.0, 19.0, 0),
3579 make_file_with_null_stats("a_low", 0.0, 9.0, 0),
3580 ])];
3581
3582 let config =
3583 FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
3584 .with_file_groups(file_groups)
3585 .with_output_ordering(vec![
3586 LexOrdering::new(vec![sort_expr.clone()]).unwrap(),
3587 ])
3588 .build();
3589
3590 let result = config.try_pushdown_sort(&[sort_expr])?;
3591 assert!(
3592 matches!(result, SortOrderPushdownResult::Exact { .. }),
3593 "Expected Exact (no NULLs), got {result:?}"
3594 );
3595 Ok(())
3596 }
3597
3598 fn make_projection(pairs: Vec<(Arc<dyn PhysicalExpr>, &str)>) -> ProjectionExprs {
3600 ProjectionExprs::new(
3601 pairs
3602 .into_iter()
3603 .map(|(expr, alias)| ProjectionExpr::new(expr, alias)),
3604 )
3605 }
3606
3607 fn make_volatile_expr() -> Arc<dyn PhysicalExpr> {
3610 use datafusion_common::config::ConfigOptions;
3611 use datafusion_expr::ScalarUDF;
3612 use datafusion_functions::math::random::RandomFunc;
3613 use datafusion_physical_expr::ScalarFunctionExpr;
3614
3615 Arc::new(ScalarFunctionExpr::new(
3616 "random",
3617 Arc::new(ScalarUDF::from(RandomFunc::new())),
3618 vec![],
3619 Arc::new(Field::new("random", DataType::Float64, false)),
3620 Arc::new(ConfigOptions::default()),
3621 ))
3622 }
3623
3624 fn make_udf_expr(args: Vec<Arc<dyn PhysicalExpr>>) -> Arc<dyn PhysicalExpr> {
3627 use datafusion_common::config::ConfigOptions;
3628 use datafusion_expr::ScalarUDF;
3629 use datafusion_functions::math::abs::AbsFunc;
3630 use datafusion_physical_expr::ScalarFunctionExpr;
3631
3632 Arc::new(ScalarFunctionExpr::new(
3633 "abs",
3634 Arc::new(ScalarUDF::from(AbsFunc::new())),
3635 args,
3636 Arc::new(Field::new("abs", DataType::Int32, false)),
3637 Arc::new(ConfigOptions::default()),
3638 ))
3639 }
3640
3641 fn make_leaf_pushable_expr() -> Arc<dyn PhysicalExpr> {
3645 use datafusion_common::config::ConfigOptions;
3646 use datafusion_expr::ScalarUDF;
3647 use datafusion_functions::core::getfield::GetFieldFunc;
3648 use datafusion_physical_expr::ScalarFunctionExpr;
3649 use datafusion_physical_expr::expressions::Literal;
3650
3651 Arc::new(ScalarFunctionExpr::new(
3652 "get_field",
3653 Arc::new(ScalarUDF::from(GetFieldFunc::new())),
3654 vec![
3655 Arc::new(Column::new("s", 0)),
3656 Arc::new(Literal::new(ScalarValue::Utf8(Some("x".to_string())))),
3657 ],
3658 Arc::new(Field::new("x", DataType::Int32, true)),
3659 Arc::new(ConfigOptions::default()),
3660 ))
3661 }
3662
3663 #[test]
3666 fn test_would_duplicate_allows_column_only_inner() {
3667 let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3668 let col_b: Arc<dyn PhysicalExpr> = Arc::new(Column::new("b", 1));
3669
3670 let inner =
3671 make_projection(vec![(Arc::clone(&col_a), "a"), (Arc::clone(&col_b), "b")]);
3672
3673 let outer = make_projection(vec![
3675 (Arc::new(Column::new("a", 0)), "x"),
3676 (Arc::new(Column::new("a", 0)), "y"),
3677 ]);
3678
3679 assert!(!would_duplicate_costly_exprs(&inner, &outer));
3680 }
3681
3682 #[test]
3685 fn test_would_duplicate_blocks_computed_multi_ref() {
3686 let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3687 let col_b: Arc<dyn PhysicalExpr> = Arc::new(Column::new("b", 1));
3688 let inner = make_projection(vec![
3690 (
3691 Arc::new(BinaryExpr::new(
3692 Arc::clone(&col_a),
3693 Operator::Plus,
3694 Arc::clone(&col_b),
3695 )),
3696 "sum",
3697 ),
3698 (Arc::clone(&col_b), "b"),
3699 ]);
3700
3701 let outer = make_projection(vec![
3703 (Arc::new(Column::new("sum", 0)), "x"),
3704 (Arc::new(Column::new("sum", 0)), "y"),
3705 ]);
3706
3707 assert!(would_duplicate_costly_exprs(&inner, &outer));
3708 }
3709
3710 #[test]
3713 fn test_would_duplicate_allows_unreferenced_volatile() {
3714 let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3715 let inner =
3717 make_projection(vec![(make_volatile_expr(), "r"), (Arc::clone(&col_a), "a")]);
3718
3719 let outer = make_projection(vec![(Arc::new(Column::new("a", 1)), "a")]);
3721
3722 assert!(!would_duplicate_costly_exprs(&inner, &outer));
3723 }
3724
3725 #[test]
3729 fn test_would_duplicate_blocks_multi_ref_volatile() {
3730 let inner = make_projection(vec![(make_volatile_expr(), "r")]);
3732
3733 let outer = make_projection(vec![
3735 (Arc::new(Column::new("r", 0)), "x"),
3736 (Arc::new(Column::new("r", 0)), "y"),
3737 ]);
3738
3739 assert!(would_duplicate_costly_exprs(&inner, &outer));
3740 }
3741
3742 #[test]
3745 fn test_would_duplicate_allows_single_ref_volatile() {
3746 let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3747 let inner =
3749 make_projection(vec![(make_volatile_expr(), "r"), (Arc::clone(&col_a), "a")]);
3750
3751 let outer = make_projection(vec![
3753 (Arc::new(Column::new("r", 0)), "x"),
3754 (Arc::new(Column::new("a", 1)), "a"),
3755 ]);
3756
3757 assert!(!would_duplicate_costly_exprs(&inner, &outer));
3758 }
3759
3760 #[test]
3763 fn test_would_duplicate_blocks_single_expr_self_ref_volatile() {
3764 let inner = make_projection(vec![(make_volatile_expr(), "r")]);
3766
3767 let outer = make_projection(vec![(
3769 Arc::new(BinaryExpr::new(
3770 Arc::new(Column::new("r", 0)),
3771 Operator::Plus,
3772 Arc::new(Column::new("r", 0)),
3773 )),
3774 "x",
3775 )]);
3776
3777 assert!(would_duplicate_costly_exprs(&inner, &outer));
3778 }
3779
3780 #[test]
3783 fn test_would_duplicate_blocks_volatile_nested_in_arithmetic() {
3784 let inner = make_projection(vec![(
3786 Arc::new(BinaryExpr::new(
3787 make_volatile_expr(),
3788 Operator::Plus,
3789 Arc::new(Literal::new(ScalarValue::Float64(Some(1.0)))),
3790 )),
3791 "expr",
3792 )]);
3793
3794 let outer = make_projection(vec![
3796 (Arc::new(Column::new("expr", 0)), "x"),
3797 (Arc::new(Column::new("expr", 0)), "y"),
3798 ]);
3799
3800 assert!(would_duplicate_costly_exprs(&inner, &outer));
3801 }
3802
3803 #[test]
3805 fn test_would_duplicate_empty_projections() {
3806 let inner = make_projection(vec![]);
3807 let outer = make_projection(vec![]);
3808 assert!(!would_duplicate_costly_exprs(&inner, &outer));
3809 }
3810
3811 #[test]
3814 fn test_would_duplicate_blocks_multi_ref_expensive() {
3815 let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3816 let inner = make_projection(vec![(make_udf_expr(vec![col_a]), "abs_a")]);
3818
3819 let outer = make_projection(vec![
3821 (Arc::new(Column::new("abs_a", 0)), "x"),
3822 (Arc::new(Column::new("abs_a", 0)), "y"),
3823 ]);
3824
3825 assert!(would_duplicate_costly_exprs(&inner, &outer));
3826 }
3827
3828 #[test]
3831 fn test_would_duplicate_allows_single_ref_expensive() {
3832 let col_a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
3833 let inner = make_projection(vec![
3835 (make_udf_expr(vec![Arc::clone(&col_a)]), "abs_a"),
3836 (Arc::clone(&col_a), "a"),
3837 ]);
3838
3839 let outer = make_projection(vec![
3841 (Arc::new(Column::new("abs_a", 0)), "out"),
3842 (Arc::new(Column::new("a", 1)), "a"),
3843 ]);
3844
3845 assert!(!would_duplicate_costly_exprs(&inner, &outer));
3846 }
3847
3848 #[test]
3853 fn test_would_duplicate_allows_leaf_pushable_scalar_function() {
3854 let inner = make_projection(vec![(make_leaf_pushable_expr(), "f")]);
3856
3857 let outer = make_projection(vec![
3859 (Arc::new(Column::new("f", 0)), "x"),
3860 (Arc::new(Column::new("f", 0)), "y"),
3861 ]);
3862
3863 assert!(!would_duplicate_costly_exprs(&inner, &outer));
3864 }
3865}