Skip to content

Commit f9231fc

Browse files
pepijnvealamb
andauthored
Move GuaranteeRewriter to datafusion_expr (#18821)
## Which issue does this PR close? - None, break out PR of changes done in #17813 ## Rationale for this change In #17813 `GuaranteeRewriter` is used from the `datafusion_expr` crate. In order to enable this the type needed to be moved from `datafusion_optimizer` to `datafusion_expr`. Additionally, during the development of #17813 some latent bugs were discovered in `GuaranteeRewriter` that have been resolved. ## What changes are included in this PR? - Move `GuaranteeRewriter` to `datafusion_expr` - Fix two bugs where rewrites of 'between' expression would fail - when one of the bounds was untyped null - when the lower bound was greater than the upper bound - Add logic to replace expressions with literal null based on provided guarantees - Split implementation into smaller functions for easier readability ## Are these changes tested? - Existing tests updated - Tests added for bugfixes ## Are there any user-facing changes? No --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent d01c0d3 commit f9231fc

File tree

6 files changed

+696
-492
lines changed

6 files changed

+696
-492
lines changed

datafusion/core/tests/optimizer/mod.rs

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ use arrow::datatypes::{
2727
DataType, Field, Fields, Schema, SchemaBuilder, SchemaRef, TimeUnit,
2828
};
2929
use datafusion_common::config::ConfigOptions;
30-
use datafusion_common::tree_node::{TransformedResult, TreeNode};
30+
use datafusion_common::tree_node::TransformedResult;
3131
use datafusion_common::{plan_err, DFSchema, Result, ScalarValue, TableReference};
3232
use datafusion_expr::interval_arithmetic::{Interval, NullableInterval};
3333
use datafusion_expr::{
@@ -37,14 +37,14 @@ use datafusion_expr::{
3737
use datafusion_functions::core::expr_ext::FieldAccessor;
3838
use datafusion_optimizer::analyzer::Analyzer;
3939
use datafusion_optimizer::optimizer::Optimizer;
40-
use datafusion_optimizer::simplify_expressions::GuaranteeRewriter;
4140
use datafusion_optimizer::{OptimizerConfig, OptimizerContext};
4241
use datafusion_sql::planner::{ContextProvider, SqlToRel};
4342
use datafusion_sql::sqlparser::ast::Statement;
4443
use datafusion_sql::sqlparser::dialect::GenericDialect;
4544
use datafusion_sql::sqlparser::parser::Parser;
4645

4746
use chrono::DateTime;
47+
use datafusion_expr::expr_rewriter::rewrite_with_guarantees;
4848
use datafusion_functions::datetime;
4949

5050
#[cfg(test)]
@@ -304,8 +304,6 @@ fn test_inequalities_non_null_bounded() {
304304
),
305305
];
306306

307-
let mut rewriter = GuaranteeRewriter::new(guarantees.iter());
308-
309307
// (original_expr, expected_simplification)
310308
let simplified_cases = &[
311309
(col("x").lt(lit(0)), false),
@@ -337,7 +335,7 @@ fn test_inequalities_non_null_bounded() {
337335
),
338336
];
339337

340-
validate_simplified_cases(&mut rewriter, simplified_cases);
338+
validate_simplified_cases(&guarantees, simplified_cases);
341339

342340
let unchanged_cases = &[
343341
col("x").gt(lit(2)),
@@ -348,26 +346,32 @@ fn test_inequalities_non_null_bounded() {
348346
col("x").not_between(lit(3), lit(10)),
349347
];
350348

351-
validate_unchanged_cases(&mut rewriter, unchanged_cases);
349+
validate_unchanged_cases(&guarantees, unchanged_cases);
352350
}
353351

354-
fn validate_simplified_cases<T>(rewriter: &mut GuaranteeRewriter, cases: &[(Expr, T)])
355-
where
352+
fn validate_simplified_cases<T>(
353+
guarantees: &[(Expr, NullableInterval)],
354+
cases: &[(Expr, T)],
355+
) where
356356
ScalarValue: From<T>,
357357
T: Clone,
358358
{
359359
for (expr, expected_value) in cases {
360-
let output = expr.clone().rewrite(rewriter).data().unwrap();
360+
let output = rewrite_with_guarantees(expr.clone(), guarantees)
361+
.data()
362+
.unwrap();
361363
let expected = lit(ScalarValue::from(expected_value.clone()));
362364
assert_eq!(
363365
output, expected,
364366
"{expr} simplified to {output}, but expected {expected}"
365367
);
366368
}
367369
}
368-
fn validate_unchanged_cases(rewriter: &mut GuaranteeRewriter, cases: &[Expr]) {
370+
fn validate_unchanged_cases(guarantees: &[(Expr, NullableInterval)], cases: &[Expr]) {
369371
for expr in cases {
370-
let output = expr.clone().rewrite(rewriter).data().unwrap();
372+
let output = rewrite_with_guarantees(expr.clone(), guarantees)
373+
.data()
374+
.unwrap();
371375
assert_eq!(
372376
&output, expr,
373377
"{expr} was simplified to {output}, but expected it to be unchanged"

0 commit comments

Comments
 (0)