alamb commented on code in PR #17900:
URL: https://github.com/apache/datafusion/pull/17900#discussion_r2402531273
##########
datafusion/core/tests/physical_optimizer/enforce_sorting.rs:
##########
@@ -91,6 +92,138 @@ fn csv_exec_sorted(
DataSourceExec::from_data_source(config)
}
+/// Runs the sort enforcement optimizer and asserts the plan
+/// against the original and expected plans
+struct EnforceSortingTest {
+ plan: Arc<dyn ExecutionPlan>,
+ repartition_sorts: bool,
+ /// If true, asserts that the input and optimized plans are the same
+ expect_no_change: bool,
+ /// A message printed into the snapshot to describe the expected output
+ expected_description: Option<String>,
+}
+
+impl EnforceSortingTest {
+ fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
+ Self {
+ plan,
+ repartition_sorts: false,
+ expect_no_change: false,
+ expected_description: None,
+ }
+ }
+
+ /// Set whether to repartition sorts
+ fn with_repartition_sorts(mut self, repartition_sorts: bool) -> Self {
+ self.repartition_sorts = repartition_sorts;
+ self
+ }
+
+ /// Set whether to expect no change in the plan
+ fn with_expect_no_change(mut self, expect_no_change: bool) -> Self {
+ self.expect_no_change = expect_no_change;
+ self
+ }
+
+ /// Add an expected output description
+ fn with_expected_description(mut self, description: &str) -> Self {
+ self.expected_description = Some(format!("{description}\n"));
+ self
+ }
+
+ /// Runs the enforce sorting test and returns a string with the input and
+ /// optimized plan as strings for snapshot comparison using insta
+ fn run(&self) -> String {
Review Comment:
this mimics / replicates the logic of `assert_optimized_plan!`
I was not able to remove that macro yet because there are several
parameterized plans that use it that I have to migrate somehow separately
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
