You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by bl...@apache.org on 2025/05/09 12:25:47 UTC
(datafusion) branch main updated: Migrate Optimizer tests to insta, part6 (#15984)
This is an automated email from the ASF dual-hosted git repository.
blaginin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 30bb548089 Migrate Optimizer tests to insta, part6 (#15984)
30bb548089 is described below
commit 30bb5480890cc9f05cfe09bf449555a12581f8a7
Author: Tommy shu <qs...@gmail.com>
AuthorDate: Fri May 9 08:25:42 2025 -0400
Migrate Optimizer tests to insta, part6 (#15984)
* migrate tests in `type_coercion.rs` to use snapshot assertions
* remove `assert_analyzed_plan_eq` and `assert_analyzed_plan_with_config_eq`
* remove unnecessary `pub`
* refactor: replace custom assertion functions with snapshot assertions in EliminateLimit tests
* format
* rename
* rename
* refactor: replace custom assertion function with macro for optimized plan equality in tests
* format macro
---
datafusion/optimizer/src/analyzer/type_coercion.rs | 713 +++++++++++++++------
datafusion/optimizer/src/eliminate_limit.rs | 158 +++--
.../src/simplify_expressions/simplify_exprs.rs | 422 +++++++-----
datafusion/optimizer/src/test/mod.rs | 46 +-
4 files changed, 902 insertions(+), 437 deletions(-)
diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs
index d47f7ea6ce..860e041bb4 100644
--- a/datafusion/optimizer/src/analyzer/type_coercion.rs
+++ b/datafusion/optimizer/src/analyzer/type_coercion.rs
@@ -1055,12 +1055,13 @@ mod test {
use arrow::datatypes::DataType::Utf8;
use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, TimeUnit};
+ use insta::assert_snapshot;
use crate::analyzer::type_coercion::{
coerce_case_expression, TypeCoercion, TypeCoercionRewriter,
};
use crate::analyzer::Analyzer;
- use crate::test::{assert_analyzed_plan_eq, assert_analyzed_plan_with_config_eq};
+ use crate::assert_analyzed_plan_with_config_eq_snapshot;
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{TransformedResult, TreeNode};
use datafusion_common::{DFSchema, DFSchemaRef, Result, ScalarValue, Spans};
@@ -1096,13 +1097,83 @@ mod test {
}))
}
+ macro_rules! assert_analyzed_plan_eq {
+ (
+ $plan: expr,
+ @ $expected: literal $(,)?
+ ) => {{
+ let options = ConfigOptions::default();
+ let rule = Arc::new(TypeCoercion::new());
+ assert_analyzed_plan_with_config_eq_snapshot!(
+ options,
+ rule,
+ $plan,
+ @ $expected,
+ )
+ }};
+ }
+
+ macro_rules! coerce_on_output_if_viewtype {
+ (
+ $is_viewtype: expr,
+ $plan: expr,
+ @ $expected: literal $(,)?
+ ) => {{
+ let mut options = ConfigOptions::default();
+ // coerce on output
+ if $is_viewtype {options.optimizer.expand_views_at_output = true;}
+ let rule = Arc::new(TypeCoercion::new());
+
+ assert_analyzed_plan_with_config_eq_snapshot!(
+ options,
+ rule,
+ $plan,
+ @ $expected,
+ )
+ }};
+ }
+
+ fn assert_type_coercion_error(
+ plan: LogicalPlan,
+ expected_substr: &str,
+ ) -> Result<()> {
+ let options = ConfigOptions::default();
+ let analyzer = Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())]);
+
+ match analyzer.execute_and_check(plan, &options, |_, _| {}) {
+ Ok(succeeded_plan) => {
+ panic!(
+ "Expected a type coercion error, but analysis succeeded: \n{:#?}",
+ succeeded_plan
+ );
+ }
+ Err(e) => {
+ let msg = e.to_string();
+ assert!(
+ msg.contains(expected_substr),
+ "Error did not contain expected substring.\n expected to find: `{}`\n actual error: `{}`",
+ expected_substr,
+ msg
+ );
+ }
+ }
+
+ Ok(())
+ }
+
#[test]
fn simple_case() -> Result<()> {
let expr = col("a").lt(lit(2_u32));
let empty = empty_with_type(DataType::Float64);
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected = "Projection: a < CAST(UInt32(2) AS Float64)\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a < CAST(UInt32(2) AS Float64)
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -1137,28 +1208,15 @@ mod test {
Arc::new(analyzed_union),
)?);
- let expected = "Projection: a\n Union\n Projection: CAST(datafusion.test.foo.a AS Int64) AS a\n EmptyRelation\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), top_level_plan, expected)
- }
-
- fn coerce_on_output_if_viewtype(plan: LogicalPlan, expected: &str) -> Result<()> {
- let mut options = ConfigOptions::default();
- options.optimizer.expand_views_at_output = true;
-
- assert_analyzed_plan_with_config_eq(
- options,
- Arc::new(TypeCoercion::new()),
- plan.clone(),
- expected,
- )
- }
-
- fn do_not_coerce_on_output(plan: LogicalPlan, expected: &str) -> Result<()> {
- assert_analyzed_plan_with_config_eq(
- ConfigOptions::default(),
- Arc::new(TypeCoercion::new()),
- plan.clone(),
- expected,
+ assert_analyzed_plan_eq!(
+ top_level_plan,
+ @r"
+ Projection: a
+ Union
+ Projection: CAST(datafusion.test.foo.a AS Int64) AS a
+ EmptyRelation
+ EmptyRelation
+ "
)
}
@@ -1172,12 +1230,26 @@ mod test {
vec![expr.clone()],
Arc::clone(&empty),
)?);
+
// Plan A: no coerce
- let if_not_coerced = "Projection: a\n EmptyRelation";
- do_not_coerce_on_output(plan.clone(), if_not_coerced)?;
+ coerce_on_output_if_viewtype!(
+ false,
+ plan.clone(),
+ @r"
+ Projection: a
+ EmptyRelation
+ "
+ )?;
+
// Plan A: coerce requested: Utf8View => LargeUtf8
- let if_coerced = "Projection: CAST(a AS LargeUtf8)\n EmptyRelation";
- coerce_on_output_if_viewtype(plan.clone(), if_coerced)?;
+ coerce_on_output_if_viewtype!(
+ true,
+ plan.clone(),
+ @r"
+ Projection: CAST(a AS LargeUtf8)
+ EmptyRelation
+ "
+ )?;
// Plan B
// scenario: outermost bool projection
@@ -1187,12 +1259,33 @@ mod test {
Arc::clone(&empty),
)?);
// Plan B: no coerce
- let if_not_coerced =
- "Projection: a < CAST(Utf8(\"foo\") AS Utf8View)\n EmptyRelation";
- do_not_coerce_on_output(bool_plan.clone(), if_not_coerced)?;
+ coerce_on_output_if_viewtype!(
+ false,
+ bool_plan.clone(),
+ @r#"
+ Projection: a < CAST(Utf8("foo") AS Utf8View)
+ EmptyRelation
+ "#
+ )?;
+
+ coerce_on_output_if_viewtype!(
+ false,
+ plan.clone(),
+ @r"
+ Projection: a
+ EmptyRelation
+ "
+ )?;
+
// Plan B: coerce requested: no coercion applied
- let if_coerced = if_not_coerced;
- coerce_on_output_if_viewtype(bool_plan, if_coerced)?;
+ coerce_on_output_if_viewtype!(
+ true,
+ plan.clone(),
+ @r"
+ Projection: CAST(a AS LargeUtf8)
+ EmptyRelation
+ "
+ )?;
// Plan C
// scenario: with a non-projection root logical plan node
@@ -1202,13 +1295,29 @@ mod test {
input: Arc::new(plan),
fetch: None,
});
+
// Plan C: no coerce
- let if_not_coerced =
- "Sort: a ASC NULLS FIRST\n Projection: a\n EmptyRelation";
- do_not_coerce_on_output(sort_plan.clone(), if_not_coerced)?;
+ coerce_on_output_if_viewtype!(
+ false,
+ sort_plan.clone(),
+ @r"
+ Sort: a ASC NULLS FIRST
+ Projection: a
+ EmptyRelation
+ "
+ )?;
+
// Plan C: coerce requested: Utf8View => LargeUtf8
- let if_coerced = "Projection: CAST(a AS LargeUtf8)\n Sort: a ASC NULLS FIRST\n Projection: a\n EmptyRelation";
- coerce_on_output_if_viewtype(sort_plan.clone(), if_coerced)?;
+ coerce_on_output_if_viewtype!(
+ true,
+ sort_plan.clone(),
+ @r"
+ Projection: CAST(a AS LargeUtf8)
+ Sort: a ASC NULLS FIRST
+ Projection: a
+ EmptyRelation
+ "
+ )?;
// Plan D
// scenario: two layers of projections with view types
@@ -1217,11 +1326,27 @@ mod test {
Arc::new(sort_plan),
)?);
// Plan D: no coerce
- let if_not_coerced = "Projection: a\n Sort: a ASC NULLS FIRST\n Projection: a\n EmptyRelation";
- do_not_coerce_on_output(plan.clone(), if_not_coerced)?;
+ coerce_on_output_if_viewtype!(
+ false,
+ plan.clone(),
+ @r"
+ Projection: a
+ Sort: a ASC NULLS FIRST
+ Projection: a
+ EmptyRelation
+ "
+ )?;
// Plan B: coerce requested: Utf8View => LargeUtf8 only on outermost
- let if_coerced = "Projection: CAST(a AS LargeUtf8)\n Sort: a ASC NULLS FIRST\n Projection: a\n EmptyRelation";
- coerce_on_output_if_viewtype(plan.clone(), if_coerced)?;
+ coerce_on_output_if_viewtype!(
+ true,
+ plan.clone(),
+ @r"
+ Projection: CAST(a AS LargeUtf8)
+ Sort: a ASC NULLS FIRST
+ Projection: a
+ EmptyRelation
+ "
+ )?;
Ok(())
}
@@ -1236,12 +1361,26 @@ mod test {
vec![expr.clone()],
Arc::clone(&empty),
)?);
+
// Plan A: no coerce
- let if_not_coerced = "Projection: a\n EmptyRelation";
- do_not_coerce_on_output(plan.clone(), if_not_coerced)?;
+ coerce_on_output_if_viewtype!(
+ false,
+ plan.clone(),
+ @r"
+ Projection: a
+ EmptyRelation
+ "
+ )?;
+
// Plan A: coerce requested: BinaryView => LargeBinary
- let if_coerced = "Projection: CAST(a AS LargeBinary)\n EmptyRelation";
- coerce_on_output_if_viewtype(plan.clone(), if_coerced)?;
+ coerce_on_output_if_viewtype!(
+ true,
+ plan.clone(),
+ @r"
+ Projection: CAST(a AS LargeBinary)
+ EmptyRelation
+ "
+ )?;
// Plan B
// scenario: outermost bool projection
@@ -1250,13 +1389,26 @@ mod test {
vec![bool_expr],
Arc::clone(&empty),
)?);
+
// Plan B: no coerce
- let if_not_coerced =
- "Projection: a < CAST(Binary(\"8,1,8,1\") AS BinaryView)\n EmptyRelation";
- do_not_coerce_on_output(bool_plan.clone(), if_not_coerced)?;
+ coerce_on_output_if_viewtype!(
+ false,
+ bool_plan.clone(),
+ @r#"
+ Projection: a < CAST(Binary("8,1,8,1") AS BinaryView)
+ EmptyRelation
+ "#
+ )?;
+
// Plan B: coerce requested: no coercion applied
- let if_coerced = if_not_coerced;
- coerce_on_output_if_viewtype(bool_plan, if_coerced)?;
+ coerce_on_output_if_viewtype!(
+ true,
+ bool_plan.clone(),
+ @r#"
+ Projection: a < CAST(Binary("8,1,8,1") AS BinaryView)
+ EmptyRelation
+ "#
+ )?;
// Plan C
// scenario: with a non-projection root logical plan node
@@ -1266,13 +1418,28 @@ mod test {
input: Arc::new(plan),
fetch: None,
});
+
// Plan C: no coerce
- let if_not_coerced =
- "Sort: a ASC NULLS FIRST\n Projection: a\n EmptyRelation";
- do_not_coerce_on_output(sort_plan.clone(), if_not_coerced)?;
+ coerce_on_output_if_viewtype!(
+ false,
+ sort_plan.clone(),
+ @r"
+ Sort: a ASC NULLS FIRST
+ Projection: a
+ EmptyRelation
+ "
+ )?;
// Plan C: coerce requested: BinaryView => LargeBinary
- let if_coerced = "Projection: CAST(a AS LargeBinary)\n Sort: a ASC NULLS FIRST\n Projection: a\n EmptyRelation";
- coerce_on_output_if_viewtype(sort_plan.clone(), if_coerced)?;
+ coerce_on_output_if_viewtype!(
+ true,
+ sort_plan.clone(),
+ @r"
+ Projection: CAST(a AS LargeBinary)
+ Sort: a ASC NULLS FIRST
+ Projection: a
+ EmptyRelation
+ "
+ )?;
// Plan D
// scenario: two layers of projections with view types
@@ -1280,12 +1447,30 @@ mod test {
vec![col("a")],
Arc::new(sort_plan),
)?);
+
// Plan D: no coerce
- let if_not_coerced = "Projection: a\n Sort: a ASC NULLS FIRST\n Projection: a\n EmptyRelation";
- do_not_coerce_on_output(plan.clone(), if_not_coerced)?;
+ coerce_on_output_if_viewtype!(
+ false,
+ plan.clone(),
+ @r"
+ Projection: a
+ Sort: a ASC NULLS FIRST
+ Projection: a
+ EmptyRelation
+ "
+ )?;
+
// Plan B: coerce requested: BinaryView => LargeBinary only on outermost
- let if_coerced = "Projection: CAST(a AS LargeBinary)\n Sort: a ASC NULLS FIRST\n Projection: a\n EmptyRelation";
- coerce_on_output_if_viewtype(plan.clone(), if_coerced)?;
+ coerce_on_output_if_viewtype!(
+ true,
+ plan.clone(),
+ @r"
+ Projection: CAST(a AS LargeBinary)
+ Sort: a ASC NULLS FIRST
+ Projection: a
+ EmptyRelation
+ "
+ )?;
Ok(())
}
@@ -1299,9 +1484,14 @@ mod test {
vec![expr.clone().or(expr)],
empty,
)?);
- let expected = "Projection: a < CAST(UInt32(2) AS Float64) OR a < CAST(UInt32(2) AS Float64)\
- \n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a < CAST(UInt32(2) AS Float64) OR a < CAST(UInt32(2) AS Float64)
+ EmptyRelation
+ "
+ )
}
#[derive(Debug, Clone)]
@@ -1340,9 +1530,14 @@ mod test {
})
.call(vec![lit(123_i32)]);
let plan = LogicalPlan::Projection(Projection::try_new(vec![udf], empty)?);
- let expected =
- "Projection: TestScalarUDF(CAST(Int32(123) AS Float32))\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: TestScalarUDF(CAST(Int32(123) AS Float32))
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -1372,9 +1567,14 @@ mod test {
vec![scalar_function_expr],
empty,
)?);
- let expected =
- "Projection: TestScalarUDF(CAST(Int64(10) AS Float32))\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: TestScalarUDF(CAST(Int64(10) AS Float32))
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -1397,8 +1597,14 @@ mod test {
None,
));
let plan = LogicalPlan::Projection(Projection::try_new(vec![udaf], empty)?);
- let expected = "Projection: MY_AVG(CAST(Int64(10) AS Float64))\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: MY_AVG(CAST(Int64(10) AS Float64))
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -1445,8 +1651,14 @@ mod test {
None,
));
let plan = LogicalPlan::Projection(Projection::try_new(vec![agg_expr], empty)?);
- let expected = "Projection: avg(Float64(12))\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: avg(Float64(12))
+ EmptyRelation
+ "
+ )?;
let empty = empty_with_type(DataType::Int32);
let agg_expr = Expr::AggregateFunction(expr::AggregateFunction::new_udf(
@@ -1458,9 +1670,14 @@ mod test {
None,
));
let plan = LogicalPlan::Projection(Projection::try_new(vec![agg_expr], empty)?);
- let expected = "Projection: avg(CAST(a AS Float64))\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: avg(CAST(a AS Float64))
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -1489,10 +1706,14 @@ mod test {
+ lit(ScalarValue::new_interval_dt(123, 456));
let empty = empty();
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected =
- "Projection: CAST(Utf8(\"1998-03-18\") AS Date32) + IntervalDayTime(\"IntervalDayTime { days: 123, milliseconds: 456 }\")\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Projection: CAST(Utf8("1998-03-18") AS Date32) + IntervalDayTime("IntervalDayTime { days: 123, milliseconds: 456 }")
+ EmptyRelation
+ "#
+ )
}
#[test]
@@ -1501,8 +1722,12 @@ mod test {
let expr = col("a").in_list(vec![lit(1_i32), lit(4_i8), lit(8_i64)], false);
let empty = empty_with_type(DataType::Int64);
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected = "Projection: a IN ([CAST(Int32(1) AS Int64), CAST(Int8(4) AS Int64), Int64(8)])\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a IN ([CAST(Int32(1) AS Int64), CAST(Int8(4) AS Int64), Int64(8)])
+ EmptyRelation
+ ")?;
// a in (1,4,8), a is decimal
let expr = col("a").in_list(vec![lit(1_i32), lit(4_i8), lit(8_i64)], false);
@@ -1514,8 +1739,12 @@ mod test {
)?),
}));
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected = "Projection: CAST(a AS Decimal128(24, 4)) IN ([CAST(Int32(1) AS Decimal128(24, 4)), CAST(Int8(4) AS Decimal128(24, 4)), CAST(Int64(8) AS Decimal128(24, 4))])\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: CAST(a AS Decimal128(24, 4)) IN ([CAST(Int32(1) AS Decimal128(24, 4)), CAST(Int8(4) AS Decimal128(24, 4)), CAST(Int64(8) AS Decimal128(24, 4))])
+ EmptyRelation
+ ")
}
#[test]
@@ -1528,10 +1757,14 @@ mod test {
);
let empty = empty_with_type(Utf8);
let plan = LogicalPlan::Filter(Filter::try_new(expr, empty)?);
- let expected =
- "Filter: CAST(a AS Date32) BETWEEN CAST(Utf8(\"2002-05-08\") AS Date32) AND CAST(Utf8(\"2002-05-08\") AS Date32) + IntervalYearMonth(\"1\")\
- \n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Filter: CAST(a AS Date32) BETWEEN CAST(Utf8("2002-05-08") AS Date32) AND CAST(Utf8("2002-05-08") AS Date32) + IntervalYearMonth("1")
+ EmptyRelation
+ "#
+ )
}
#[test]
@@ -1544,11 +1777,15 @@ mod test {
);
let empty = empty_with_type(Utf8);
let plan = LogicalPlan::Filter(Filter::try_new(expr, empty)?);
+
// TODO: we should cast col(a).
- let expected =
- "Filter: CAST(a AS Date32) BETWEEN CAST(Utf8(\"2002-05-08\") AS Date32) + IntervalYearMonth(\"1\") AND CAST(Utf8(\"2002-12-08\") AS Date32)\
- \n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Filter: CAST(a AS Date32) BETWEEN CAST(Utf8("2002-05-08") AS Date32) + IntervalYearMonth("1") AND CAST(Utf8("2002-12-08") AS Date32)
+ EmptyRelation
+ "#
+ )
}
#[test]
@@ -1556,10 +1793,14 @@ mod test {
let expr = lit(ScalarValue::Null).between(lit(ScalarValue::Null), lit(2i64));
let empty = empty();
let plan = LogicalPlan::Filter(Filter::try_new(expr, empty)?);
- let expected =
- "Filter: CAST(NULL AS Int64) BETWEEN CAST(NULL AS Int64) AND Int64(2)\
- \n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Filter: CAST(NULL AS Int64) BETWEEN CAST(NULL AS Int64) AND Int64(2)
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -1569,37 +1810,60 @@ mod test {
let empty = empty_with_type(DataType::Boolean);
let plan =
LogicalPlan::Projection(Projection::try_new(vec![expr.clone()], empty)?);
- let expected = "Projection: a IS TRUE\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a IS TRUE
+ EmptyRelation
+ "
+ )?;
let empty = empty_with_type(DataType::Int64);
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let ret = assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, "");
- let err = ret.unwrap_err().to_string();
- assert!(err.contains("Cannot infer common argument type for comparison operation Int64 IS DISTINCT FROM Boolean"), "{err}");
+ assert_type_coercion_error(
+ plan,
+ "Cannot infer common argument type for comparison operation Int64 IS DISTINCT FROM Boolean"
+ )?;
// is not true
let expr = col("a").is_not_true();
let empty = empty_with_type(DataType::Boolean);
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected = "Projection: a IS NOT TRUE\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a IS NOT TRUE
+ EmptyRelation
+ "
+ )?;
// is false
let expr = col("a").is_false();
let empty = empty_with_type(DataType::Boolean);
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected = "Projection: a IS FALSE\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a IS FALSE
+ EmptyRelation
+ "
+ )?;
// is not false
let expr = col("a").is_not_false();
let empty = empty_with_type(DataType::Boolean);
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected = "Projection: a IS NOT FALSE\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a IS NOT FALSE
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -1610,27 +1874,38 @@ mod test {
let like_expr = Expr::Like(Like::new(false, expr, pattern, None, false));
let empty = empty_with_type(Utf8);
let plan = LogicalPlan::Projection(Projection::try_new(vec![like_expr], empty)?);
- let expected = "Projection: a LIKE Utf8(\"abc\")\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Projection: a LIKE Utf8("abc")
+ EmptyRelation
+ "#
+ )?;
let expr = Box::new(col("a"));
let pattern = Box::new(lit(ScalarValue::Null));
let like_expr = Expr::Like(Like::new(false, expr, pattern, None, false));
let empty = empty_with_type(Utf8);
let plan = LogicalPlan::Projection(Projection::try_new(vec![like_expr], empty)?);
- let expected = "Projection: a LIKE CAST(NULL AS Utf8)\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a LIKE CAST(NULL AS Utf8)
+ EmptyRelation
+ "
+ )?;
let expr = Box::new(col("a"));
let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
let like_expr = Expr::Like(Like::new(false, expr, pattern, None, false));
let empty = empty_with_type(DataType::Int64);
let plan = LogicalPlan::Projection(Projection::try_new(vec![like_expr], empty)?);
- let err = assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected);
- assert!(err.is_err());
- assert!(err.unwrap_err().to_string().contains(
- "There isn't a common type to coerce Int64 and Utf8 in LIKE expression"
- ));
+ assert_type_coercion_error(
+ plan,
+ "There isn't a common type to coerce Int64 and Utf8 in LIKE expression",
+ )?;
// ilike
let expr = Box::new(col("a"));
@@ -1638,27 +1913,39 @@ mod test {
let ilike_expr = Expr::Like(Like::new(false, expr, pattern, None, true));
let empty = empty_with_type(Utf8);
let plan = LogicalPlan::Projection(Projection::try_new(vec![ilike_expr], empty)?);
- let expected = "Projection: a ILIKE Utf8(\"abc\")\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Projection: a ILIKE Utf8("abc")
+ EmptyRelation
+ "#
+ )?;
let expr = Box::new(col("a"));
let pattern = Box::new(lit(ScalarValue::Null));
let ilike_expr = Expr::Like(Like::new(false, expr, pattern, None, true));
let empty = empty_with_type(Utf8);
let plan = LogicalPlan::Projection(Projection::try_new(vec![ilike_expr], empty)?);
- let expected = "Projection: a ILIKE CAST(NULL AS Utf8)\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a ILIKE CAST(NULL AS Utf8)
+ EmptyRelation
+ "
+ )?;
let expr = Box::new(col("a"));
let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
let ilike_expr = Expr::Like(Like::new(false, expr, pattern, None, true));
let empty = empty_with_type(DataType::Int64);
let plan = LogicalPlan::Projection(Projection::try_new(vec![ilike_expr], empty)?);
- let err = assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected);
- assert!(err.is_err());
- assert!(err.unwrap_err().to_string().contains(
- "There isn't a common type to coerce Int64 and Utf8 in ILIKE expression"
- ));
+ assert_type_coercion_error(
+ plan,
+ "There isn't a common type to coerce Int64 and Utf8 in ILIKE expression",
+ )?;
+
Ok(())
}
@@ -1669,23 +1956,34 @@ mod test {
let empty = empty_with_type(DataType::Boolean);
let plan =
LogicalPlan::Projection(Projection::try_new(vec![expr.clone()], empty)?);
- let expected = "Projection: a IS UNKNOWN\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a IS UNKNOWN
+ EmptyRelation
+ "
+ )?;
let empty = empty_with_type(Utf8);
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let ret = assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected);
- let err = ret.unwrap_err().to_string();
- assert!(err.contains("Cannot infer common argument type for comparison operation Utf8 IS DISTINCT FROM Boolean"), "{err}");
+ assert_type_coercion_error(
+ plan,
+ "Cannot infer common argument type for comparison operation Utf8 IS DISTINCT FROM Boolean"
+ )?;
// is not unknown
let expr = col("a").is_not_unknown();
let empty = empty_with_type(DataType::Boolean);
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected = "Projection: a IS NOT UNKNOWN\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Projection: a IS NOT UNKNOWN
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -1694,21 +1992,19 @@ mod test {
let args = [col("a"), lit("b"), lit(true), lit(false), lit(13)];
// concat-type signature
- {
- let expr = ScalarUDF::new_from_impl(TestScalarUDF {
- signature: Signature::variadic(vec![Utf8], Volatility::Immutable),
- })
- .call(args.to_vec());
- let plan = LogicalPlan::Projection(Projection::try_new(
- vec![expr],
- Arc::clone(&empty),
- )?);
- let expected =
- "Projection: TestScalarUDF(a, Utf8(\"b\"), CAST(Boolean(true) AS Utf8), CAST(Boolean(false) AS Utf8), CAST(Int32(13) AS Utf8))\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- }
-
- Ok(())
+ let expr = ScalarUDF::new_from_impl(TestScalarUDF {
+ signature: Signature::variadic(vec![Utf8], Volatility::Immutable),
+ })
+ .call(args.to_vec());
+ let plan =
+ LogicalPlan::Projection(Projection::try_new(vec![expr], Arc::clone(&empty))?);
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Projection: TestScalarUDF(a, Utf8("b"), CAST(Boolean(true) AS Utf8), CAST(Boolean(false) AS Utf8), CAST(Int32(13) AS Utf8))
+ EmptyRelation
+ "#
+ )
}
#[test]
@@ -1758,10 +2054,14 @@ mod test {
.eq(cast(lit("1998-03-18"), DataType::Date32));
let empty = empty();
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected =
- "Projection: CAST(Utf8(\"1998-03-18\") AS Timestamp(Nanosecond, None)) = CAST(CAST(Utf8(\"1998-03-18\") AS Date32) AS Timestamp(Nanosecond, None))\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Projection: CAST(Utf8("1998-03-18") AS Timestamp(Nanosecond, None)) = CAST(CAST(Utf8("1998-03-18") AS Date32) AS Timestamp(Nanosecond, None))
+ EmptyRelation
+ "#
+ )
}
fn cast_if_not_same_type(
@@ -1882,12 +2182,9 @@ mod test {
else_expr: Some(Box::new(col("string"))),
};
let err = coerce_case_expression(case, &schema).unwrap_err();
- assert_eq!(
+ assert_snapshot!(
err.strip_backtrace(),
- "Error during planning: \
- Failed to coerce case (Interval(MonthDayNano)) and \
- when ([Float32, Binary, Utf8]) to common types in \
- CASE WHEN expression"
+ @"Error during planning: Failed to coerce case (Interval(MonthDayNano)) and when ([Float32, Binary, Utf8]) to common types in CASE WHEN expression"
);
let case = Case {
@@ -1900,12 +2197,9 @@ mod test {
else_expr: Some(Box::new(col("timestamp"))),
};
let err = coerce_case_expression(case, &schema).unwrap_err();
- assert_eq!(
+ assert_snapshot!(
err.strip_backtrace(),
- "Error during planning: \
- Failed to coerce then ([Date32, Float32, Binary]) and \
- else (Some(Timestamp(Nanosecond, None))) to common types \
- in CASE WHEN expression"
+ @"Error during planning: Failed to coerce then ([Date32, Float32, Binary]) and else (Some(Timestamp(Nanosecond, None))) to common types in CASE WHEN expression"
);
Ok(())
@@ -2108,12 +2402,14 @@ mod test {
let expr = col("a").eq(cast(col("a"), may_type_cutsom));
let empty = empty_with_type(map_type_entries);
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected = "Projection: a = CAST(CAST(a AS Map(Field { name: \"key_value\", data_type: Struct([Field { name: \"key\", data_type: Utf8, \
- nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: \"value\", data_type: Float64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }]), \
- nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, false)) AS Map(Field { name: \"entries\", data_type: Struct([Field { name: \"key\", data_type: Utf8, nullable: false, \
- dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: \"value\", data_type: Float64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, false))\n \
- EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Projection: a = CAST(CAST(a AS Map(Field { name: "key_value", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: Float64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, false)) AS Map(Field { name: "entries", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, [...]
+ EmptyRelation
+ "#
+ )
}
#[test]
@@ -2129,9 +2425,14 @@ mod test {
));
let empty = empty();
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected = "Projection: IntervalYearMonth(\"12\") + CAST(Utf8(\"2000-01-01T00:00:00\") AS Timestamp(Nanosecond, None))\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Projection: IntervalYearMonth("12") + CAST(Utf8("2000-01-01T00:00:00") AS Timestamp(Nanosecond, None))
+ EmptyRelation
+ "#
+ )
}
#[test]
@@ -2149,10 +2450,14 @@ mod test {
));
let empty = empty();
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
- let expected =
- "Projection: CAST(Utf8(\"1998-03-18\") AS Timestamp(Nanosecond, None)) - CAST(Utf8(\"1998-03-18\") AS Timestamp(Nanosecond, None))\n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r#"
+ Projection: CAST(Utf8("1998-03-18") AS Timestamp(Nanosecond, None)) - CAST(Utf8("1998-03-18") AS Timestamp(Nanosecond, None))
+ EmptyRelation
+ "#
+ )
}
#[test]
@@ -2171,14 +2476,17 @@ mod test {
));
let plan = LogicalPlan::Filter(Filter::try_new(in_subquery_expr, empty_int64)?);
// add cast for subquery
- let expected = "\
- Filter: a IN (<subquery>)\
- \n Subquery:\
- \n Projection: CAST(a AS Int64)\
- \n EmptyRelation\
- \n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Filter: a IN (<subquery>)
+ Subquery:
+ Projection: CAST(a AS Int64)
+ EmptyRelation
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -2196,14 +2504,17 @@ mod test {
false,
));
let plan = LogicalPlan::Filter(Filter::try_new(in_subquery_expr, empty_int32)?);
+
// add cast for subquery
- let expected = "\
- Filter: CAST(a AS Int64) IN (<subquery>)\
- \n Subquery:\
- \n EmptyRelation\
- \n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Filter: CAST(a AS Int64) IN (<subquery>)
+ Subquery:
+ EmptyRelation
+ EmptyRelation
+ "
+ )
}
#[test]
@@ -2221,13 +2532,17 @@ mod test {
false,
));
let plan = LogicalPlan::Filter(Filter::try_new(in_subquery_expr, empty_outside)?);
+
// add cast for subquery
- let expected = "Filter: CAST(a AS Decimal128(13, 8)) IN (<subquery>)\
- \n Subquery:\
- \n Projection: CAST(a AS Decimal128(13, 8))\
- \n EmptyRelation\
- \n EmptyRelation";
- assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
- Ok(())
+ assert_analyzed_plan_eq!(
+ plan,
+ @r"
+ Filter: CAST(a AS Decimal128(13, 8)) IN (<subquery>)
+ Subquery:
+ Projection: CAST(a AS Decimal128(13, 8))
+ EmptyRelation
+ EmptyRelation
+ "
+ )
}
}
diff --git a/datafusion/optimizer/src/eliminate_limit.rs b/datafusion/optimizer/src/eliminate_limit.rs
index 5d3a1b223b..2007e0c820 100644
--- a/datafusion/optimizer/src/eliminate_limit.rs
+++ b/datafusion/optimizer/src/eliminate_limit.rs
@@ -90,7 +90,6 @@ impl OptimizerRule for EliminateLimit {
#[cfg(test)]
mod tests {
use super::*;
- use crate::optimizer::Optimizer;
use crate::test::*;
use crate::OptimizerContext;
use datafusion_common::Column;
@@ -100,36 +99,43 @@ mod tests {
};
use std::sync::Arc;
+ use crate::assert_optimized_plan_eq_snapshot;
use crate::push_down_limit::PushDownLimit;
use datafusion_expr::test::function_stub::sum;
- fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
- fn assert_optimized_plan_eq(plan: LogicalPlan, expected: &str) -> Result<()> {
- let optimizer = Optimizer::with_rules(vec![Arc::new(EliminateLimit::new())]);
- let optimized_plan =
- optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
-
- let formatted_plan = format!("{optimized_plan}");
- assert_eq!(formatted_plan, expected);
- Ok(())
+ macro_rules! assert_optimized_plan_equal {
+ (
+ $plan:expr,
+ @ $expected:literal $(,)?
+ ) => {{
+ let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(EliminateLimit::new())];
+ let optimizer_ctx = OptimizerContext::new();
+ assert_optimized_plan_eq_snapshot!(
+ optimizer_ctx,
+ rules,
+ $plan,
+ @ $expected,
+ )
+ }};
}
- fn assert_optimized_plan_eq_with_pushdown(
- plan: LogicalPlan,
- expected: &str,
- ) -> Result<()> {
- fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
- let config = OptimizerContext::new().with_max_passes(1);
- let optimizer = Optimizer::with_rules(vec![
- Arc::new(PushDownLimit::new()),
- Arc::new(EliminateLimit::new()),
- ]);
- let optimized_plan = optimizer
- .optimize(plan, &config, observe)
- .expect("failed to optimize plan");
- let formatted_plan = format!("{optimized_plan}");
- assert_eq!(formatted_plan, expected);
- Ok(())
+ macro_rules! assert_optimized_plan_eq_with_pushdown {
+ (
+ $plan:expr,
+ @ $expected:literal $(,)?
+ ) => {{
+ let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+ let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![
+ Arc::new(PushDownLimit::new()),
+ Arc::new(EliminateLimit::new())
+ ];
+ assert_optimized_plan_eq_snapshot!(
+ optimizer_ctx,
+ rules,
+ $plan,
+ @ $expected,
+ )
+ }};
}
#[test]
@@ -140,8 +146,10 @@ mod tests {
.limit(0, Some(0))?
.build()?;
// No aggregate / scan / limit
- let expected = "EmptyRelation";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"EmptyRelation"
+ )
}
#[test]
@@ -157,11 +165,15 @@ mod tests {
.build()?;
// Left side is removed
- let expected = "Union\
- \n EmptyRelation\
- \n Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Union
+ EmptyRelation
+ Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]
+ TableScan: test
+ "
+ )
}
#[test]
@@ -174,8 +186,10 @@ mod tests {
.build()?;
// No aggregate / scan / limit
- let expected = "EmptyRelation";
- assert_optimized_plan_eq_with_pushdown(plan, expected)
+ assert_optimized_plan_eq_with_pushdown!(
+ plan,
+ @ "EmptyRelation"
+ )
}
#[test]
@@ -190,12 +204,16 @@ mod tests {
// After remove global-state, we don't record the parent <skip, fetch>
// So, bottom don't know parent info, so can't eliminate.
- let expected = "Limit: skip=2, fetch=1\
- \n Sort: test.a ASC NULLS LAST, fetch=3\
- \n Limit: skip=0, fetch=2\
- \n Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]\
- \n TableScan: test";
- assert_optimized_plan_eq_with_pushdown(plan, expected)
+ assert_optimized_plan_eq_with_pushdown!(
+ plan,
+ @ r"
+ Limit: skip=2, fetch=1
+ Sort: test.a ASC NULLS LAST, fetch=3
+ Limit: skip=0, fetch=2
+ Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]
+ TableScan: test
+ "
+ )
}
#[test]
@@ -208,12 +226,16 @@ mod tests {
.limit(0, Some(1))?
.build()?;
- let expected = "Limit: skip=0, fetch=1\
- \n Sort: test.a ASC NULLS LAST\
- \n Limit: skip=0, fetch=2\
- \n Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Limit: skip=0, fetch=1
+ Sort: test.a ASC NULLS LAST
+ Limit: skip=0, fetch=2
+ Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]
+ TableScan: test
+ "
+ )
}
#[test]
@@ -226,12 +248,16 @@ mod tests {
.limit(3, Some(1))?
.build()?;
- let expected = "Limit: skip=3, fetch=1\
- \n Sort: test.a ASC NULLS LAST\
- \n Limit: skip=2, fetch=1\
- \n Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Limit: skip=3, fetch=1
+ Sort: test.a ASC NULLS LAST
+ Limit: skip=2, fetch=1
+ Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]
+ TableScan: test
+ "
+ )
}
#[test]
@@ -248,12 +274,16 @@ mod tests {
.limit(3, Some(1))?
.build()?;
- let expected = "Limit: skip=3, fetch=1\
- \n Inner Join: Using test.a = test1.a\
- \n Limit: skip=2, fetch=1\
- \n TableScan: test\
- \n TableScan: test1";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Limit: skip=3, fetch=1
+ Inner Join: Using test.a = test1.a
+ Limit: skip=2, fetch=1
+ TableScan: test
+ TableScan: test1
+ "
+ )
}
#[test]
@@ -264,8 +294,12 @@ mod tests {
.limit(0, None)?
.build()?;
- let expected = "Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b)]]
+ TableScan: test
+ "
+ )
}
}
diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
index 6314209dc7..ccf90893e1 100644
--- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
+++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
@@ -155,12 +155,12 @@ mod tests {
use arrow::datatypes::{DataType, Field, Schema};
use chrono::{DateTime, Utc};
- use crate::optimizer::Optimizer;
use datafusion_expr::logical_plan::builder::table_scan_with_filters;
use datafusion_expr::logical_plan::table_scan;
use datafusion_expr::*;
use datafusion_functions_aggregate::expr_fn::{max, min};
+ use crate::assert_optimized_plan_eq_snapshot;
use crate::test::{assert_fields_eq, test_table_scan_with_name};
use crate::OptimizerContext;
@@ -180,15 +180,20 @@ mod tests {
.expect("building plan")
}
- fn assert_optimized_plan_eq(plan: LogicalPlan, expected: &str) -> Result<()> {
- // Use Optimizer to do plan traversal
- fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
- let optimizer = Optimizer::with_rules(vec![Arc::new(SimplifyExpressions::new())]);
- let optimized_plan =
- optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
- let formatted_plan = format!("{optimized_plan}");
- assert_eq!(formatted_plan, expected);
- Ok(())
+ macro_rules! assert_optimized_plan_equal {
+ (
+ $plan:expr,
+ @ $expected:literal $(,)?
+ ) => {{
+ let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(SimplifyExpressions::new())];
+ let optimizer_ctx = OptimizerContext::new();
+ assert_optimized_plan_eq_snapshot!(
+ optimizer_ctx,
+ rules,
+ $plan,
+ @ $expected,
+ )
+ }};
}
#[test]
@@ -211,9 +216,10 @@ mod tests {
assert_eq!(1, table_scan.schema().fields().len());
assert_fields_eq(&table_scan, vec!["a"]);
- let expected = "TableScan: test projection=[a], full_filters=[Boolean(true)]";
-
- assert_optimized_plan_eq(table_scan, expected)
+ assert_optimized_plan_equal!(
+ table_scan,
+ @ r"TableScan: test projection=[a], full_filters=[Boolean(true)]"
+ )
}
#[test]
@@ -224,12 +230,13 @@ mod tests {
.filter(and(col("b").gt(lit(1)), col("b").gt(lit(1))))?
.build()?;
- assert_optimized_plan_eq(
+ assert_optimized_plan_equal!(
plan,
- "\
- Filter: test.b > Int32(1)\
- \n Projection: test.a\
- \n TableScan: test",
+ @ r"
+ Filter: test.b > Int32(1)
+ Projection: test.a
+ TableScan: test
+ "
)
}
@@ -241,12 +248,13 @@ mod tests {
.filter(and(col("b").gt(lit(1)), col("b").gt(lit(1))))?
.build()?;
- assert_optimized_plan_eq(
+ assert_optimized_plan_equal!(
plan,
- "\
- Filter: test.b > Int32(1)\
- \n Projection: test.a\
- \n TableScan: test",
+ @ r"
+ Filter: test.b > Int32(1)
+ Projection: test.a
+ TableScan: test
+ "
)
}
@@ -258,12 +266,13 @@ mod tests {
.filter(or(col("b").gt(lit(1)), col("b").gt(lit(1))))?
.build()?;
- assert_optimized_plan_eq(
+ assert_optimized_plan_equal!(
plan,
- "\
- Filter: test.b > Int32(1)\
- \n Projection: test.a\
- \n TableScan: test",
+ @ r"
+ Filter: test.b > Int32(1)
+ Projection: test.a
+ TableScan: test
+ "
)
}
@@ -279,12 +288,13 @@ mod tests {
))?
.build()?;
- assert_optimized_plan_eq(
+ assert_optimized_plan_equal!(
plan,
- "\
- Filter: test.a > Int32(5) AND test.b < Int32(6)\
- \n Projection: test.a, test.b\
- \n TableScan: test",
+ @ r"
+ Filter: test.a > Int32(5) AND test.b < Int32(6)
+ Projection: test.a, test.b
+ TableScan: test
+ "
)
}
@@ -297,13 +307,15 @@ mod tests {
.project(vec![col("a")])?
.build()?;
- let expected = "\
- Projection: test.a\
- \n Filter: NOT test.c\
- \n Filter: test.b\
- \n TableScan: test";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Projection: test.a
+ Filter: NOT test.c
+ Filter: test.b
+ TableScan: test
+ "
+ )
}
#[test]
@@ -316,14 +328,16 @@ mod tests {
.project(vec![col("a")])?
.build()?;
- let expected = "\
- Projection: test.a\
- \n Limit: skip=0, fetch=1\
- \n Filter: test.c\
- \n Filter: NOT test.b\
- \n TableScan: test";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Projection: test.a
+ Limit: skip=0, fetch=1
+ Filter: test.c
+ Filter: NOT test.b
+ TableScan: test
+ "
+ )
}
#[test]
@@ -334,12 +348,14 @@ mod tests {
.project(vec![col("a")])?
.build()?;
- let expected = "\
- Projection: test.a\
- \n Filter: NOT test.b AND test.c\
- \n TableScan: test";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Projection: test.a
+ Filter: NOT test.b AND test.c
+ TableScan: test
+ "
+ )
}
#[test]
@@ -350,12 +366,14 @@ mod tests {
.project(vec![col("a")])?
.build()?;
- let expected = "\
- Projection: test.a\
- \n Filter: NOT test.b OR NOT test.c\
- \n TableScan: test";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Projection: test.a
+ Filter: NOT test.b OR NOT test.c
+ TableScan: test
+ "
+ )
}
#[test]
@@ -366,12 +384,14 @@ mod tests {
.project(vec![col("a")])?
.build()?;
- let expected = "\
- Projection: test.a\
- \n Filter: test.b\
- \n TableScan: test";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Projection: test.a
+ Filter: test.b
+ TableScan: test
+ "
+ )
}
#[test]
@@ -381,11 +401,13 @@ mod tests {
.project(vec![col("a"), col("d"), col("b").eq(lit(false))])?
.build()?;
- let expected = "\
- Projection: test.a, test.d, NOT test.b AS test.b = Boolean(false)\
- \n TableScan: test";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Projection: test.a, test.d, NOT test.b AS test.b = Boolean(false)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -399,12 +421,14 @@ mod tests {
)?
.build()?;
- let expected = "\
- Aggregate: groupBy=[[test.a, test.c]], aggr=[[max(test.b) AS max(test.b = Boolean(true)), min(test.b)]]\
- \n Projection: test.a, test.c, test.b\
- \n TableScan: test";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Aggregate: groupBy=[[test.a, test.c]], aggr=[[max(test.b) AS max(test.b = Boolean(true)), min(test.b)]]
+ Projection: test.a, test.c, test.b
+ TableScan: test
+ "
+ )
}
#[test]
@@ -422,10 +446,10 @@ mod tests {
let values = vec![vec![expr1, expr2]];
let plan = LogicalPlanBuilder::values(values)?.build()?;
- let expected = "\
- Values: (Int32(3) AS Int32(1) + Int32(2), Int32(1) AS Int32(2) - Int32(1))";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ "Values: (Int32(3) AS Int32(1) + Int32(2), Int32(1) AS Int32(2) - Int32(1))"
+ )
}
fn get_optimized_plan_formatted(
@@ -482,10 +506,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("d").gt(lit(10)).not())?
.build()?;
- let expected = "Filter: test.d <= Int32(10)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d <= Int32(10)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -495,10 +523,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("d").gt(lit(10)).and(col("d").lt(lit(100))).not())?
.build()?;
- let expected = "Filter: test.d <= Int32(10) OR test.d >= Int32(100)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d <= Int32(10) OR test.d >= Int32(100)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -508,10 +540,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("d").gt(lit(10)).or(col("d").lt(lit(100))).not())?
.build()?;
- let expected = "Filter: test.d <= Int32(10) AND test.d >= Int32(100)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d <= Int32(10) AND test.d >= Int32(100)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -521,10 +557,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("d").gt(lit(10)).not().not())?
.build()?;
- let expected = "Filter: test.d > Int32(10)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d > Int32(10)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -534,10 +574,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("e").is_null().not())?
.build()?;
- let expected = "Filter: test.e IS NOT NULL\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.e IS NOT NULL
+ TableScan: test
+ "
+ )
}
#[test]
@@ -547,10 +591,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("e").is_not_null().not())?
.build()?;
- let expected = "Filter: test.e IS NULL\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.e IS NULL
+ TableScan: test
+ "
+ )
}
#[test]
@@ -560,11 +608,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("d").in_list(vec![lit(1), lit(2), lit(3)], false).not())?
.build()?;
- let expected =
- "Filter: test.d != Int32(1) AND test.d != Int32(2) AND test.d != Int32(3)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d != Int32(1) AND test.d != Int32(2) AND test.d != Int32(3)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -574,11 +625,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("d").in_list(vec![lit(1), lit(2), lit(3)], true).not())?
.build()?;
- let expected =
- "Filter: test.d = Int32(1) OR test.d = Int32(2) OR test.d = Int32(3)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d = Int32(1) OR test.d = Int32(2) OR test.d = Int32(3)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -589,10 +643,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(qual.not())?
.build()?;
- let expected = "Filter: test.d < Int32(1) OR test.d > Int32(10)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d < Int32(1) OR test.d > Int32(10)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -603,10 +661,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(qual.not())?
.build()?;
- let expected = "Filter: test.d >= Int32(1) AND test.d <= Int32(10)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d >= Int32(1) AND test.d <= Int32(10)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -623,10 +685,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("a").like(col("b")).not())?
.build()?;
- let expected = "Filter: test.a NOT LIKE test.b\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.a NOT LIKE test.b
+ TableScan: test
+ "
+ )
}
#[test]
@@ -643,10 +709,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("a").not_like(col("b")).not())?
.build()?;
- let expected = "Filter: test.a LIKE test.b\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.a LIKE test.b
+ TableScan: test
+ "
+ )
}
#[test]
@@ -663,10 +733,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("a").ilike(col("b")).not())?
.build()?;
- let expected = "Filter: test.a NOT ILIKE test.b\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.a NOT ILIKE test.b
+ TableScan: test
+ "
+ )
}
#[test]
@@ -676,10 +750,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(binary_expr(col("d"), Operator::IsDistinctFrom, lit(10)).not())?
.build()?;
- let expected = "Filter: test.d IS NOT DISTINCT FROM Int32(10)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d IS NOT DISTINCT FROM Int32(10)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -689,10 +767,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(binary_expr(col("d"), Operator::IsNotDistinctFrom, lit(10)).not())?
.build()?;
- let expected = "Filter: test.d IS DISTINCT FROM Int32(10)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.d IS DISTINCT FROM Int32(10)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -714,11 +796,14 @@ mod tests {
// before simplify: t1.a + CAST(Int64(1), UInt32) = t2.a + CAST(Int64(2), UInt32)
// after simplify: t1.a + UInt32(1) = t2.a + UInt32(2) AS t1.a + Int64(1) = t2.a + Int64(2)
- let expected = "Inner Join: t1.a + UInt32(1) = t2.a + UInt32(2)\
- \n TableScan: t1\
- \n TableScan: t2";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Inner Join: t1.a + UInt32(1) = t2.a + UInt32(2)
+ TableScan: t1
+ TableScan: t2
+ "
+ )
}
#[test]
@@ -728,10 +813,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("d").is_not_null())?
.build()?;
- let expected = "Filter: Boolean(true)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: Boolean(true)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -741,10 +830,14 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.filter(col("d").is_null())?
.build()?;
- let expected = "Filter: Boolean(false)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: Boolean(false)
+ TableScan: test
+ "
+ )
}
#[test]
@@ -761,10 +854,13 @@ mod tests {
)?
.build()?;
- let expected = "Aggregate: groupBy=[[GROUPING SETS ((Int32(43) AS age, test.a), (Boolean(false) AS cond), (test.d AS e, Int32(3) AS Int32(1) + Int32(2)))]], aggr=[[]]\
- \n TableScan: test";
-
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Aggregate: groupBy=[[GROUPING SETS ((Int32(43) AS age, test.a), (Boolean(false) AS cond), (test.d AS e, Int32(3) AS Int32(1) + Int32(2)))]], aggr=[[]]
+ TableScan: test
+ "
+ )
}
#[test]
@@ -779,19 +875,27 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan.clone())
.filter(binary_expr(col("a"), Operator::RegexMatch, lit(".*")))?
.build()?;
- let expected = "Filter: test.a IS NOT NULL\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)?;
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: test.a IS NOT NULL
+ TableScan: test
+ "
+ )?;
// Test `!= ".*"` transforms to checking if the column is empty
let plan = LogicalPlanBuilder::from(table_scan.clone())
.filter(binary_expr(col("a"), Operator::RegexNotMatch, lit(".*")))?
.build()?;
- let expected = "Filter: test.a = Utf8(\"\")\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)?;
+ assert_optimized_plan_equal!(
+ plan,
+ @ r#"
+ Filter: test.a = Utf8("")
+ TableScan: test
+ "#
+ )?;
// Test case-insensitive versions
@@ -799,18 +903,26 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan.clone())
.filter(binary_expr(col("b"), Operator::RegexIMatch, lit(".*")))?
.build()?;
- let expected = "Filter: Boolean(true)\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)?;
+ assert_optimized_plan_equal!(
+ plan,
+ @ r"
+ Filter: Boolean(true)
+ TableScan: test
+ "
+ )?;
// Test `!~ ".*"` (case-insensitive) transforms to checking if the column is empty
let plan = LogicalPlanBuilder::from(table_scan.clone())
.filter(binary_expr(col("a"), Operator::RegexNotIMatch, lit(".*")))?
.build()?;
- let expected = "Filter: test.a = Utf8(\"\")\
- \n TableScan: test";
- assert_optimized_plan_eq(plan, expected)
+ assert_optimized_plan_equal!(
+ plan,
+ @ r#"
+ Filter: test.a = Utf8("")
+ TableScan: test
+ "#
+ )
}
}
diff --git a/datafusion/optimizer/src/test/mod.rs b/datafusion/optimizer/src/test/mod.rs
index 5193746021..20e6d2b612 100644
--- a/datafusion/optimizer/src/test/mod.rs
+++ b/datafusion/optimizer/src/test/mod.rs
@@ -99,29 +99,20 @@ pub fn get_tpch_table_schema(table: &str) -> Schema {
}
}
-pub fn assert_analyzed_plan_eq(
- rule: Arc<dyn AnalyzerRule + Send + Sync>,
- plan: LogicalPlan,
- expected: &str,
-) -> Result<()> {
- let options = ConfigOptions::default();
- assert_analyzed_plan_with_config_eq(options, rule, plan, expected)?;
+#[macro_export]
+macro_rules! assert_analyzed_plan_with_config_eq_snapshot {
+ (
+ $options:expr,
+ $rule:expr,
+ $plan:expr,
+ @ $expected:literal $(,)?
+ ) => {{
+ let analyzed_plan = $crate::Analyzer::with_rules(vec![$rule]).execute_and_check($plan, &$options, |_, _| {})?;
- Ok(())
-}
+ insta::assert_snapshot!(analyzed_plan, @ $expected);
-pub fn assert_analyzed_plan_with_config_eq(
- options: ConfigOptions,
- rule: Arc<dyn AnalyzerRule + Send + Sync>,
- plan: LogicalPlan,
- expected: &str,
-) -> Result<()> {
- let analyzed_plan =
- Analyzer::with_rules(vec![rule]).execute_and_check(plan, &options, |_, _| {})?;
- let formatted_plan = format!("{analyzed_plan}");
- assert_eq!(formatted_plan, expected);
-
- Ok(())
+ Ok::<(), datafusion_common::DataFusionError>(())
+ }};
}
pub fn assert_analyzer_check_err(
@@ -144,6 +135,19 @@ fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
#[macro_export]
macro_rules! assert_optimized_plan_eq_snapshot {
+ (
+ $optimizer_context:expr,
+ $rules:expr,
+ $plan:expr,
+ @ $expected:literal $(,)?
+ ) => {{
+ let optimizer = $crate::Optimizer::with_rules($rules);
+ let optimized_plan = optimizer.optimize($plan, &$optimizer_context, |_, _| {})?;
+ insta::assert_snapshot!(optimized_plan, @ $expected);
+
+ Ok::<(), datafusion_common::DataFusionError>(())
+ }};
+
(
$rule:expr,
$plan:expr,
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org