Repository: dask-contrib/dask-expr
Pull request #874 (2024-02-14 10:38:37)

Fix some meta calculations and rest index projections
Nonempty meta objects can actually screw us because pandas aligns under the hood


Patch:
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index bdd821f..9d7535d 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"]
@@ -1205,7 +1205,7 @@ class RenameSeries(Elemwise):
         args = [
             meta_nonempty(op._meta) if isinstance(op, Expr) else op for op in self._args
         ]
-        return self.operation(*args, **self._kwargs)
+        return make_meta(self.operation(*args, **self._kwargs))
 
     @staticmethod
     def operation(df, index, sorted_index):
@@ -2099,12 +2099,16 @@ class ResetIndex(Elemwise):
             return self._filter_simplification(parent, predicate)
 
         if isinstance(parent, Projection):
-            if self.frame.ndim == 1 and not self.drop and not isinstance(parent, list):
+            if self.frame.ndim == 1 and not self.drop:
+                if isinstance(parent.operand("columns"), list):
+                    # Don't bother, dimensionality changes are tricky here and
+                    # potential improvement is tiny
+                    return
                 col = parent.operand("columns")
-                if col in (self.name, "index"):
+                if col in (self.name, "index", self.frame._meta.index.name):
                     return
                 if all(
-                    isinstance(d, Projection) and d.operand("columns") == col
+                    isinstance(d(), Projection) and d().operand("columns") == col
                     for d in dependents[self._name]
                 ):
                     return type(self)(self.frame, True, self.name)
@@ -2961,7 +2965,7 @@ class Diff(MapOverlap):
 
     @functools.cached_property
     def _meta(self):
-        return meta_nonempty(self.frame._meta).diff(**self.kwargs)
+        return make_meta(meta_nonempty(self.frame._meta).diff(**self.kwargs))
 
     def _simplify_up(self, parent, dependents):
         if isinstance(parent, Projection):
@@ -3069,7 +3073,7 @@ class Shift(MapOverlap):
 
     @functools.cached_property
     def _meta(self):
-        return meta_nonempty(self.frame._meta).shift(**self.kwargs)
+        return make_meta(meta_nonempty(self.frame._meta).shift(**self.kwargs))
 
     @functools.cached_property
     def kwargs(self):
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_collection.py b/dask_expr/tests/test_collection.py
index 247893f..bbdee63 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -402,6 +402,10 @@ def test_reset_index_projections(pdf):
     expected = expected + expected.x.sum()
     assert_eq(q, expected, check_index=False)
 
+    pdf = pdf.to_frame()
+    df = from_pandas(pdf, npartitions=10)
+    assert_eq(df.y.reset_index()[["y"]], pdf.y.reset_index()[["y"]], check_index=False)
+
 
 def test_series_map_meta():
     ser = pd.Series(
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