Skip to content

Commit

Permalink
Fix bug in simplify expressions (#214) (#7699)
Browse files Browse the repository at this point in the history
Co-authored-by: Dan Harris <[email protected]>
  • Loading branch information
Dandandan and thinkharderdev committed Sep 29, 2023
1 parent c6d29e7 commit 2abacf4
Showing 1 changed file with 48 additions and 3 deletions.
51 changes: 48 additions & 3 deletions datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,10 +65,13 @@ impl SimplifyExpressions {
) -> Result<LogicalPlan> {
let schema = if !plan.inputs().is_empty() {
DFSchemaRef::new(merge_schema(plan.inputs()))
} else if let LogicalPlan::TableScan(_) = plan {
} else if let LogicalPlan::TableScan(scan) = plan {
// When predicates are pushed into a table scan, there needs to be
// a schema to resolve the fields against.
Arc::clone(plan.schema())
Arc::new(DFSchema::try_from_qualified_schema(
&scan.table_name,
&scan.source.schema(),
)?)
} else {
Arc::new(DFSchema::empty())
};
Expand Down Expand Up @@ -111,7 +114,7 @@ mod tests {
use crate::simplify_expressions::utils::for_test::{
cast_to_int64_expr, now_expr, to_timestamp_expr,
};
use crate::test::test_table_scan_with_name;
use crate::test::{assert_fields_eq, test_table_scan_with_name};

use super::*;
use arrow::datatypes::{DataType, Field, Schema};
Expand Down Expand Up @@ -174,6 +177,48 @@ mod tests {
Ok(())
}

#[test]
fn test_simplify_table_full_filter_in_scan() -> Result<()> {
let fields = vec![
Field::new("a", DataType::UInt32, false),
Field::new("b", DataType::UInt32, false),
Field::new("c", DataType::UInt32, false),
];

let schema = Schema::new(fields);

let table_scan = table_scan_with_filters(
Some("test"),
&schema,
Some(vec![0]),
vec![col("b").is_not_null()],
)?
.build()?;
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) AS b IS NOT NULL]";

assert_optimized_plan_eq(&table_scan, expected)
}

#[test]
fn test_simplify_filter_pushdown() -> Result<()> {
let table_scan = test_table_scan();
let plan = LogicalPlanBuilder::from(table_scan)
.project(vec![col("a")])?
.filter(and(col("b").gt(lit(1)), col("b").gt(lit(1))))?
.build()?;

assert_optimized_plan_eq(
&plan,
"\
Filter: test.b > Int32(1)\
\n Projection: test.a\
\n TableScan: test",
)
}

#[test]
fn test_simplify_optimized_plan() -> Result<()> {
let table_scan = test_table_scan();
Expand Down

0 comments on commit 2abacf4

Please sign in to comment.