Repository: dask-contrib/dask-expr
Pull request #876 (2024-02-14 14:24:08)

Push filters through set_index, sort_values and repartition


Patch:
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index 9d7535d..655d5a1 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -108,12 +108,9 @@ class Expr(core.Expr):
 
     def __getitem__(self, other):
         if isinstance(other, Expr):
-            frame = self
             if not are_co_aligned(self, other):
-                frame, other = maybe_align_partitions(
-                    frame, other, divisions=calc_divisions_for_align(frame, other)
-                )
-            return Filter(frame, other)
+                return FilterAlign(self, other)
+            return Filter(self, other)
         else:
             return Projection(self, other)  # df[["a", "b", "c"]]
 
@@ -1819,7 +1816,9 @@ class Filter(Blockwise):
                     type(self)(self.frame, result), *parent.operands[1:]
                 )
 
-        if isinstance(parent, Filter) and not isinstance(self.frame, Filter):
+        if isinstance(parent, (FilterAlign, Filter)) and not isinstance(
+            self.frame, (FilterAlign, Filter)
+        ):
             if not self.frame._filter_passthrough_available(self, dependents):
                 # We want to collect filters again when we can't move them
                 # anymore. Otherwise, a chain of Filters might block Projections
@@ -1842,7 +1841,7 @@ class Filter(Blockwise):
                 # We can't push Projections through filters if the preceding operation
                 # allows us to push filters further down the graph because Projections
                 # block filter pushdown
-                if not isinstance(self.frame, Filter):
+                if not isinstance(self.frame, (FilterAlign, Filter)):
                     return
                 elif is_filter_pushdown_available(
                     self.frame, self, dependents, allow_reduction=False
@@ -3200,6 +3199,16 @@ class CombineFrameAlign(MaybeAlignPartitions):
     _expr_cls = CombineSeries
 
 
+class FilterAlign(MaybeAlignPartitions):
+    _projection_passthrough = True
+    _filter_passthrough = True
+    _parameters = ["frame", "predicate"]
+    _expr_cls = Filter
+
+    def _simplify_up(self, parent, dependents):
+        return Filter._simplify_up(self, parent, dependents)
+
+
 class AssignAlign(MaybeAlignPartitions):
     _parameters = ["frame", "column", "value"]
     _expr_cls = Assign
diff --git a/dask_expr/_repartition.py b/dask_expr/_repartition.py
index 1e845b5..6085866 100644
--- a/dask_expr/_repartition.py
+++ b/dask_expr/_repartition.py
@@ -13,7 +13,7 @@ from dask.utils import iter_chunks, parse_bytes
 from pandas.api.types import is_datetime64_any_dtype, is_numeric_dtype
 from tlz import unique
 
-from dask_expr._expr import Expr, Projection, plain_column_projection
+from dask_expr._expr import Expr, Filter, Projection, plain_column_projection
 from dask_expr._reductions import TotalMemoryUsageFrame
 from dask_expr._util import LRU
 
@@ -35,6 +35,7 @@ class Repartition(Expr):
         "partition_size": None,
     }
     _is_length_preserving = True
+    _filter_passthrough = True
 
     @functools.cached_property
     def _meta(self):
@@ -122,6 +123,10 @@ class Repartition(Expr):
             raise NotImplementedError()
 
     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):
             return plain_column_projection(self, parent, dependents)
 
diff --git a/dask_expr/_shuffle.py b/dask_expr/_shuffle.py
index 2827031..7b9bcb1 100644
--- a/dask_expr/_shuffle.py
+++ b/dask_expr/_shuffle.py
@@ -38,6 +38,7 @@ from dask_expr._expr import (
     Projection,
     ToSeriesIndex,
     determine_column_projection,
+    is_filter_pushdown_available,
 )
 from dask_expr._reductions import (
     All,
@@ -787,6 +788,7 @@ class SetIndex(BaseSetIndexSortValues):
         "options": None,
         "append": False,
     }
+    _filter_passthrough = True
 
     @property
     def _projection_columns(self):
@@ -856,7 +858,7 @@ class SetIndex(BaseSetIndexSortValues):
         )
 
     def _simplify_up(self, parent, dependents):
