Skip to content

Commit

Permalink
tests: simplify tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@tshauck
tshauck committed Jun 18, 2024
1 parent 07b02d1 commit d8bd11a
Show file tree
Hide file tree
Showing 2 changed files with 129 additions and 147 deletions.
235 changes: 91 additions & 144 deletions datafusion/optimizer/src/propagate_empty_relation.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand All @@ -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)
Expand Down Expand Up @@ -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]
Expand Down
41 changes: 38 additions & 3 deletions datafusion/optimizer/src/test/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,21 @@ pub fn assert_analyzed_plan_eq(

Ok(())
}

pub fn assert_analyzed_plan_ne(
rule: Arc<dyn AnalyzerRule + Send + Sync>,
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<dyn AnalyzerRule + Send + Sync>,
plan: LogicalPlan,
Expand Down Expand Up @@ -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<Arc<dyn OptimizerRule + Send + Sync>>,
plan: LogicalPlan,
expected: &str,
) -> Result<()> {
) -> LogicalPlan {
fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
let config = &mut OptimizerContext::new()
.with_max_passes(1)
Expand All @@ -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<Arc<dyn OptimizerRule + Send + Sync>>,
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<Arc<dyn OptimizerRule + Send + Sync>>,
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<dyn OptimizerRule + Send + Sync>,
plan: LogicalPlan,
Expand Down

0 comments on commit d8bd11a

Please sign in to comment.