fix[rust]: ensure all predicates use same key function when inserting…

… in hashmap (#5034)
pola-rs · Sep 30, 2022 · 0e6bff2 · 0e6bff2
1 parent c3fe475
commit 0e6bff2
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 50 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/predicate_pushdown/mod.rs b/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/predicate_pushdown/mod.rs
@@ -7,7 +7,7 @@ use utils::*;
 use super::*;
 use crate::dsl::function_expr::FunctionExpr;
 use crate::logical_plan::{optimizer, Context};
-use crate::utils::{aexpr_to_leaf_names, aexprs_to_schema, check_input_node, has_aexpr};
+use crate::utils::{aexprs_to_schema, check_input_node, has_aexpr};
 
 #[derive(Default)]
 pub struct PredicatePushDown {}
@@ -91,12 +91,11 @@ impl PredicatePushDown {
                     let input_schema = lp_arena.get(node).schema(lp_arena);
                     let mut pushdown_predicates =
                         optimizer::init_hashmap(Some(acc_predicates.len()));
-                    for (name, &predicate) in acc_predicates.iter() {
+                    for (_, &predicate) in acc_predicates.iter() {
                         // we can pushdown the predicate
                         if check_input_node(predicate, &input_schema, expr_arena) {
                             insert_and_combine_predicate(
                                 &mut pushdown_predicates,
-                                name.clone(),
                                 predicate,
                                 expr_arena,
                             )
@@ -157,7 +156,7 @@ impl PredicatePushDown {
     ///
     /// * `AlogicalPlan` - Arena based logical plan tree representing the query.
     /// * `acc_predicates` - The predicates we accumulate during tree traversal.
-    ///                      The hashmap maps from root-column name to predicates on that column.
+    ///                      The hashmap maps from leaf-column name to predicates on that column.
     ///                      If the key is already taken we combine the predicate with a bitand operation.
     ///                      The `Node`s are indexes in the `expr_arena`
     /// * `lp_arena` - The local memory arena for the logical plan.
@@ -178,8 +177,7 @@ impl PredicatePushDown {
                 // we remove it and apply it locally
                 let local_predicates = transfer_to_local_by_node(&mut acc_predicates, |node| predicate_is_pushdown_boundary(node, expr_arena));
 
-                let name = roots_to_key(&aexpr_to_leaf_names(predicate, expr_arena));
-                insert_and_combine_predicate(&mut acc_predicates, name, predicate, expr_arena);
+                insert_and_combine_predicate(&mut acc_predicates, predicate, expr_arena);
                 let alp = lp_arena.take(input);
                 let new_input = self.push_down(alp, acc_predicates, lp_arena, expr_arena)?;
 
@@ -451,12 +449,9 @@ impl PredicatePushDown {
                     // be influenced by join
                     #[allow(clippy::suspicious_else_formatting)]
                     if !predicate_is_pushdown_boundary(predicate, expr_arena) {
-                        // no else if. predicate can be in both tables.
                         if check_input_node(predicate, &schema_left, expr_arena) {
-                            let name = get_insertion_name(expr_arena, predicate, &schema_left);
                             insert_and_combine_predicate(
                                 &mut pushdown_left,
-                                name,
                                 predicate,
                                 expr_arena,
                             );
@@ -467,10 +462,8 @@ impl PredicatePushDown {
                         // in that case we should not push down as the user wants to filter on `x`
                         // not on `x_rhs`.
                         else if check_input_node(predicate, &schema_right, expr_arena)  {
-                            let name = get_insertion_name(expr_arena, predicate, &schema_right);
                             insert_and_combine_predicate(
                                 &mut pushdown_right,
-                                name,
                                 predicate,
                                 expr_arena,
                             );
@@ -577,12 +570,7 @@ mod test {
 
         let predicate_expr = col("foo").gt(col("bar"));
         let predicate = to_aexpr(predicate_expr.clone(), &mut expr_arena);
-        insert_and_combine_predicate(
-            &mut acc_predicates,
-            Arc::from("foo"),
-            predicate,
-            &mut expr_arena,
-        );
+        insert_and_combine_predicate(&mut acc_predicates, predicate, &mut expr_arena);
         let root = *acc_predicates.get("foo").unwrap();
         let expr = node_to_expr(root, &expr_arena);
         assert_eq!(format!("{:?}", &expr), format!("{:?}", predicate_expr));

diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/predicate_pushdown/utils.rs b/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/predicate_pushdown/utils.rs
@@ -22,10 +22,11 @@ impl Dsl for Node {
 /// Don't overwrite predicates but combine them.
 pub(super) fn insert_and_combine_predicate(
     acc_predicates: &mut PlHashMap<Arc<str>, Node>,
-    name: Arc<str>,
     predicate: Node,
     arena: &mut Arena<AExpr>,
 ) {
+    let name = predicate_to_key(predicate, arena);
+
     acc_predicates
         .entry(name)
         .and_modify(|existing_predicate| {
@@ -77,37 +78,30 @@ pub(super) fn predicate_at_scan(
 // an invisible ascii token we use as delimiter
 const HIDDEN_DELIMITER: char = '\u{1D17A}';
 
-/// Determine the hashmap key by combining all the root column names of a predicate
-pub(super) fn roots_to_key(roots: &[Arc<str>]) -> Arc<str> {
-    if roots.len() == 1 {
-        roots[0].clone()
-    } else {
-        let mut new = String::with_capacity(32 * roots.len());
-        for (i, name) in roots.iter().enumerate() {
-            if i > 0 {
-                new.push(HIDDEN_DELIMITER)
+/// Determine the hashmap key by combining all the leaf column names of a predicate
+pub(super) fn predicate_to_key(predicate: Node, expr_arena: &Arena<AExpr>) -> Arc<str> {
+    let mut iter = aexpr_to_leaf_names_iter(predicate, expr_arena);
+    if let Some(first) = iter.next() {
+        if let Some(second) = iter.next() {
+            let mut new = String::with_capacity(32 * iter.size_hint().0);
+            new.push_str(&first);
+            new.push(HIDDEN_DELIMITER);
+            new.push_str(&second);
+
+            for name in iter {
+                new.push(HIDDEN_DELIMITER);
+                new.push_str(&name);
             }
-            new.push_str(name);
+            return Arc::from(new);
         }
-        Arc::from(new)
+        first
+    } else {
+        let mut s = String::new();
+        s.push(HIDDEN_DELIMITER);
+        Arc::from(s)
     }
 }
 
-pub(super) fn get_insertion_name(
-    expr_arena: &Arena<AExpr>,
-    predicate: Node,
-    schema: &Schema,
-) -> Arc<str> {
-    Arc::from(
-        expr_arena
-            .get(predicate)
-            .to_field(schema, Context::Default, expr_arena)
-            .unwrap()
-            .name()
-            .as_ref(),
-    )
-}
-
 // this checks if a predicate from a node upstream can pass
 // the predicate in this filter
 // Cases where this cannot be the case:
@@ -255,12 +249,7 @@ where
                             projection_roots[0].clone(),
                         );
 
-                        insert_and_combine_predicate(
-                            acc_predicates,
-                            projection_roots[0].clone(),
-                            predicate,
-                            expr_arena,
-                        );
+                        insert_and_combine_predicate(acc_predicates, predicate, expr_arena);
                     } else {
                         // this may be a complex binary function. The predicate may only be valid
                         // on this projected column so we do filter locally.