Skip to content

Commit

Permalink
feat(rust, python)!: error on string <-> date cmp (#6498)
Browse files Browse the repository at this point in the history
  • Loading branch information
@ritchie46
ritchie46 authored Jan 27, 2023
1 parent 5047917 commit 7950c3c
Show file tree
Hide file tree
Showing 9 changed files with 153 additions and 92 deletions.
4 changes: 2 additions & 2 deletions polars/polars-lazy/polars-plan/src/logical_plan/optimizer/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ pub fn optimize(
// we must clean up the predicates, because the agg_scan_projection
// uses them in the hashtable to determine duplicates.
let simplify_bools = &mut [Box::new(SimplifyBooleanRule {}) as Box<dyn OptimizationRule>];
lp_top = opt.optimize_loop(simplify_bools, expr_arena, lp_arena, lp_top);
lp_top = opt.optimize_loop(simplify_bools, expr_arena, lp_arena, lp_top)?;

// scan the LP to aggregate all the column used in scans
// these columns will be added to the state of the AggScanProjection rule
Expand All @@ -176,7 +176,7 @@ pub fn optimize(

rules.push(Box::new(ReplaceDropNulls {}));

lp_top = opt.optimize_loop(&mut rules, expr_arena, lp_arena, lp_top);
lp_top = opt.optimize_loop(&mut rules, expr_arena, lp_arena, lp_top)?;

// during debug we check if the optimizations have not modified the final schema
#[cfg(debug_assertions)]
Expand Down
79 changes: 43 additions & 36 deletions polars/polars-lazy/polars-plan/src/logical_plan/optimizer/simplify_expr.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use crate::prelude::function_expr::FunctionExpr;
macro_rules! eval_binary_same_type {
($lhs:expr, $operand: tt, $rhs:expr) => {{
if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) = ($lhs, $rhs) {
return match (lit_left, lit_right) {
match (lit_left, lit_right) {
(LiteralValue::Float32(x), LiteralValue::Float32(y)) => {
Some(AExpr::Literal(LiteralValue::Float32(x $operand y)))
}
Expand Down Expand Up @@ -46,17 +46,18 @@ macro_rules! eval_binary_same_type {
Some(AExpr::Literal(LiteralValue::UInt64(x $operand y)))
}
_ => None,
};
}
} else {
None
}
None

}}
}

macro_rules! eval_binary_bool_type {
($lhs:expr, $operand: tt, $rhs:expr) => {{
if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) = ($lhs, $rhs) {
return match (lit_left, lit_right) {
match (lit_left, lit_right) {
(LiteralValue::Float32(x), LiteralValue::Float32(y)) => {
Some(AExpr::Literal(LiteralValue::Boolean(x $operand y)))
}
Expand Down Expand Up @@ -95,9 +96,10 @@ macro_rules! eval_binary_bool_type {
Some(AExpr::Literal(LiteralValue::Boolean(x $operand y)))
}
_ => None,
};
}
} else {
None
}
None

}}
}
Expand All @@ -111,9 +113,9 @@ impl OptimizationRule for SimplifyBooleanRule {
expr_node: Node,
_: &Arena<ALogicalPlan>,
_: Node,
) -> Option<AExpr> {
) -> PolarsResult<Option<AExpr>> {
let expr = expr_arena.get(expr_node);
match expr {
let out = match expr {
// true AND x => x
AExpr::BinaryExpr {
left,
Expand Down Expand Up @@ -273,7 +275,8 @@ impl OptimizationRule for SimplifyBooleanRule {
}
}
_ => None,
}
};
Ok(out)
}
}