-        from dask_expr._expr import Filter, Head, Index, Tail
+        from dask_expr._expr import Filter, Head, Tail
 
         # TODO, handle setting index with other frame
         if (
@@ -889,14 +891,18 @@ class SetIndex(BaseSetIndexSortValues):
                 type(self)(self.frame[columns], *self.operands[1:]),
                 parent.operand("columns"),
             )
+        if isinstance(parent, Filter) and self._filter_passthrough_available(
+            parent, dependents
+        ):
+            return self._filter_simplification(parent)
+
+    def _filter_passthrough_available(self, parent, dependents):
+        if is_filter_pushdown_available(self, parent, dependents):
+            from dask_expr._expr import Index
 
-        if isinstance(parent, Filter):
             p = parent.predicate
-            if any(isinstance(x, Index) for x in p.walk()):
-                # Punt on cases where the new index is part of the filter
-                return
-            predicate = parent.substitute(self, self.frame).predicate
-            return type(self)(self.frame[predicate], *self.operands[1:])
+            return not any(isinstance(x, Index) for x in p.walk())
+        return False
 
 
 class SortValues(BaseSetIndexSortValues):
@@ -925,6 +931,7 @@ class SortValues(BaseSetIndexSortValues):
         "ignore_index": False,
         "shuffle_method": None,
     }
+    _filter_passthrough = True
 
     def _divisions(self):
         if self.frame.npartitions == 1:
@@ -1030,11 +1037,10 @@ class SortValues(BaseSetIndexSortValues):
                 self.frame, n=parent.n, _columns=self.by, ascending=self.ascending
             )
 
-        if isinstance(parent, Filter):
-            return SortValues(
-                Filter(self.frame, parent.predicate.substitute(self, self.frame)),
-                *self.operands[1:],
-            )
+        if isinstance(parent, Filter) and self._filter_passthrough_available(
+            parent, dependents
+        ):
+            return self._filter_simplification(parent)
         if isinstance(parent, Projection):
             columns = determine_column_projection(
                 self, parent, dependents, additional_columns=self.by


Test patch:
diff --git a/dask_expr/tests/test_indexing.py b/dask_expr/tests/test_indexing.py
index 6cfea63..535f3a2 100644
--- a/dask_expr/tests/test_indexing.py
+++ b/dask_expr/tests/test_indexing.py
@@ -141,3 +141,16 @@ def test_loc_with_function(df, pdf):
         return _df.columns.str.contains("y")
 
     assert_eq(df.loc[:, _col_loc_fun], pdf.loc[:, _col_loc_fun])
+
+
+def test_getitem_align():
+    df = pd.DataFrame(
+        {
+            "A": [1, 2, 3, 4, 5, 6, 7, 8, 9],
+            "B": [9, 8, 7, 6, 5, 4, 3, 2, 1],
+            "C": [True, False, True] * 3,
+        },
+        columns=list("ABC"),
+    )
+    ddf = from_pandas(df, 2)
+    assert_eq(ddf[ddf.C.repartition([0, 2, 5, 8])], df[df.C])
diff --git a/dask_expr/tests/test_repartition.py b/dask_expr/tests/test_repartition.py
index fcd64ed..7cb4cf3 100644
--- a/dask_expr/tests/test_repartition.py
+++ b/dask_expr/tests/test_repartition.py
@@ -110,3 +110,24 @@ def test_repartition_empty_partitions_dtype():
         df[df.x < 5].repartition(npartitions=1),
         pdf[pdf.x < 5],
     )
+
+
+def test_repartition_filter_pushdown():
+    pdf = pd.DataFrame({"x": [1, 2, 3, 4, 5, 6, 7, 8] * 10, "y": 1, "z": 2})
+    df = from_pandas(pdf, npartitions=10)
+    result = df.repartition(npartitions=5)
+    result = result[result.x > 5.0]
+    expected = df[df.x > 5.0].repartition(npartitions=5)
+    assert result.simplify()._name == expected._name
+
+    result = df.repartition(npartitions=5)
+    result = result[result.x > 5.0][["x", "y"]]
+    expected = df[["x", "y"]]
+    expected = expected[expected.x > 5.0].repartition(npartitions=5)
+    assert result.simplify()._name == expected.simplify()._name
+
+    result = df.repartition(npartitions=5)[["x", "y"]]
+    result = result[result.x > 5.0]
+    expected = df[["x", "y"]]
+    expected = expected[expected.x > 5.0].repartition(npartitions=5)
+    assert result.simplify()._name == expected.simplify()._name
diff --git a/dask_expr/tests/test_shuffle.py b/dask_expr/tests/test_shuffle.py
index fb4ecc2..c7b9772 100644
--- a/dask_expr/tests/test_shuffle.py
+++ b/dask_expr/tests/test_shuffle.py
@@ -5,10 +5,10 @@ import numpy as np
 import pytest
 
 from dask_expr import from_pandas, new_collection
