alamb commented on code in PR #16978: URL: https://github.com/apache/datafusion/pull/16978#discussion_r2254815057
########## datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs: ########## @@ -59,128 +61,62 @@ use url::Url; /// /// # Parameters /// -/// * `$EXPECTED_PLAN_LINES`: Expected input plan. -/// * `EXPECTED_OPTIMIZED_PLAN_LINES`: Optimized plan when the flag -/// `prefer_existing_sort` is `false`. -/// * `EXPECTED_PREFER_SORT_ON_OPTIMIZED_PLAN_LINES`: Optimized plan when -/// the flag `prefer_existing_sort` is `true`. /// * `$PLAN`: The plan to optimize. -macro_rules! assert_optimized_prefer_sort_on_off { - ($EXPECTED_PLAN_LINES: expr, $EXPECTED_OPTIMIZED_PLAN_LINES: expr, $EXPECTED_PREFER_SORT_ON_OPTIMIZED_PLAN_LINES: expr, $PLAN: expr, $PREFER_EXISTING_SORT: expr, $SOURCE_UNBOUNDED: expr) => { - if $PREFER_EXISTING_SORT { - assert_optimized!( - $EXPECTED_PLAN_LINES, - $EXPECTED_PREFER_SORT_ON_OPTIMIZED_PLAN_LINES, - $PLAN, - $PREFER_EXISTING_SORT, - $SOURCE_UNBOUNDED - ); - } else { - assert_optimized!( - $EXPECTED_PLAN_LINES, - $EXPECTED_OPTIMIZED_PLAN_LINES, - $PLAN, - $PREFER_EXISTING_SORT, - $SOURCE_UNBOUNDED +/// * `$PREFER_EXISTING_SORT`: Value of the `prefer_existing_sort` flag. +#[macro_export] +macro_rules! get_plan_string { + ($PLAN: expr, $PREFER_EXISTING_SORT: expr, $SOURCE_UNBOUNDED: expr) => {{ + let physical_plan = $PLAN; + let formatted = displayable(physical_plan.as_ref()).indent(true).to_string(); + + // Run the rule top-down + let mut config = ConfigOptions::new(); + config.optimizer.prefer_existing_sort = $PREFER_EXISTING_SORT; + let plan_with_pipeline_fixer = + OrderPreservationContext::new_default(physical_plan); + let parallel = plan_with_pipeline_fixer + .transform_up(|plan_with_pipeline_fixer| { + replace_with_order_preserving_variants( + plan_with_pipeline_fixer, + false, + false, + &config, + ) + }) + .data() + .and_then(check_integrity)?; + let optimized_physical_plan = parallel.plan; + + // Get string representation of the plan + let formatted_optimized = displayable(optimized_physical_plan.as_ref()) + .indent(true) + .to_string(); + + if !$SOURCE_UNBOUNDED { + let ctx = SessionContext::new(); + let object_store = InMemory::new(); + object_store + .put( + &object_store::path::Path::from("file_path"), + bytes::Bytes::from("").into(), + ) + .await?; + ctx.register_object_store( + &Url::parse("test://").unwrap(), + Arc::new(object_store), ); - } - }; -} - -/// Runs the `replace_with_order_preserving_variants` sub-rule and asserts -/// the plan against the original and expected plans for both bounded and -/// unbounded cases. -/// -/// # Parameters -/// -/// * `EXPECTED_UNBOUNDED_PLAN_LINES`: Expected input unbounded plan. -/// * `EXPECTED_BOUNDED_PLAN_LINES`: Expected input bounded plan. -/// * `EXPECTED_UNBOUNDED_OPTIMIZED_PLAN_LINES`: Optimized plan, which is -/// the same regardless of the value of the `prefer_existing_sort` flag. -/// * `EXPECTED_BOUNDED_OPTIMIZED_PLAN_LINES`: Optimized plan when the flag -/// `prefer_existing_sort` is `false` for bounded cases. -/// * `EXPECTED_BOUNDED_PREFER_SORT_ON_OPTIMIZED_PLAN_LINES`: Optimized plan -/// when the flag `prefer_existing_sort` is `true` for bounded cases. -/// * `$PLAN`: The plan to optimize. -/// * `$SOURCE_UNBOUNDED`: Whether the given plan contains an unbounded source. -macro_rules! assert_optimized_in_all_boundedness_situations { Review Comment: oh yeah, I remember this macro -- this is crazy -- 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-unsubscr...@datafusion.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For additional commands, e-mail: github-h...@datafusion.apache.org