Skip to content

Commit 168f63f

Browse files
committed
Rewrite SimplifyExprs to avoid a Plan copy
fixup
1 parent 0bccb70 commit 168f63f

File tree

2 files changed

+98
-78
lines changed

2 files changed

+98
-78
lines changed

datafusion/core/tests/simplification.rs

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ use datafusion_expr::{
3030
expr, table_scan, BuiltinScalarFunction, Cast, ColumnarValue, Expr, ExprSchemable,
3131
LogicalPlan, LogicalPlanBuilder, ScalarUDF, Volatility,
3232
};
33+
use datafusion_optimizer::optimizer::Optimizer;
3334
use datafusion_optimizer::simplify_expressions::{ExprSimplifier, SimplifyExpressions};
3435
use datafusion_optimizer::{OptimizerContext, OptimizerRule};
3536
use std::sync::Arc;
@@ -107,14 +108,14 @@ fn test_table_scan() -> LogicalPlan {
107108
.expect("building plan")
108109
}
109110

110-
fn get_optimized_plan_formatted(plan: &LogicalPlan, date_time: &DateTime<Utc>) -> String {
111+
fn get_optimized_plan_formatted(plan: LogicalPlan, date_time: &DateTime<Utc>) -> String {
111112
let config = OptimizerContext::new().with_query_execution_start_time(*date_time);
112-
let rule = SimplifyExpressions::new();
113113

114-
let optimized_plan = rule
115-
.try_optimize(plan, &config)
116-
.unwrap()
117-
.expect("failed to optimize plan");
114+
// Use Optimizer to do plan traversal
115+
fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
116+
let optimizer = Optimizer::with_rules(vec![Arc::new(SimplifyExpressions::new())]);
117+
let optimized_plan = optimizer.optimize(plan, &config, observe).unwrap();
118+
118119
format!("{optimized_plan:?}")
119120
}
120121

@@ -236,7 +237,7 @@ fn to_timestamp_expr_folded() -> Result<()> {
236237
let expected = "Projection: TimestampNanosecond(1599566400000000000, None) AS to_timestamp(Utf8(\"2020-09-08T12:00:00+00:00\"))\
237238
\n TableScan: test"
238239
.to_string();
239-
let actual = get_optimized_plan_formatted(&plan, &Utc::now());
240+
let actual = get_optimized_plan_formatted(plan, &Utc::now());
240241
assert_eq!(expected, actual);
241242
Ok(())
242243
}
@@ -260,7 +261,7 @@ fn now_less_than_timestamp() -> Result<()> {
260261
// expression down to a single constant (true)
261262
let expected = "Filter: Boolean(true)\
262263
\n TableScan: test";
263-
let actual = get_optimized_plan_formatted(&plan, &time);
264+
let actual = get_optimized_plan_formatted(plan, &time);
264265

265266
assert_eq!(expected, actual);
266267
Ok(())
@@ -288,7 +289,7 @@ fn select_date_plus_interval() -> Result<()> {
288289
// expression down to a single constant (true)
289290
let expected = r#"Projection: Date32("18636") AS to_timestamp(Utf8("2020-09-08T12:05:00+00:00")) + IntervalDayTime("528280977408")
290291
TableScan: test"#;
291-
let actual = get_optimized_plan_formatted(&plan, &time);
292+
let actual = get_optimized_plan_formatted(plan, &time);
292293

293294
assert_eq!(expected, actual);
294295
Ok(())
@@ -420,7 +421,7 @@ fn multiple_now() -> Result<()> {
420421
.build()?;
421422

422423
// expect the same timestamp appears in both exprs
423-
let actual = get_optimized_plan_formatted(&plan, &time);
424+
let actual = get_optimized_plan_formatted(plan, &time);
424425
let expected = format!(
425426
"Projection: TimestampNanosecond({}, Some(\"+00:00\")) AS now(), TimestampNanosecond({}, Some(\"+00:00\")) AS t2\
426427
\n TableScan: test",

0 commit comments

Comments
 (0)