You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2023/01/17 01:03:52 UTC
[GitHub] [arrow-datafusion] alamb opened a new pull request, #4928: Reduce redundancy in sort_enforcement tests
alamb opened a new pull request, #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928
# Which issue does this PR close?
# Rationale for this change
Rationale: I am working on a bug related to sort enforcement and wanted to add a new test and the current structure was very hard to do without a lot of copy / paste.
# What changes are included in this PR?
refactor repetition in sort_enforcement.rs tests into a macro
# Are these changes tested?
Only tests
# Are there any user-facing changes?
No -- there is no intended change in behavior
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] mustafasrepo commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
mustafasrepo commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1071822908
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -664,32 +677,33 @@ mod tests {
Arc::new(WindowFrame::new(true)),
schema.as_ref(),
)?],
- sort_exec.clone(),
- sort_exec.schema(),
+ sort.clone(),
+ sort.schema(),
vec![],
Some(sort_exprs),
)?) as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", window_agg_exec.schema().as_ref()).unwrap(),
- options: SortOptions {
+
+ let sort_exprs = vec![sort_expr_options(
+ "non_nullable_col",
+ &window_agg_exec.schema(),
+ SortOptions {
descending: false,
nulls_first: false,
},
- }];
- let sort_exec = Arc::new(SortExec::try_new(
- sort_exprs.clone(),
- window_agg_exec,
- None,
- )?) as Arc<dyn ExecutionPlan>;
+ )];
+
+ let sort = sort_exec(sort_exprs.clone(), window_agg_exec);
+
// Add dummy layer propagating Sort above, to test whether sort can be removed from multi layer before
- let filter_exec = Arc::new(FilterExec::try_new(
+ let filter = filter_exec(
Arc::new(NotExpr::new(
col("non_nullable_col", schema.as_ref()).unwrap(),
)),
- sort_exec,
- )?) as Arc<dyn ExecutionPlan>;
+ sort,
+ );
+
// let filter_exec = sort_exec;
- let window_agg_exec = Arc::new(WindowAggExec::try_new(
+ let physical_plan = Arc::new(WindowAggExec::try_new(
Review Comment:
given util function `window_exec` is available we can construct `physical_plan` with
``` rust
let physical_plan = window_exec(
filter.clone(),
filter.schema(),
&sort_exprs,
"non_nullable_col",
)?;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] mustafasrepo commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
mustafasrepo commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1071835456
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -699,214 +713,147 @@ mod tests {
Arc::new(WindowFrame::new(true)),
schema.as_ref(),
)?],
- filter_exec.clone(),
- filter_exec.schema(),
+ filter.clone(),
+ filter.schema(),
vec![],
Some(sort_exprs),
)?) as Arc<dyn ExecutionPlan>;
- let physical_plan = window_agg_exec;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "WindowAggExec: 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 }]",
- " FilterExec: NOT non_nullable_col@1",
- " SortExec: [non_nullable_col@1 ASC NULLS LAST]",
- " WindowAggExec: 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 }]",
- " SortExec: [non_nullable_col@1 DESC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "WindowAggExec: 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: CurrentRow, end_bound: Following(NULL) }]",
- " FilterExec: NOT non_nullable_col@1",
- " WindowAggExec: 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 }]",
- " SortExec: [non_nullable_col@1 DESC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+
+ let expected_input = vec![
+ "WindowAggExec: 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 }]",
+ " FilterExec: NOT non_nullable_col@1",
+ " SortExec: [non_nullable_col@1 ASC NULLS LAST]",
+ " WindowAggExec: 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 }]",
+ " SortExec: [non_nullable_col@1 DESC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+
+ let expected_optimized = vec![
+ "WindowAggExec: 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: CurrentRow, end_bound: Following(NULL) }]",
+ " FilterExec: NOT non_nullable_col@1",
+ " WindowAggExec: 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 }]",
+ " SortExec: [non_nullable_col@1 DESC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_add_required_sort() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let physical_plan = Arc::new(SortPreservingMergeExec::new(sort_exprs, source))
- as Arc<dyn ExecutionPlan>;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = { vec!["SortPreservingMergeExec: [nullable_col@0 ASC]"] };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
+
+ let sort_exprs = vec![sort_expr("nullable_col", &schema)];
+
+ let physical_plan = sort_preserving_merge_exec(sort_exprs, source);
+
+ let expected_input = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_remove_unnecessary_sort1() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs.clone(), source, None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_preserving_merge_exec =
- Arc::new(SortPreservingMergeExec::new(sort_exprs, sort_exec))
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(
- sort_exprs.clone(),
- sort_preserving_merge_exec,
- None,
- )?) as Arc<dyn ExecutionPlan>;
- let sort_preserving_merge_exec =
- Arc::new(SortPreservingMergeExec::new(sort_exprs, sort_exec))
- as Arc<dyn ExecutionPlan>;
- let physical_plan = sort_preserving_merge_exec;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
+ let sort_exprs = vec![sort_expr("nullable_col", &schema)];
+ let sort = sort_exec(sort_exprs.clone(), source);
+ let spm = sort_preserving_merge_exec(sort_exprs, sort);
+
+ let sort_exprs = vec![sort_expr("nullable_col", &schema)];
+ let sort = sort_exec(sort_exprs.clone(), spm);
+ let physical_plan = sort_preserving_merge_exec(sort_exprs, sort);
+ let expected_input = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_change_wrong_sorting() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
+ let source = memory_exec(&schema);
let sort_exprs = vec![
- PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- },
- PhysicalSortExpr {
- expr: col("non_nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- },
+ sort_expr("nullable_col", &schema),
+ sort_expr("non_nullable_col", &schema),
];
- let sort_exec = Arc::new(SortExec::try_new(
- vec![sort_exprs[0].clone()],
- source,
- None,
- )?) as Arc<dyn ExecutionPlan>;
- let sort_preserving_merge_exec =
- Arc::new(SortPreservingMergeExec::new(sort_exprs, sort_exec))
- as Arc<dyn ExecutionPlan>;
- let physical_plan = sort_preserving_merge_exec;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
- " SortExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let sort = sort_exec(vec![sort_exprs[0].clone()], source);
+ let physical_plan = sort_preserving_merge_exec(sort_exprs, sort);
+ let expected_input = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
+ " SortExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
+
+ /// make PhysicalSortExpr with default options
+ fn sort_expr(name: &str, schema: &Schema) -> PhysicalSortExpr {
+ sort_expr_options(name, schema, SortOptions::default())
+ }
+
+ /// PhysicalSortExpr with specified options
+ fn sort_expr_options(
+ name: &str,
+ schema: &Schema,
+ options: SortOptions,
+ ) -> PhysicalSortExpr {
+ PhysicalSortExpr {
+ expr: col(name, schema).unwrap(),
+ options,
+ }
+ }
+
+ fn memory_exec(schema: &SchemaRef) -> Arc<dyn ExecutionPlan> {
+ Arc::new(MemoryExec::try_new(&[], schema.clone(), None).unwrap())
+ }
+
+ fn sort_exec(
+ sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
+ input: Arc<dyn ExecutionPlan>,
+ ) -> Arc<dyn ExecutionPlan> {
+ let sort_exprs = sort_exprs.into_iter().collect();
+ Arc::new(SortExec::try_new(sort_exprs, input, None).unwrap())
+ }
+
+ fn sort_preserving_merge_exec(
+ sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
+ input: Arc<dyn ExecutionPlan>,
+ ) -> Arc<dyn ExecutionPlan> {
+ let sort_exprs = sort_exprs.into_iter().collect();
+ Arc::new(SortPreservingMergeExec::new(sort_exprs, input))
+ }
+
Review Comment:
By the way this is just a suggestion. I think, with or without this change PR is ready to merge
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] ursabot commented on pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
ursabot commented on PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#issuecomment-1385345781
Benchmark runs are scheduled for baseline = 4e08117b1d3a2d17a2f416ac73120a1a3edac1da and contender = e7c2ef09242711f31f50349e524381d004c26938. e7c2ef09242711f31f50349e524381d004c26938 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Skipped :warning: Benchmarking of arrow-datafusion-commits is not supported on ec2-t3-xlarge-us-east-2] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/daba6a80243a4b4e90ac9c92290cfdaa...25f6c80779884619b1fc00816a48463a/)
[Skipped :warning: Benchmarking of arrow-datafusion-commits is not supported on test-mac-arm] [test-mac-arm](https://conbench.ursa.dev/compare/runs/1f797c8937274629b58a3f75e499e29c...fb1be5d878894adeb2f6548c3ec0422e/)
[Skipped :warning: Benchmarking of arrow-datafusion-commits is not supported on ursa-i9-9960x] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/45a62c6914464d139e51e2155a4403b6...0e41736e647840c3ab98ca678dee04d5/)
[Skipped :warning: Benchmarking of arrow-datafusion-commits is not supported on ursa-thinkcentre-m75q] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/2a0df24c1d82417da738e4dce44d84b9...d72ff241e3b04233bf0718ed2ee4aac9/)
Buildkite builds:
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] mustafasrepo commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
mustafasrepo commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1071822908
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -664,32 +677,33 @@ mod tests {
Arc::new(WindowFrame::new(true)),
schema.as_ref(),
)?],
- sort_exec.clone(),
- sort_exec.schema(),
+ sort.clone(),
+ sort.schema(),
vec![],
Some(sort_exprs),
)?) as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", window_agg_exec.schema().as_ref()).unwrap(),
- options: SortOptions {
+
+ let sort_exprs = vec![sort_expr_options(
+ "non_nullable_col",
+ &window_agg_exec.schema(),
+ SortOptions {
descending: false,
nulls_first: false,
},
- }];
- let sort_exec = Arc::new(SortExec::try_new(
- sort_exprs.clone(),
- window_agg_exec,
- None,
- )?) as Arc<dyn ExecutionPlan>;
+ )];
+
+ let sort = sort_exec(sort_exprs.clone(), window_agg_exec);
+
// Add dummy layer propagating Sort above, to test whether sort can be removed from multi layer before
- let filter_exec = Arc::new(FilterExec::try_new(
+ let filter = filter_exec(
Arc::new(NotExpr::new(
col("non_nullable_col", schema.as_ref()).unwrap(),
)),
- sort_exec,
- )?) as Arc<dyn ExecutionPlan>;
+ sort,
+ );
+
// let filter_exec = sort_exec;
- let window_agg_exec = Arc::new(WindowAggExec::try_new(
+ let physical_plan = Arc::new(WindowAggExec::try_new(
Review Comment:
given util function `window_exec` is available we can construct physical plan with
``` rust
let physical_plan = window_exec(
filter.clone(),
filter.schema(),
&sort_exprs,
"non_nullable_col",
)?;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] mustafasrepo commented on pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
mustafasrepo commented on PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#issuecomment-1384914613
This is much cleaner and more readable. Thanks @alamb.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] mustafasrepo commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
mustafasrepo commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1071822908
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -664,32 +677,33 @@ mod tests {
Arc::new(WindowFrame::new(true)),
schema.as_ref(),
)?],
- sort_exec.clone(),
- sort_exec.schema(),
+ sort.clone(),
+ sort.schema(),
vec![],
Some(sort_exprs),
)?) as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", window_agg_exec.schema().as_ref()).unwrap(),
- options: SortOptions {
+
+ let sort_exprs = vec![sort_expr_options(
+ "non_nullable_col",
+ &window_agg_exec.schema(),
+ SortOptions {
descending: false,
nulls_first: false,
},
- }];
- let sort_exec = Arc::new(SortExec::try_new(
- sort_exprs.clone(),
- window_agg_exec,
- None,
- )?) as Arc<dyn ExecutionPlan>;
+ )];
+
+ let sort = sort_exec(sort_exprs.clone(), window_agg_exec);
+
// Add dummy layer propagating Sort above, to test whether sort can be removed from multi layer before
- let filter_exec = Arc::new(FilterExec::try_new(
+ let filter = filter_exec(
Arc::new(NotExpr::new(
col("non_nullable_col", schema.as_ref()).unwrap(),
)),
- sort_exec,
- )?) as Arc<dyn ExecutionPlan>;
+ sort,
+ );
+
// let filter_exec = sort_exec;
- let window_agg_exec = Arc::new(WindowAggExec::try_new(
+ let physical_plan = Arc::new(WindowAggExec::try_new(
Review Comment:
given util function below is available we can construct physical plan with
``` rust
let physical_plan = window_exec(
filter.clone(),
filter.schema(),
&sort_exprs,
"non_nullable_col",
)?;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
alamb commented on PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#issuecomment-1384712409
cc @mustafasrepo and @mingmwang who I think have worked on this code
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
alamb commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1072135633
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -589,71 +589,84 @@ mod tests {
Ok(())
}
+ /// Runs the sort enforcement optimizer and asserts the plan
+ /// against the original and expected plans
+ ///
+ /// `$EXPECTED_PLAN_LINES`: input plan
+ /// `$EXPECTED_OPTIMIZED_PLAN_LINES`: optimized plan
+ /// `$PLAN`: the plan to optimized
+ ///
+ macro_rules! assert_optimized {
+ ($EXPECTED_PLAN_LINES: expr, $EXPECTED_OPTIMIZED_PLAN_LINES: expr, $PLAN: expr) => {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+
+ let physical_plan = $PLAN;
+ let formatted = displayable(physical_plan.as_ref()).indent().to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+
+ let expected_plan_lines: Vec<&str> = $EXPECTED_PLAN_LINES
+ .iter().map(|s| *s).collect();
+
+ assert_eq!(
+ expected_plan_lines, actual,
+ "\n**Original Plan Mismatch\n\nexpected:\n\n{expected_plan_lines:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ let expected_optimized_lines: Vec<&str> = $EXPECTED_OPTIMIZED_PLAN_LINES
+ .iter().map(|s| *s).collect();
+
+ // Run the actual optimizer
+ let optimized_physical_plan =
+ EnforceSorting::new().optimize(physical_plan, state.config_options())?;
+
+ let formatted = displayable(optimized_physical_plan.as_ref())
+ .indent()
+ .to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+ assert_eq!(
+ expected_optimized_lines, actual,
+ "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected_optimized_lines:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ };
+ }
+
#[tokio::test]
async fn test_remove_unnecessary_sort() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs, source, None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let physical_plan = Arc::new(SortExec::try_new(sort_exprs, sort_exec, None)?)
- as Arc<dyn ExecutionPlan>;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortExec: [nullable_col@0 ASC]",
- " SortExec: [non_nullable_col@1 ASC]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = { vec!["SortExec: [nullable_col@0 ASC]"] };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
+ let input = sort_exec(vec![sort_expr("non_nullable_col", &schema)], source);
+ let physical_plan = sort_exec(vec![sort_expr("nullable_col", &schema)], input);
+
+ let expected_input = vec![
+ "SortExec: [nullable_col@0 ASC]",
+ " SortExec: [non_nullable_col@1 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_remove_unnecessary_sort_window_multilayer() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", source.schema().as_ref()).unwrap(),
- options: SortOptions {
+ let source = memory_exec(&schema);
+
+ let sort_exprs = vec![sort_expr_options(
+ "non_nullable_col",
+ &source.schema(),
+ SortOptions {
descending: true,
nulls_first: true,
},
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs.clone(), source, None)?)
- as Arc<dyn ExecutionPlan>;
+ )];
+ let sort = sort_exec(sort_exprs.clone(), source);
+
let window_agg_exec = Arc::new(WindowAggExec::try_new(
Review Comment:
in https://github.com/apache/arrow-datafusion/pull/4945
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb merged pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
alamb merged PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
alamb commented on PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#issuecomment-1385323134
Update https://github.com/apache/arrow-datafusion/issues/4943 is the bug I am working on with sort enforcement
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] mingmwang commented on pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
mingmwang commented on PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#issuecomment-1384930514
Nice !!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
alamb commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1072131145
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -589,71 +589,84 @@ mod tests {
Ok(())
}
+ /// Runs the sort enforcement optimizer and asserts the plan
+ /// against the original and expected plans
+ ///
+ /// `$EXPECTED_PLAN_LINES`: input plan
+ /// `$EXPECTED_OPTIMIZED_PLAN_LINES`: optimized plan
+ /// `$PLAN`: the plan to optimized
+ ///
+ macro_rules! assert_optimized {
+ ($EXPECTED_PLAN_LINES: expr, $EXPECTED_OPTIMIZED_PLAN_LINES: expr, $PLAN: expr) => {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+
+ let physical_plan = $PLAN;
+ let formatted = displayable(physical_plan.as_ref()).indent().to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+
+ let expected_plan_lines: Vec<&str> = $EXPECTED_PLAN_LINES
+ .iter().map(|s| *s).collect();
+
+ assert_eq!(
+ expected_plan_lines, actual,
+ "\n**Original Plan Mismatch\n\nexpected:\n\n{expected_plan_lines:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ let expected_optimized_lines: Vec<&str> = $EXPECTED_OPTIMIZED_PLAN_LINES
+ .iter().map(|s| *s).collect();
+
+ // Run the actual optimizer
+ let optimized_physical_plan =
+ EnforceSorting::new().optimize(physical_plan, state.config_options())?;
+
+ let formatted = displayable(optimized_physical_plan.as_ref())
+ .indent()
+ .to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+ assert_eq!(
+ expected_optimized_lines, actual,
+ "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected_optimized_lines:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ };
+ }
+
#[tokio::test]
async fn test_remove_unnecessary_sort() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs, source, None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let physical_plan = Arc::new(SortExec::try_new(sort_exprs, sort_exec, None)?)
- as Arc<dyn ExecutionPlan>;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortExec: [nullable_col@0 ASC]",
- " SortExec: [non_nullable_col@1 ASC]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = { vec!["SortExec: [nullable_col@0 ASC]"] };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
+ let input = sort_exec(vec![sort_expr("non_nullable_col", &schema)], source);
+ let physical_plan = sort_exec(vec![sort_expr("nullable_col", &schema)], input);
+
+ let expected_input = vec![
+ "SortExec: [nullable_col@0 ASC]",
+ " SortExec: [non_nullable_col@1 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_remove_unnecessary_sort_window_multilayer() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", source.schema().as_ref()).unwrap(),
- options: SortOptions {
+ let source = memory_exec(&schema);
+
+ let sort_exprs = vec![sort_expr_options(
+ "non_nullable_col",
+ &source.schema(),
+ SortOptions {
descending: true,
nulls_first: true,
},
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs.clone(), source, None)?)
- as Arc<dyn ExecutionPlan>;
+ )];
+ let sort = sort_exec(sort_exprs.clone(), source);
+
let window_agg_exec = Arc::new(WindowAggExec::try_new(
Review Comment:
Thank you -- I was being lazy -- I will do so
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
alamb commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1071656208
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -589,71 +589,84 @@ mod tests {
Ok(())
}
+ /// Runs the sort enforcement optimizer and asserts the plan
+ /// against the original and expected plans
+ ///
+ /// `$EXPECTED_PLAN_LINES`: input plan
+ /// `$EXPECTED_OPTIMIZED_PLAN_LINES`: optimized plan
+ /// `$PLAN`: the plan to optimized
+ ///
+ macro_rules! assert_optimized {
+ ($EXPECTED_PLAN_LINES: expr, $EXPECTED_OPTIMIZED_PLAN_LINES: expr, $PLAN: expr) => {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+
+ let physical_plan = $PLAN;
+ let formatted = displayable(physical_plan.as_ref()).indent().to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+
+ let expected_plan_lines: Vec<&str> = $EXPECTED_PLAN_LINES
+ .iter().map(|s| *s).collect();
+
+ assert_eq!(
+ expected_plan_lines, actual,
+ "\n**Original Plan Mismatch\n\nexpected:\n\n{expected_plan_lines:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ let expected_optimized_lines: Vec<&str> = $EXPECTED_OPTIMIZED_PLAN_LINES
+ .iter().map(|s| *s).collect();
+
+ // Run the actual optimizer
+ let optimized_physical_plan =
+ EnforceSorting::new().optimize(physical_plan, state.config_options())?;
+
+ let formatted = displayable(optimized_physical_plan.as_ref())
+ .indent()
+ .to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+ assert_eq!(
+ expected_optimized_lines, actual,
+ "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected_optimized_lines:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ };
+ }
+
#[tokio::test]
async fn test_remove_unnecessary_sort() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs, source, None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let physical_plan = Arc::new(SortExec::try_new(sort_exprs, sort_exec, None)?)
- as Arc<dyn ExecutionPlan>;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortExec: [nullable_col@0 ASC]",
- " SortExec: [non_nullable_col@1 ASC]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = { vec!["SortExec: [nullable_col@0 ASC]"] };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
Review Comment:
I think the new structure is much clearer about what the tests are doing -- if you look at the whitespace blind diff https://github.com/apache/arrow-datafusion/pull/4928/files?w=1 you can see that all the plans (original and optimized) are the same
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] mustafasrepo commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
mustafasrepo commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1071819266
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -699,214 +713,147 @@ mod tests {
Arc::new(WindowFrame::new(true)),
schema.as_ref(),
)?],
- filter_exec.clone(),
- filter_exec.schema(),
+ filter.clone(),
+ filter.schema(),
vec![],
Some(sort_exprs),
)?) as Arc<dyn ExecutionPlan>;
- let physical_plan = window_agg_exec;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "WindowAggExec: 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 }]",
- " FilterExec: NOT non_nullable_col@1",
- " SortExec: [non_nullable_col@1 ASC NULLS LAST]",
- " WindowAggExec: 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 }]",
- " SortExec: [non_nullable_col@1 DESC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "WindowAggExec: 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: CurrentRow, end_bound: Following(NULL) }]",
- " FilterExec: NOT non_nullable_col@1",
- " WindowAggExec: 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 }]",
- " SortExec: [non_nullable_col@1 DESC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+
+ let expected_input = vec![
+ "WindowAggExec: 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 }]",
+ " FilterExec: NOT non_nullable_col@1",
+ " SortExec: [non_nullable_col@1 ASC NULLS LAST]",
+ " WindowAggExec: 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 }]",
+ " SortExec: [non_nullable_col@1 DESC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+
+ let expected_optimized = vec![
+ "WindowAggExec: 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: CurrentRow, end_bound: Following(NULL) }]",
+ " FilterExec: NOT non_nullable_col@1",
+ " WindowAggExec: 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 }]",
+ " SortExec: [non_nullable_col@1 DESC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_add_required_sort() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let physical_plan = Arc::new(SortPreservingMergeExec::new(sort_exprs, source))
- as Arc<dyn ExecutionPlan>;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = { vec!["SortPreservingMergeExec: [nullable_col@0 ASC]"] };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
+
+ let sort_exprs = vec![sort_expr("nullable_col", &schema)];
+
+ let physical_plan = sort_preserving_merge_exec(sort_exprs, source);
+
+ let expected_input = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_remove_unnecessary_sort1() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs.clone(), source, None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_preserving_merge_exec =
- Arc::new(SortPreservingMergeExec::new(sort_exprs, sort_exec))
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(
- sort_exprs.clone(),
- sort_preserving_merge_exec,
- None,
- )?) as Arc<dyn ExecutionPlan>;
- let sort_preserving_merge_exec =
- Arc::new(SortPreservingMergeExec::new(sort_exprs, sort_exec))
- as Arc<dyn ExecutionPlan>;
- let physical_plan = sort_preserving_merge_exec;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
+ let sort_exprs = vec![sort_expr("nullable_col", &schema)];
+ let sort = sort_exec(sort_exprs.clone(), source);
+ let spm = sort_preserving_merge_exec(sort_exprs, sort);
+
+ let sort_exprs = vec![sort_expr("nullable_col", &schema)];
+ let sort = sort_exec(sort_exprs.clone(), spm);
+ let physical_plan = sort_preserving_merge_exec(sort_exprs, sort);
+ let expected_input = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_change_wrong_sorting() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
+ let source = memory_exec(&schema);
let sort_exprs = vec![
- PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- },
- PhysicalSortExpr {
- expr: col("non_nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- },
+ sort_expr("nullable_col", &schema),
+ sort_expr("non_nullable_col", &schema),
];
- let sort_exec = Arc::new(SortExec::try_new(
- vec![sort_exprs[0].clone()],
- source,
- None,
- )?) as Arc<dyn ExecutionPlan>;
- let sort_preserving_merge_exec =
- Arc::new(SortPreservingMergeExec::new(sort_exprs, sort_exec))
- as Arc<dyn ExecutionPlan>;
- let physical_plan = sort_preserving_merge_exec;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
- " SortExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let sort = sort_exec(vec![sort_exprs[0].clone()], source);
+ let physical_plan = sort_preserving_merge_exec(sort_exprs, sort);
+ let expected_input = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
+ " SortExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
+
+ /// make PhysicalSortExpr with default options
+ fn sort_expr(name: &str, schema: &Schema) -> PhysicalSortExpr {
+ sort_expr_options(name, schema, SortOptions::default())
+ }
+
+ /// PhysicalSortExpr with specified options
+ fn sort_expr_options(
+ name: &str,
+ schema: &Schema,
+ options: SortOptions,
+ ) -> PhysicalSortExpr {
+ PhysicalSortExpr {
+ expr: col(name, schema).unwrap(),
+ options,
+ }
+ }
+
+ fn memory_exec(schema: &SchemaRef) -> Arc<dyn ExecutionPlan> {
+ Arc::new(MemoryExec::try_new(&[], schema.clone(), None).unwrap())
+ }
+
+ fn sort_exec(
+ sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
+ input: Arc<dyn ExecutionPlan>,
+ ) -> Arc<dyn ExecutionPlan> {
+ let sort_exprs = sort_exprs.into_iter().collect();
+ Arc::new(SortExec::try_new(sort_exprs, input, None).unwrap())
+ }
+
+ fn sort_preserving_merge_exec(
+ sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
+ input: Arc<dyn ExecutionPlan>,
+ ) -> Arc<dyn ExecutionPlan> {
+ let sort_exprs = sort_exprs.into_iter().collect();
+ Arc::new(SortPreservingMergeExec::new(sort_exprs, input))
+ }
+
Review Comment:
Maybe we can add here one more function to encapsulate window exec creation
```rust
fn window_exec(
input: Arc<dyn ExecutionPlan>,
schema: SchemaRef,
sort_exprs: &[PhysicalSortExpr],
arg_column_name: &str,
) -> Result<Arc<dyn ExecutionPlan>> {
Ok(Arc::new(WindowAggExec::try_new(
vec![create_window_expr(
&WindowFunction::AggregateFunction(AggregateFunction::Count),
"count".to_owned(),
&[col(arg_column_name, &schema)?],
&[],
sort_exprs,
Arc::new(WindowFrame::new(true)),
schema.as_ref(),
)?],
input,
schema,
vec![],
Some(sort_exprs.to_vec()),
)?) as Arc<dyn ExecutionPlan>)
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] mustafasrepo commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
mustafasrepo commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1071819266
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -699,214 +713,147 @@ mod tests {
Arc::new(WindowFrame::new(true)),
schema.as_ref(),
)?],
- filter_exec.clone(),
- filter_exec.schema(),
+ filter.clone(),
+ filter.schema(),
vec![],
Some(sort_exprs),
)?) as Arc<dyn ExecutionPlan>;
- let physical_plan = window_agg_exec;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "WindowAggExec: 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 }]",
- " FilterExec: NOT non_nullable_col@1",
- " SortExec: [non_nullable_col@1 ASC NULLS LAST]",
- " WindowAggExec: 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 }]",
- " SortExec: [non_nullable_col@1 DESC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "WindowAggExec: 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: CurrentRow, end_bound: Following(NULL) }]",
- " FilterExec: NOT non_nullable_col@1",
- " WindowAggExec: 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 }]",
- " SortExec: [non_nullable_col@1 DESC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+
+ let expected_input = vec![
+ "WindowAggExec: 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 }]",
+ " FilterExec: NOT non_nullable_col@1",
+ " SortExec: [non_nullable_col@1 ASC NULLS LAST]",
+ " WindowAggExec: 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 }]",
+ " SortExec: [non_nullable_col@1 DESC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+
+ let expected_optimized = vec![
+ "WindowAggExec: 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: CurrentRow, end_bound: Following(NULL) }]",
+ " FilterExec: NOT non_nullable_col@1",
+ " WindowAggExec: 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 }]",
+ " SortExec: [non_nullable_col@1 DESC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_add_required_sort() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let physical_plan = Arc::new(SortPreservingMergeExec::new(sort_exprs, source))
- as Arc<dyn ExecutionPlan>;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = { vec!["SortPreservingMergeExec: [nullable_col@0 ASC]"] };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
+
+ let sort_exprs = vec![sort_expr("nullable_col", &schema)];
+
+ let physical_plan = sort_preserving_merge_exec(sort_exprs, source);
+
+ let expected_input = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_remove_unnecessary_sort1() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs.clone(), source, None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_preserving_merge_exec =
- Arc::new(SortPreservingMergeExec::new(sort_exprs, sort_exec))
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(
- sort_exprs.clone(),
- sort_preserving_merge_exec,
- None,
- )?) as Arc<dyn ExecutionPlan>;
- let sort_preserving_merge_exec =
- Arc::new(SortPreservingMergeExec::new(sort_exprs, sort_exec))
- as Arc<dyn ExecutionPlan>;
- let physical_plan = sort_preserving_merge_exec;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortPreservingMergeExec: [nullable_col@0 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
+ let sort_exprs = vec![sort_expr("nullable_col", &schema)];
+ let sort = sort_exec(sort_exprs.clone(), source);
+ let spm = sort_preserving_merge_exec(sort_exprs, sort);
+
+ let sort_exprs = vec![sort_expr("nullable_col", &schema)];
+ let sort = sort_exec(sort_exprs.clone(), spm);
+ let physical_plan = sort_preserving_merge_exec(sort_exprs, sort);
+ let expected_input = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortPreservingMergeExec: [nullable_col@0 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_change_wrong_sorting() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
+ let source = memory_exec(&schema);
let sort_exprs = vec![
- PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- },
- PhysicalSortExpr {
- expr: col("non_nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- },
+ sort_expr("nullable_col", &schema),
+ sort_expr("non_nullable_col", &schema),
];
- let sort_exec = Arc::new(SortExec::try_new(
- vec![sort_exprs[0].clone()],
- source,
- None,
- )?) as Arc<dyn ExecutionPlan>;
- let sort_preserving_merge_exec =
- Arc::new(SortPreservingMergeExec::new(sort_exprs, sort_exec))
- as Arc<dyn ExecutionPlan>;
- let physical_plan = sort_preserving_merge_exec;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
- " SortExec: [nullable_col@0 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = {
- vec![
- "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
- " SortExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
- " MemoryExec: partitions=0, partition_sizes=[]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- assert_eq!(
- expected, actual,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let sort = sort_exec(vec![sort_exprs[0].clone()], source);
+ let physical_plan = sort_preserving_merge_exec(sort_exprs, sort);
+ let expected_input = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
+ " SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortPreservingMergeExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
+ " SortExec: [nullable_col@0 ASC,non_nullable_col@1 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
+
+ /// make PhysicalSortExpr with default options
+ fn sort_expr(name: &str, schema: &Schema) -> PhysicalSortExpr {
+ sort_expr_options(name, schema, SortOptions::default())
+ }
+
+ /// PhysicalSortExpr with specified options
+ fn sort_expr_options(
+ name: &str,
+ schema: &Schema,
+ options: SortOptions,
+ ) -> PhysicalSortExpr {
+ PhysicalSortExpr {
+ expr: col(name, schema).unwrap(),
+ options,
+ }
+ }
+
+ fn memory_exec(schema: &SchemaRef) -> Arc<dyn ExecutionPlan> {
+ Arc::new(MemoryExec::try_new(&[], schema.clone(), None).unwrap())
+ }
+
+ fn sort_exec(
+ sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
+ input: Arc<dyn ExecutionPlan>,
+ ) -> Arc<dyn ExecutionPlan> {
+ let sort_exprs = sort_exprs.into_iter().collect();
+ Arc::new(SortExec::try_new(sort_exprs, input, None).unwrap())
+ }
+
+ fn sort_preserving_merge_exec(
+ sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
+ input: Arc<dyn ExecutionPlan>,
+ ) -> Arc<dyn ExecutionPlan> {
+ let sort_exprs = sort_exprs.into_iter().collect();
+ Arc::new(SortPreservingMergeExec::new(sort_exprs, input))
+ }
+
Review Comment:
Maybe we can add here one more function to encapsulate window exec creation
```rust
fn window_exec(
input: Arc<dyn ExecutionPlan>,
schema: SchemaRef,
sort_exprs: &[PhysicalSortExpr],
fn_arg_column_name: &str,
) -> Result<Arc<dyn ExecutionPlan>> {
let fn_arg = col(fn_arg_column_name, &schema)?;
Ok(Arc::new(WindowAggExec::try_new(
vec![create_window_expr(
&WindowFunction::AggregateFunction(AggregateFunction::Count),
"count".to_owned(),
&[fn_arg],
&[],
sort_exprs,
Arc::new(WindowFrame::new(true)),
schema.as_ref(),
)?],
input,
schema,
vec![],
Some(sort_exprs.to_vec()),
)?) as Arc<dyn ExecutionPlan>)
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
alamb commented on PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#issuecomment-1385335264
Will implement suggestions in a follow on PR
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] mustafasrepo commented on a diff in pull request #4928: Reduce redundancy in sort_enforcement tests
Posted by GitBox <gi...@apache.org>.
mustafasrepo commented on code in PR #4928:
URL: https://github.com/apache/arrow-datafusion/pull/4928#discussion_r1071825011
##########
datafusion/core/src/physical_optimizer/sort_enforcement.rs:
##########
@@ -589,71 +589,84 @@ mod tests {
Ok(())
}
+ /// Runs the sort enforcement optimizer and asserts the plan
+ /// against the original and expected plans
+ ///
+ /// `$EXPECTED_PLAN_LINES`: input plan
+ /// `$EXPECTED_OPTIMIZED_PLAN_LINES`: optimized plan
+ /// `$PLAN`: the plan to optimized
+ ///
+ macro_rules! assert_optimized {
+ ($EXPECTED_PLAN_LINES: expr, $EXPECTED_OPTIMIZED_PLAN_LINES: expr, $PLAN: expr) => {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+
+ let physical_plan = $PLAN;
+ let formatted = displayable(physical_plan.as_ref()).indent().to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+
+ let expected_plan_lines: Vec<&str> = $EXPECTED_PLAN_LINES
+ .iter().map(|s| *s).collect();
+
+ assert_eq!(
+ expected_plan_lines, actual,
+ "\n**Original Plan Mismatch\n\nexpected:\n\n{expected_plan_lines:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ let expected_optimized_lines: Vec<&str> = $EXPECTED_OPTIMIZED_PLAN_LINES
+ .iter().map(|s| *s).collect();
+
+ // Run the actual optimizer
+ let optimized_physical_plan =
+ EnforceSorting::new().optimize(physical_plan, state.config_options())?;
+
+ let formatted = displayable(optimized_physical_plan.as_ref())
+ .indent()
+ .to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+ assert_eq!(
+ expected_optimized_lines, actual,
+ "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected_optimized_lines:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ };
+ }
+
#[tokio::test]
async fn test_remove_unnecessary_sort() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs, source, None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("nullable_col", schema.as_ref()).unwrap(),
- options: SortOptions::default(),
- }];
- let physical_plan = Arc::new(SortExec::try_new(sort_exprs, sort_exec, None)?)
- as Arc<dyn ExecutionPlan>;
- let formatted = displayable(physical_plan.as_ref()).indent().to_string();
- let expected = {
- vec![
- "SortExec: [nullable_col@0 ASC]",
- " SortExec: [non_nullable_col@1 ASC]",
- ]
- };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
- let optimized_physical_plan =
- EnforceSorting::new().optimize(physical_plan, state.config_options())?;
- let formatted = displayable(optimized_physical_plan.as_ref())
- .indent()
- .to_string();
- let expected = { vec!["SortExec: [nullable_col@0 ASC]"] };
- let actual: Vec<&str> = formatted.trim().lines().collect();
- let actual_len = actual.len();
- let actual_trim_last = &actual[..actual_len - 1];
- assert_eq!(
- expected, actual_trim_last,
- "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
- );
+ let source = memory_exec(&schema);
+ let input = sort_exec(vec![sort_expr("non_nullable_col", &schema)], source);
+ let physical_plan = sort_exec(vec![sort_expr("nullable_col", &schema)], input);
+
+ let expected_input = vec![
+ "SortExec: [nullable_col@0 ASC]",
+ " SortExec: [non_nullable_col@1 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ let expected_optimized = vec![
+ "SortExec: [nullable_col@0 ASC]",
+ " MemoryExec: partitions=0, partition_sizes=[]",
+ ];
+ assert_optimized!(expected_input, expected_optimized, physical_plan);
Ok(())
}
#[tokio::test]
async fn test_remove_unnecessary_sort_window_multilayer() -> Result<()> {
- let session_ctx = SessionContext::new();
- let state = session_ctx.state();
let schema = create_test_schema()?;
- let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
- as Arc<dyn ExecutionPlan>;
- let sort_exprs = vec![PhysicalSortExpr {
- expr: col("non_nullable_col", source.schema().as_ref()).unwrap(),
- options: SortOptions {
+ let source = memory_exec(&schema);
+
+ let sort_exprs = vec![sort_expr_options(
+ "non_nullable_col",
+ &source.schema(),
+ SortOptions {
descending: true,
nulls_first: true,
},
- }];
- let sort_exec = Arc::new(SortExec::try_new(sort_exprs.clone(), source, None)?)
- as Arc<dyn ExecutionPlan>;
+ )];
+ let sort = sort_exec(sort_exprs.clone(), source);
+
let window_agg_exec = Arc::new(WindowAggExec::try_new(
Review Comment:
given util function `window_exec` is available, we can use below snippet to create `window_agg_exec`
```rust
let window_agg_exec =
window_exec(sort.clone(), sort.schema(), &sort_exprs, "non_nullable_col")?;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org