fix[rust]: block predicate on join if already in left table (#4791)

pola-rs · Sep 9, 2022 · 5dae949 · 5dae949
1 parent bc24180
commit 5dae949
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 3 deletions.
diff --git a/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs b/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs
@@ -447,6 +447,7 @@ impl PredicatePushDown {
 
                     // predicate should not have an aggregation or window function as that would
                     // 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) {
@@ -459,8 +460,11 @@ impl PredicatePushDown {
                             );
                             filter_left = true;
                         }
-
-                        if check_input_node(predicate, &schema_right, expr_arena)  {
+                        // this is `else if` because if the predicate is in the left hand side
+                        // the right hand side should be renamed with the suffix.
+                        // 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,
@@ -472,7 +476,7 @@ impl PredicatePushDown {
                         }
                     }
                     match (filter_left, filter_right, &options.how) {
-                        // if not pushed down on of the tables we have to do it locally.
+                        // if not pushed down on one of the tables we have to do it locally.
                         (false, false, _) |
                         // if left join and predicate only available in right table,
                         // 'we should not filter right, because that would lead to

diff --git a/polars/polars-lazy/src/tests/predicate_queries.rs b/polars/polars-lazy/src/tests/predicate_queries.rs
@@ -237,3 +237,30 @@ fn test_predicate_pd_apply() -> Result<()> {
     assert!(predicate_at_scan(q.clone()));
     Ok(())
 }
+#[test]
+fn test_predicate_on_join_suffix_4788() -> Result<()> {
+    let lf = df![
+      "x" => [1, 2],
+      "y" => [1, 1],
+    ]?
+    .lazy();
+
+    let q = (lf.clone().join_builder().with(lf))
+        .left_on([col("y")])
+        .right_on([col("y")])
+        .suffix("_")
+        .finish()
+        .filter(col("x").eq(1));
+
+    // the left hand side should have a predicate
+    assert!(predicate_at_scan(q.clone()));
+
+    let expected = df![
+        "x" => [1, 1],
+        "y" => [1, 1],
+        "x_" => [1, 2],
+    ]?;
+    assert_eq!(q.collect()?, expected);
+
+    Ok(())
+}