From d8bd11a6428987808f2ffe9f7beec315983c03cd Mon Sep 17 00:00:00 2001 From: Trent Hauck Date: Tue, 18 Jun 2024 14:38:43 -0700 Subject: [PATCH] tests: simplify tests --- .../optimizer/src/propagate_empty_relation.rs | 235 +++++++----------- datafusion/optimizer/src/test/mod.rs | 41 ++- 2 files changed, 129 insertions(+), 147 deletions(-) diff --git a/datafusion/optimizer/src/propagate_empty_relation.rs b/datafusion/optimizer/src/propagate_empty_relation.rs index 9253e34d095c..6b5b1a91cabc 100644 --- a/datafusion/optimizer/src/propagate_empty_relation.rs +++ b/datafusion/optimizer/src/propagate_empty_relation.rs @@ -255,8 +255,9 @@ mod tests { use crate::eliminate_filter::EliminateFilter; use crate::eliminate_nested_union::EliminateNestedUnion; use crate::test::{ - assert_optimized_plan_eq, assert_optimized_plan_eq_with_rules, test_table_scan, - test_table_scan_fields, test_table_scan_with_name, + assert_optimized_plan_eq, assert_optimized_plan_eq_with_rules, + assert_optimized_plan_ne_with_rules, test_table_scan, test_table_scan_fields, + test_table_scan_with_name, }; use super::*; @@ -280,6 +281,21 @@ mod tests { ) } + fn assert_together_optimized_plan_ne( + plan: LogicalPlan, + expected: &str, + ) -> Result<()> { + assert_optimized_plan_ne_with_rules( + vec![ + Arc::new(EliminateFilter::new()), + Arc::new(EliminateNestedUnion::new()), + Arc::new(PropagateEmptyRelation::new()), + ], + plan, + expected, + ) + } + #[test] fn propagate_empty() -> Result<()> { let plan = LogicalPlanBuilder::empty(false) @@ -433,180 +449,111 @@ mod tests { assert_together_optimized_plan_eq(plan, expected) } - #[test] - fn test_left_join_empty_left_table() -> Result<()> { - let left_table_scan = test_table_scan()?; - let left = LogicalPlanBuilder::from(left_table_scan) - .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? - .build()?; + fn empty_left_and_right_lp( + left_empty: bool, + right_empty: bool, + join_type: JoinType, + eq: bool, + ) -> Result<()> { + let left_lp = if left_empty { + let left_table_scan = test_table_scan()?; - let right_table_scan = test_table_scan_with_name("test2")?; - let right = LogicalPlanBuilder::from(right_table_scan).build()?; + let left = LogicalPlanBuilder::from(left_table_scan) + .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? + .build()?; - let plan = LogicalPlanBuilder::from(left) - .join_using( - right, - JoinType::Left, - vec![Column::from_name("a".to_string())], - )? - .build()?; + left + } else { + let scan = test_table_scan_with_name("left").unwrap(); + LogicalPlanBuilder::from(scan).build().unwrap() + }; - let expected = "EmptyRelation"; - assert_together_optimized_plan_eq(plan, expected) - } + let right_lp = if right_empty { + let right_table_scan = test_table_scan_with_name("right")?; - #[test] - fn test_right_join_empty_right_table() -> Result<()> { - let left_table_scan = test_table_scan()?; - let left = LogicalPlanBuilder::from(left_table_scan).build()?; + let right = LogicalPlanBuilder::from(right_table_scan) + .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? + .build()?; - let right_table_scan = test_table_scan_with_name("test2")?; - let right = LogicalPlanBuilder::from(right_table_scan) - .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? - .build()?; + right + } else { + let scan = test_table_scan_with_name("right").unwrap(); + LogicalPlanBuilder::from(scan).build().unwrap() + }; - let plan = LogicalPlanBuilder::from(left) + let plan = LogicalPlanBuilder::from(left_lp) .join_using( - right, - JoinType::Right, + right_lp, + join_type, vec![Column::from_name("a".to_string())], - )? - .build()?; + ) + .unwrap() + .build() + .unwrap(); let expected = "EmptyRelation"; - assert_together_optimized_plan_eq(plan, expected) - } - - #[test] - fn test_left_semi_join_empty_left_table() -> Result<()> { - let left_table_scan = test_table_scan()?; - let left = LogicalPlanBuilder::from(left_table_scan) - .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? - .build()?; - - let right_table_scan = test_table_scan_with_name("test2")?; - let right = LogicalPlanBuilder::from(right_table_scan).build()?; - - let plan = LogicalPlanBuilder::from(left) - .join_using( - right, - JoinType::LeftSemi, - vec![Column::from_name("a".to_string())], - )? - .build()?; - let expected = "EmptyRelation"; - assert_together_optimized_plan_eq(plan, expected) + if eq { + assert_together_optimized_plan_eq(plan, expected) + } else { + assert_together_optimized_plan_ne(plan, expected) + } } #[test] - fn test_left_semi_join_empty_right_table() -> Result<()> { - let left_table_scan = test_table_scan()?; - let left = LogicalPlanBuilder::from(left_table_scan).build()?; + fn test_join_empty_propagation_rules() -> Result<()> { + // test left join with empty left + empty_left_and_right_lp(true, false, JoinType::Left, true)?; - let right_table_scan = test_table_scan_with_name("test2")?; - let right = LogicalPlanBuilder::from(right_table_scan) - .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? - .build()?; + // test right join with empty right + empty_left_and_right_lp(false, true, JoinType::Right, true)?; - let plan = LogicalPlanBuilder::from(left) - .join_using( - right, - JoinType::LeftSemi, - vec![Column::from_name("a".to_string())], - )? - .build()?; + // test left semi join with empty left + empty_left_and_right_lp(true, false, JoinType::LeftSemi, true)?; - let expected = "EmptyRelation"; - assert_together_optimized_plan_eq(plan, expected) - } + // test left semi join with empty right + empty_left_and_right_lp(false, true, JoinType::LeftSemi, true)?; - #[test] - fn test_right_semi_join_empty_right_table() -> Result<()> { - let left_table_scan = test_table_scan()?; - let left = LogicalPlanBuilder::from(left_table_scan).build()?; + // test right semi join with empty left + empty_left_and_right_lp(true, false, JoinType::RightSemi, true)?; - let right_table_scan = test_table_scan_with_name("test2")?; - let right = LogicalPlanBuilder::from(right_table_scan) - .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? - .build()?; + // test right semi join with empty right + empty_left_and_right_lp(false, true, JoinType::RightSemi, true)?; - let plan = LogicalPlanBuilder::from(left) - .join_using( - right, - JoinType::RightSemi, - vec![Column::from_name("a".to_string())], - )? - .build()?; + // test left anti join empty left + empty_left_and_right_lp(true, false, JoinType::LeftAnti, true)?; - let expected = "EmptyRelation"; - assert_together_optimized_plan_eq(plan, expected) + // test right anti join empty right + empty_left_and_right_lp(false, true, JoinType::RightAnti, true) } #[test] - fn test_right_semi_join_empty_left_table() -> Result<()> { - let left_table_scan = test_table_scan()?; - let left = LogicalPlanBuilder::from(left_table_scan) - .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? - .build()?; - - let right_table_scan = test_table_scan_with_name("test2")?; - let right = LogicalPlanBuilder::from(right_table_scan).build()?; - - let plan = LogicalPlanBuilder::from(left) - .join_using( - right, - JoinType::RightSemi, - vec![Column::from_name("a".to_string())], - )? - .build()?; + fn test_join_empty_propagation_rules_noop() -> Result<()> { + // these cases should not result in an empty relation - let expected = "EmptyRelation"; - assert_together_optimized_plan_eq(plan, expected) - } + // test left join with empty right + empty_left_and_right_lp(false, true, JoinType::Left, false)?; - #[test] - fn test_right_anti_join_empty_right_table() -> Result<()> { - let left_table_scan = test_table_scan()?; - let left = LogicalPlanBuilder::from(left_table_scan).build()?; + // test right join with empty left + empty_left_and_right_lp(true, false, JoinType::Right, false)?; - let right_table_scan = test_table_scan_with_name("test2")?; - let right = LogicalPlanBuilder::from(right_table_scan) - .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? - .build()?; - - let plan = LogicalPlanBuilder::from(left) - .join_using( - right, - JoinType::RightAnti, - vec![Column::from_name("a".to_string())], - )? - .build()?; + // test left semi with non-empty left and right + empty_left_and_right_lp(false, false, JoinType::LeftSemi, false)?; - let expected = "EmptyRelation"; - assert_together_optimized_plan_eq(plan, expected) - } + // test right semi with non-empty left and right + empty_left_and_right_lp(false, false, JoinType::RightSemi, false)?; - #[test] - fn test_left_anti_join_empty_left_table() -> Result<()> { - let left_table_scan = test_table_scan()?; - let left = LogicalPlanBuilder::from(left_table_scan) - .filter(Expr::Literal(ScalarValue::Boolean(Some(false))))? - .build()?; + // test left anti join with non-empty left and right + empty_left_and_right_lp(false, false, JoinType::LeftAnti, false)?; - let right_table_scan = test_table_scan_with_name("test2")?; - let right = LogicalPlanBuilder::from(right_table_scan).build()?; + // test left anti with non-empty left and empty right + empty_left_and_right_lp(false, true, JoinType::LeftAnti, false)?; - let plan = LogicalPlanBuilder::from(left) - .join_using( - right, - JoinType::LeftAnti, - vec![Column::from_name("a".to_string())], - )? - .build()?; + // test right anti join with non-empty left and right + empty_left_and_right_lp(false, false, JoinType::RightAnti, false)?; - let expected = "EmptyRelation"; - assert_together_optimized_plan_eq(plan, expected) + // test right anti with empty left and non-empty right + empty_left_and_right_lp(true, false, JoinType::RightAnti, false) } #[test] diff --git a/datafusion/optimizer/src/test/mod.rs b/datafusion/optimizer/src/test/mod.rs index 98d19956df3c..de775f56e786 100644 --- a/datafusion/optimizer/src/test/mod.rs +++ b/datafusion/optimizer/src/test/mod.rs @@ -121,6 +121,21 @@ pub fn assert_analyzed_plan_eq( Ok(()) } + +pub fn assert_analyzed_plan_ne( + rule: Arc, + plan: LogicalPlan, + expected: &str, +) -> Result<()> { + let options = ConfigOptions::default(); + let analyzed_plan = + Analyzer::with_rules(vec![rule]).execute_and_check(plan, &options, |_, _| {})?; + let formatted_plan = format!("{analyzed_plan:?}"); + assert_ne!(formatted_plan, expected); + + Ok(()) +} + pub fn assert_analyzed_plan_eq_display_indent( rule: Arc, plan: LogicalPlan, @@ -169,11 +184,10 @@ pub fn assert_optimized_plan_eq( Ok(()) } -pub fn assert_optimized_plan_eq_with_rules( +fn generate_optimized_plan_with_rules( rules: Vec>, plan: LogicalPlan, - expected: &str, -) -> Result<()> { +) -> LogicalPlan { fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {} let config = &mut OptimizerContext::new() .with_max_passes(1) @@ -182,11 +196,32 @@ pub fn assert_optimized_plan_eq_with_rules( let optimized_plan = optimizer .optimize(plan, config, observe) .expect("failed to optimize plan"); + + optimized_plan +} + +pub fn assert_optimized_plan_eq_with_rules( + rules: Vec>, + plan: LogicalPlan, + expected: &str, +) -> Result<()> { + let optimized_plan = generate_optimized_plan_with_rules(rules, plan); let formatted_plan = format!("{optimized_plan:?}"); assert_eq!(formatted_plan, expected); Ok(()) } +pub fn assert_optimized_plan_ne_with_rules( + rules: Vec>, + plan: LogicalPlan, + expected: &str, +) -> Result<()> { + let optimized_plan = generate_optimized_plan_with_rules(rules, plan); + let formatted_plan = format!("{optimized_plan:?}"); + assert_ne!(formatted_plan, expected); + Ok(()) +} + pub fn assert_optimized_plan_eq_display_indent( rule: Arc, plan: LogicalPlan,