Repository: dask-contrib/dask-expr
Pull request #875 (2024-02-14 11:09:32)

Push filters through shuffles


Patch:
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index bdd821f..a5b6cca 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -421,6 +421,11 @@ class Expr(core.Expr):
     def dtypes(self):
         return self._meta.dtypes
 
+    def _filter_simplification(self, parent, predicate=None):
+        if predicate is None:
+            predicate = parent.predicate.substitute(self, self.frame)
+        return type(self)(self.frame[predicate], *self.operands[1:])
+
 
 class Literal(Expr):
     """Represent a literal (known) value as an `Expr`"""
@@ -1155,11 +1160,6 @@ class Elemwise(Blockwise):
             return self._filter_simplification(parent, predicate)
         return super()._simplify_up(parent, dependents)
 
-    def _filter_simplification(self, parent, predicate=None):
-        if predicate is None:
-            predicate = parent.predicate.substitute(self, self.frame)
-        return type(self)(self.frame[predicate], *self.operands[1:])
-
 
 class RenameFrame(Elemwise):
     _parameters = ["frame", "columns"]
diff --git a/dask_expr/_shuffle.py b/dask_expr/_shuffle.py
index e34eee6..9586cb3 100644
--- a/dask_expr/_shuffle.py
+++ b/dask_expr/_shuffle.py
@@ -33,6 +33,7 @@ from dask_expr._expr import (
     Assign,
     Blockwise,
     Expr,
+    Filter,
     PartitionsFiltered,
     Projection,
     ToSeriesIndex,
@@ -83,6 +84,7 @@ class ShuffleBase(Expr):
         "index_shuffle": None,
     }
     _is_length_preserving = True
+    _filter_passthrough = True
 
     def __str__(self):
         return f"Shuffle({self._name[-7:]})"
@@ -91,6 +93,10 @@ class ShuffleBase(Expr):
         return [self.frame, self.partitioning_index]
 
     def _simplify_up(self, parent, dependents):
+        if isinstance(parent, Filter) and self._filter_passthrough_available(
+            parent, dependents
+        ):
+            return self._filter_simplification(parent)
         if isinstance(parent, Projection):
             # Move the column projection to come
             # before the abstract Shuffle


Test patch:
diff --git a/dask_expr/tests/test_shuffle.py b/dask_expr/tests/test_shuffle.py
index 63afeca..2ce8bf5 100644
--- a/dask_expr/tests/test_shuffle.py
+++ b/dask_expr/tests/test_shuffle.py
@@ -695,3 +695,24 @@ def test_shuffle_no_assign(df, pdf):
     result = df.shuffle(df.x)
     q = result.optimize(fuse=False)
     assert len([x for x in q.walk() if isinstance(x, Assign)]) == 0
+
+
+def test_shuffle_filter_pushdown(pdf):
+    pdf["z"] = 1
+    df = from_pandas(pdf, npartitions=10)
+    result = df.shuffle("x")
+    result = result[result.x > 5.0]
+    expected = df[df.x > 5.0].shuffle("x")
+    assert result.simplify()._name == expected._name
+
+    result = df.shuffle("x")
+    result = result[result.x > 5.0][["x", "y"]]
+    expected = df[["x", "y"]]
+    expected = expected[expected.x > 5.0].shuffle("x")
+    assert result.simplify()._name == expected.simplify()._name
+
+    result = df.shuffle("x")[["x", "y"]]
+    result = result[result.x > 5.0]
+    expected = df[["x", "y"]]
+    expected = expected[expected.x > 5.0].shuffle("x")
+    assert result.simplify()._name == expected.simplify()._name