-from dask_expr._expr import Assign, Blockwise
+from dask_expr._expr import Assign, Blockwise, Filter
 from dask_expr._reductions import NFirst, NLast
 from dask_expr._repartition import RepartitionToFewer
-from dask_expr._shuffle import TaskShuffle, divisions_lru
+from dask_expr._shuffle import BaseSetIndexSortValues, TaskShuffle, divisions_lru
 from dask_expr.io import FromPandas
 from dask_expr.tests._util import _backend_library, assert_eq, xfail_gpu
 
@@ -689,22 +689,55 @@ def test_shuffle_no_assign(df, pdf):
     assert len([x for x in q.walk() if isinstance(x, Assign)]) == 0
 
 
-def test_shuffle_filter_pushdown(pdf):
+@pytest.mark.parametrize("meth", ["shuffle", "sort_values"])
+def test_shuffle_filter_pushdown(pdf, meth):
     pdf["z"] = 1
     df = from_pandas(pdf, npartitions=10)
-    result = df.shuffle("x")
+    result = getattr(df, meth)("x")
     result = result[result.x > 5.0]
-    expected = df[df.x > 5.0].shuffle("x")
+    expected = getattr(df[df.x > 5.0], meth)("x")
     assert result.simplify()._name == expected._name
 
-    result = df.shuffle("x")
+    result = getattr(df, meth)("x")
     result = result[result.x > 5.0][["x", "y"]]
     expected = df[["x", "y"]]
-    expected = expected[expected.x > 5.0].shuffle("x")
+    expected = getattr(expected[expected.x > 5.0], meth)("x")
     assert result.simplify()._name == expected.simplify()._name
 
-    result = df.shuffle("x")[["x", "y"]]
+    result = getattr(df, meth)("x")[["x", "y"]]
     result = result[result.x > 5.0]
     expected = df[["x", "y"]]
-    expected = expected[expected.x > 5.0].shuffle("x")
+    expected = getattr(expected[expected.x > 5.0], meth)("x")
+    assert result.simplify()._name == expected.simplify()._name
+
+
+@pytest.mark.parametrize("meth", ["set_index", "sort_values"])
+def test_sort_values_avoid_overeager_filter_pushdown(meth):
+    pdf1 = pd.DataFrame({"a": [4, 2, 3], "b": [1, 2, 3]})
+    df = from_pandas(pdf1, npartitions=2)
+    df = getattr(df, meth)("a")
+    df = df[df.b > 2] + df.b.sum()
+    result = df.simplify()
+    assert isinstance(result.expr.left, Filter)
+    assert isinstance(result.expr.left.frame, BaseSetIndexSortValues)
+
+
+def test_set_index_filter_pushdown():
+    pdf = pd.DataFrame({"x": [1, 2, 3, 4, 5, 6, 7, 8] * 10, "y": 1, "z": 2})
+    df = from_pandas(pdf, npartitions=10)
+    result = df.set_index("x")
+    result = result[result.y == 1]
+    expected = df[df.y == 1].set_index("x")
+    assert result.simplify()._name == expected._name
+
+    result = df.set_index("x")
+    result = result[result.y == 1][["y"]]
+    expected = df[["x", "y"]]
+    expected = expected[expected.y == 1].set_index("x")
+    assert result.simplify()._name == expected.simplify()._name
+
+    result = df.set_index("x")[["y"]]
+    result = result[result.y == 1]
+    expected = df[["x", "y"]]
+    expected = expected[expected.y == 1].set_index("x")
     assert result.simplify()._name == expected.simplify()._name