Repository: dask-contrib/dask-expr
Pull request #888 (2024-02-21 11:45:47)

Fix merge predicate pushdowns with weird predicates
I thought a little about how we do predicate pushdowns and stumbled upon this, we currently don't cover weird predicates very well, we have to walk the predicate until we arrive at the merge to determine from which side they are from


Patch:
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index 3c911f9..07788ec 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -1172,6 +1172,7 @@ class RenameFrame(Elemwise):
             reverse_mapping = {val: key for key, val in self.operand("columns").items()}
 
             columns = determine_column_projection(self, parent, dependents)
+            columns = _convert_to_list(columns)
             columns = [
                 reverse_mapping[col] if col in reverse_mapping else col
                 for col in columns
diff --git a/dask_expr/_merge.py b/dask_expr/_merge.py
index a7d3f4d..3e47e10 100644
--- a/dask_expr/_merge.py
+++ b/dask_expr/_merge.py
@@ -26,6 +26,8 @@ from dask_expr._expr import (  # noqa: F401
     PartitionsFiltered,
     Projection,
     Unaryop,
+    _DelayedExpr,
+    are_co_aligned,
     determine_column_projection,
     is_filter_pushdown_available,
 )
@@ -108,6 +110,45 @@ class Merge(Expr):
             return True
         return False
 
+    def _predicate_columns(self, predicate):
+        if isinstance(predicate, (Projection, Unaryop, Isin)):
+            return self._get_original_predicate_columns(predicate)
+        elif isinstance(predicate, Binop):
+            if isinstance(predicate, And):
+                return None
+
+            if not isinstance(predicate.right, Expr):
+                return self._get_original_predicate_columns(predicate.left)
+            elif isinstance(predicate.right, Elemwise):
+                return self._get_original_predicate_columns(predicate)
+            else:
+                return None
+        else:
+            # Unsupported predicate type
+            return None
+
+    def _get_original_predicate_columns(self, predicate):
+        predicate_columns = set()
+        stack = [predicate]
+        seen = set()
+        while stack:
+            e = stack.pop()
+            if self._name == e._name:
+                continue
+
+            if e._name in seen:
+                continue
+            seen.add(e._name)
+
+            if isinstance(e, _DelayedExpr):
+                continue
+
+            dependencies = e.dependencies()
+            stack.extend(dependencies)
+            if any(d._name == self._name for d in dependencies):
+                predicate_columns.update(e.columns)
+        return predicate_columns
+
     def __str__(self):
         return f"Merge({self._name[-7:]})"
 
@@ -388,23 +429,6 @@ class Merge(Expr):
         # Blockwise merge
         return BlockwiseMerge(left, right, **self.kwargs)
 
-    def _predicate_columns(self, predicate):
-        if isinstance(predicate, (Projection, Unaryop, Isin)):
-            return set(predicate.columns)
-        elif isinstance(predicate, Binop):
-            if isinstance(predicate, And):
-                return None
-
-            if not isinstance(predicate.right, Expr):
-                return set(predicate.left.columns)
-            elif isinstance(predicate.right, Elemwise):
-                return set(predicate.left.columns) | set(predicate.right.columns)
-            else:
-                return None
-        else:
-            # Unsupported predicate type
-            return None
-
     def _simplify_up(self, parent, dependents):
         if isinstance(parent, Filter):
             if not self._filter_passthrough_available(parent, dependents):


Test patch:
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index 9edb0d4..8ae6674 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -1927,6 +1927,12 @@ def test_assign_squash_together(df, pdf):
     assert "Assign: a=1, b=2" in [line for line in result.expr._tree_repr_lines()]
 
 
+def test_rename_series_reduction(pdf, df):
+    result = df[["x"]].rename(columns={"x": "a"})["a"]
+    expected = pdf[["x"]].rename(columns={"x": "a"})["a"]
+    assert_eq(result, expected)
+
+
 def test_are_co_aligned(pdf, df):
     df2 = df.reset_index()
     assert are_co_aligned(df.expr, df2.expr)
diff --git a/dask_expr/tests/test_merge.py b/dask_expr/tests/test_merge.py
index 2a86d3d..d228df3 100644
--- a/dask_expr/tests/test_merge.py
+++ b/dask_expr/tests/test_merge.py
@@ -7,6 +7,7 @@ from dask_expr import Merge, from_pandas, merge, repartition
 from dask_expr._expr import Filter, Projection
 from dask_expr._merge import BroadcastJoin
 from dask_expr._shuffle import Shuffle
+from dask_expr.io import FromPandas
 from dask_expr.tests._util import _backend_library, assert_eq
 
 # Set DataFrame backend for this module
@@ -796,6 +797,26 @@ def test_filter_merge_suffixes():
     assert_eq(result, expected[expected.b_left < 2], check_index=False)
 
 
+def test_merge_filter_renaming_columns():
+    pdf1 = pd.DataFrame({"a": [1, 2, 3, 4], "left": 4})
+    df1 = from_pandas(pdf1, npartitions=2)
+
+    pdf2 = pd.DataFrame({"a": [1, 2, 3, 4], "right": [6, 6, 5, 6]})
+    df2 = from_pandas(pdf2, npartitions=2)
+
+    result = df1.merge(df2)
+    result = result[result[["right"]].rename(columns={"right": "left"})["left"] == 6]
+    expected = pdf1.merge(pdf2)
+    expected = expected[
+        expected[["right"]].rename(columns={"right": "left"})["left"] == 6
+    ]
+    assert_eq(result, expected, check_index=False)
+
+    q = result.simplify()
+    assert isinstance(q.expr.right, Filter)
+    assert isinstance(q.expr.left, FromPandas)
+
+
 def test_merge_avoid_overeager_filter_pushdown():
     df = pd.DataFrame({"a": [1, 2, 3], "b": 1})
     ddf = from_pandas(df, npartitions=2)