VTX-4075: fix more tests

coralogix · Jan 26, 2024 · 7608327 · 7608327
1 parent 0220a95
commit 7608327
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs
@@ -3718,7 +3718,7 @@ pub(crate) mod tests {
         let expected = &[
             "SortRequiredExec: [c@2 ASC]",
             // Since this projection is trivial, increasing parallelism is not beneficial
-            "ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]",
+            "ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c], sort_expr=[c@2 ASC]",
             "ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
         ];
 
@@ -4250,14 +4250,14 @@ pub(crate) mod tests {
             sort_preserving_merge_exec(sort_key_after_projection, proj_parquet);
         let expected = &[
             "SortPreservingMergeExec: [c2@1 ASC]",
-            "  ProjectionExec: expr=[a@0 as a2, c@2 as c2]",
+            "  ProjectionExec: expr=[a@0 as a2, c@2 as c2], sort_expr=[c2@1 ASC]",
             "    ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
         ];
         plans_matches_expected!(expected, &plan_parquet);
 
         // data should not be repartitioned / resorted
         let expected_parquet = &[
-            "ProjectionExec: expr=[a@0 as a2, c@2 as c2]",
+            "ProjectionExec: expr=[a@0 as a2, c@2 as c2], sort_expr=[c2@1 ASC]",
             "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
         ];
 
@@ -4289,14 +4289,14 @@ pub(crate) mod tests {
         let plan_csv = sort_preserving_merge_exec(sort_key_after_projection, proj_csv);
         let expected = &[
             "SortPreservingMergeExec: [c2@1 ASC]",
-            "  ProjectionExec: expr=[a@0 as a2, c@2 as c2]",
+            "  ProjectionExec: expr=[a@0 as a2, c@2 as c2], sort_expr=[c2@1 ASC]",
             "    CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], has_header=false",
         ];
         plans_matches_expected!(expected, &plan_csv);
 
         // data should not be repartitioned / resorted
         let expected_csv = &[
-            "ProjectionExec: expr=[a@0 as a2, c@2 as c2]",
+            "ProjectionExec: expr=[a@0 as a2, c@2 as c2], sort_expr=[c2@1 ASC]",
             "CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], has_header=false",
         ];
 

diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs
@@ -2162,7 +2162,7 @@ mod tests {
         ];
         let expected_optimized = [
             "SortPreservingMergeExec: [a@0 ASC]",
-            "  SortPreservingRepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, sort_exprs=a@0 ASC",
+            "  RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC",
             "    RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
             "      CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC], has_header=false",
         ];
@@ -2191,7 +2191,7 @@ mod tests {
         ];
         let expected_optimized = [
             "SortPreservingMergeExec: [a@0 ASC]",
-            "  SortPreservingRepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, sort_exprs=a@0 ASC",
+            "  RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC",
             "    RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
             "      CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC], has_header=false",
         ];
@@ -2263,7 +2263,7 @@ mod tests {
         let expected_input = [
             "BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow }], mode=[Sorted]",
             "  SortPreservingMergeExec: [a@0 ASC,b@1 ASC]",
-            "    SortPreservingRepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=10, sort_exprs=a@0 ASC,b@1 ASC",
+            "    RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC,b@1 ASC",
             "      RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
             "        SortExec: expr=[a@0 ASC,b@1 ASC]",
             "          CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], has_header=false",

diff --git a/datafusion/core/src/physical_optimizer/projection_pushdown.rs b/datafusion/core/src/physical_optimizer/projection_pushdown.rs
@@ -2200,7 +2200,7 @@ mod tests {
 
         let initial = get_plan_string(&projection);
         let expected_initial = [
-            "ProjectionExec: expr=[c@2 as c, a@0 as new_a, b@1 as b]",
+            "ProjectionExec: expr=[c@2 as c, a@0 as new_a, b@1 as b], sort_expr=[b@2 ASC,c@0 + new_a@1 ASC]",
             "  SortExec: expr=[b@1 ASC,c@2 + a@0 ASC]",
             "    CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], has_header=false"
             ];