Expand Down Expand Up @@ -436,10 +439,10 @@ impl OptimizationRule for SimplifyExprRule {
expr_node: Node,
_lp_arena: &Arena<ALogicalPlan>,
_lp_node: Node,
) -> Option<AExpr> {
) -> PolarsResult<Option<AExpr>> {
let expr = expr_arena.get(expr_node);

match expr {
let out = match expr {
// lit(left) + lit(right) => lit(left + right)
// and null propagation
AExpr::BinaryExpr { left, op, right } => {
Expand All @@ -449,36 +452,38 @@ impl OptimizationRule for SimplifyExprRule {
// lit(left) + lit(right) => lit(left + right)
#[allow(clippy::manual_map)]
let out = match op {
Operator::Plus => match eval_binary_same_type!(left_aexpr, +, right_aexpr) {
Some(new) => Some(new),
None => {
// try to replace addition of string columns with `concat_str`
#[cfg(all(feature = "strings", feature = "concat_str"))]
{
string_addition_to_linear_concat(
_lp_arena,
_lp_node,
expr_arena,
*left,
*right,
left_aexpr,
right_aexpr,
)
}
#[cfg(not(all(feature = "strings", feature = "concat_str")))]
{
None
Operator::Plus => {
match dbg!(eval_binary_same_type!(left_aexpr, +, right_aexpr)) {
Some(new) => Some(new),
None => {
// try to replace addition of string columns with `concat_str`
#[cfg(all(feature = "strings", feature = "concat_str"))]
{
string_addition_to_linear_concat(
_lp_arena,
_lp_node,
expr_arena,
*left,
*right,
left_aexpr,
right_aexpr,
)
}
#[cfg(not(all(feature = "strings", feature = "concat_str")))]
{
None
}
}
}
},
}
Operator::Minus => eval_binary_same_type!(left_aexpr, -, right_aexpr),
Operator::Multiply => eval_binary_same_type!(left_aexpr, *, right_aexpr),
Operator::Divide => eval_binary_same_type!(left_aexpr, /, right_aexpr),
Operator::TrueDivide => {
if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) =
(left_aexpr, right_aexpr)
{
return match (lit_left, lit_right) {
match (lit_left, lit_right) {
(LiteralValue::Float32(x), LiteralValue::Float32(y)) => {
Some(AExpr::Literal(LiteralValue::Float32(x / y)))
}
Expand Down Expand Up @@ -514,9 +519,10 @@ impl OptimizationRule for SimplifyExprRule {
AExpr::Literal(LiteralValue::Float64(*x as f64 / *y as f64)),
),
_ => None,
};
}
} else {
None
}
None
}
Operator::FloorDivide => None,
Operator::Modulus => eval_binary_same_type!(left_aexpr, %, right_aexpr),
Expand All @@ -531,7 +537,7 @@ impl OptimizationRule for SimplifyExprRule {
Operator::Xor => eval_bitwise(left_aexpr, right_aexpr, |l, r| l ^ r),
};
if out.is_some() {
return out;
return Ok(out);
}

// Null propagation.
Expand Down Expand Up @@ -644,7 +650,8 @@ impl OptimizationRule for SimplifyExprRule {
}

_ => None,
}
};
Ok(out)
}
}

Expand Down
9 changes: 5 additions & 4 deletions polars/polars-lazy/polars-plan/src/logical_plan/optimizer/slice_pushdown_expr.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ impl OptimizationRule for SlicePushDown {
expr_node: Node,
_lp_arena: &Arena<ALogicalPlan>,
_lp_node: Node,
) -> Option<AExpr> {
) -> PolarsResult<Option<AExpr>> {
if let AExpr::Slice {
input,
offset,
Expand All @@ -26,7 +26,7 @@ impl OptimizationRule for SlicePushDown {
let length = *length;

use AExpr::*;
match expr_arena.get(*input) {
let out = match expr_arena.get(*input) {
m @ Alias(..) | m @ Cast { .. } => {
let m = m.clone();
let input = m.get_input();
Expand Down Expand Up @@ -94,9 +94,10 @@ impl OptimizationRule for SlicePushDown {
}
}
_ => None,
}
};
Ok(out)
} else {
None
Ok(None)
}
}
}
12 changes: 7 additions & 5 deletions polars/polars-lazy/polars-plan/src/logical_plan/optimizer/stack_opt.rs
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use polars_core::prelude::PolarsResult;

use crate::logical_plan::aexpr::AExpr;
use crate::logical_plan::alp::ALogicalPlan;
use crate::prelude::{Arena, Node};
Expand All @@ -12,7 +14,7 @@ impl StackOptimizer {
expr_arena: &mut Arena<AExpr>,
lp_arena: &mut Arena<ALogicalPlan>,
lp_top: Node,
) -> Node {
) -> PolarsResult<Node> {
let mut changed = true;

let mut plans = Vec::with_capacity(32);
Expand Down Expand Up @@ -50,7 +52,7 @@ impl StackOptimizer {
current_expr_node,
lp_arena,
current_node,
) {
)? {
expr_arena.replace(current_expr_node, x);
changed = true;
}
Expand All @@ -62,7 +64,7 @@ impl StackOptimizer {
}
}
}
lp_top
Ok(lp_top)
}
}

Expand All @@ -86,7 +88,7 @@ pub trait OptimizationRule {
_expr_node: Node,
_lp_arena: &Arena<ALogicalPlan>,
_lp_node: Node,
) -> Option<AExpr> {
None
) -> PolarsResult<Option<AExpr>> {
Ok(None)
}
}
Loading

0 comments on commit 7950c3c

Please sign in to comment.