Repository: dask-contrib/dask-expr
Pull request #858 (2024-02-08 11:48:08)

Simplify sort_values simplify_up step
This became unnecessary complex, the performance differences is 50 vs 100ms on a df with 1 million rows and 10 columns, which doesn't matter for us


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/_reductions.py b/dask_expr/_reductions.py
index f8072f9..3b562cf 100644
--- a/dask_expr/_reductions.py
+++ b/dask_expr/_reductions.py
@@ -1250,14 +1250,6 @@ class NLargest(ReductionConstantDim):
         return self.chunk_kwargs
 
 
-def _nsmallest_slow(df, columns, n):
-    return df.sort_values(by=columns).head(n)
-
-
-def _nlargest_slow(df, columns, n):
-    return df.sort_values(by=columns).tail(n)
-
-
 def _nfirst(df, columns, n, ascending):
     return df.sort_values(by=columns, ascending=ascending).head(n)
 
@@ -1269,28 +1261,17 @@ def _nlast(df, columns, n, ascending):
 class NFirst(NLargest):
     _parameters = ["frame", "n", "_columns", "ascending", "split_every"]
     _defaults = {"n": 5, "_columns": None, "ascending": None, "split_every": None}
-    reduction_chunk = _nfirst
-    reduction_aggregate = _nfirst
-
-    @property
-    def chunk_kwargs(self):
-        return {"ascending": self.ascending, **super().chunk_kwargs}
-
-
-class NLast(NLargest):
-    _parameters = ["frame", "n", "_columns", "ascending", "split_every"]
-    _defaults = {"n": 5, "_columns": None, "ascending": None, "split_every": None}
-    reduction_chunk = _nlast
-    reduction_aggregate = _nlast
+    reduction_chunk = staticmethod(_nfirst)
+    reduction_aggregate = staticmethod(_nfirst)
 
     @property
     def chunk_kwargs(self):
         return {"ascending": self.ascending, **super().chunk_kwargs}
 
 
-class NLargestSlow(NLargest):
-    reduction_chunk = _nlargest_slow
-    reduction_aggregate = _nlargest_slow
+class NLast(NFirst):
+    reduction_chunk = staticmethod(_nlast)
+    reduction_aggregate = staticmethod(_nlast)
 
 
 class NSmallest(NLargest):
@@ -1298,11 +1279,6 @@ class NSmallest(NLargest):
     reduction_aggregate = M.nsmallest
 
 
-class NSmallestSlow(NLargest):
-    reduction_chunk = _nsmallest_slow
-    reduction_aggregate = _nsmallest_slow
-
-
 class ValueCounts(ReductionConstantDim):
     _defaults = {
         "sort": None,
diff --git a/dask_expr/_shuffle.py b/dask_expr/_shuffle.py
index e34eee6..2827031 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,
@@ -52,10 +53,8 @@ from dask_expr._reductions import (
     NBytes,
     NFirst,
     NLargest,
-    NLargestSlow,
     NLast,
     NSmallest,
-    NSmallestSlow,
     Prod,
     Size,
     Sum,
@@ -63,7 +62,7 @@ from dask_expr._reductions import (
     ValueCounts,
 )
 from dask_expr._repartition import Repartition, RepartitionToFewer
-from dask_expr._util import LRU, _convert_to_list, is_valid_nth_dtype
+from dask_expr._util import LRU, _convert_to_list
 
 
 class ShuffleBase(Expr):
@@ -83,6 +82,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 +91,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
@@ -860,10 +864,7 @@ class SetIndex(BaseSetIndexSortValues):
             and isinstance(self._other, (int, str))
             and self._other in self.frame.columns
         ):
-            if is_valid_nth_dtype(self.frame._meta.dtypes[self._other]):
-                head = NSmallest(self.frame, n=parent.n, _columns=self._other)
-            else:
-                head = NSmallestSlow(self.frame, n=parent.n, _columns=self._other)
+            head = NFirst(self.frame, n=parent.n, _columns=self._other, ascending=True)
             return SetIndex(head, _other=self._other)
 
         if (
@@ -871,10 +872,7 @@ class SetIndex(BaseSetIndexSortValues):
             and isinstance(self._other, (int, str))
             and self._other in self.frame.columns
         ):
-            if is_valid_nth_dtype(self.frame._meta.dtypes[self._other]):
-                tail = NLargest(self.frame, n=parent.n, _columns=self._other)
-            else:
-                tail = NLargestSlow(self.frame, n=parent.n, _columns=self._other)
+            tail = NLast(self.frame, n=parent.n, _columns=self._other, ascending=True)
             return SetIndex(tail, _other=self._other)
 
         if isinstance(parent, Projection):
@@ -1022,48 +1020,15 @@ class SortValues(BaseSetIndexSortValues):
     def _simplify_up(self, parent, dependents):
         from dask_expr._expr import Filter, Head, Tail
 
-        _all_ascending = (
-            self.ascending
-            if isinstance(self.ascending, bool)
-            else all(ascending for ascending in self.ascending)
-        )
-        _all_descending = (
-            not self.ascending
-            if isinstance(self.ascending, bool)
-            else all(not ascending for ascending in self.ascending)
-        )
-
         if isinstance(parent, Head):
-            if _all_ascending:
-                if is_valid_nth_dtype(self._meta_by_dtype):
-                    return NSmallest(self.frame, n=parent.n, _columns=self.by)
-                else:
-                    return NSmallestSlow(self.frame, n=parent.n, _columns=self.by)
-            elif _all_descending:
-                if is_valid_nth_dtype(self._meta_by_dtype):
-                    return NLargest(self.frame, n=parent.n, _columns=self.by)
-                else:
-                    return NLargestSlow(self.frame, n=parent.n, _columns=self.by)
-            else:
-                return NFirst(
-                    self.frame, n=parent.n, _columns=self.by, ascending=self.ascending
-                )
+            return NFirst(
+                self.frame, n=parent.n, _columns=self.by, ascending=self.ascending
+            )
 
         if isinstance(parent, Tail):
-            if _all_ascending:
-                if is_valid_nth_dtype(self._meta_by_dtype):
-                    return NLargest(self.frame, n=parent.n, _columns=self.by)
-                else:
-                    return NLargestSlow(self.frame, n=parent.n, _columns=self.by)
-            elif _all_descending:
-                if is_valid_nth_dtype(self._meta_by_dtype):
-                    return NSmallest(self.frame, n=parent.n, _columns=self.by)
-                else:
-                    return NSmallestSlow(self.frame, n=parent.n, _columns=self.by)
-            else:
-                return NLast(
-                    self.frame, n=parent.n, _columns=self.by, ascending=self.ascending
-                )
+            return NLast(
+                self.frame, n=parent.n, _columns=self.by, ascending=self.ascending
+            )
 
         if isinstance(parent, Filter):
             return SortValues(
diff --git a/dask_expr/_util.py b/dask_expr/_util.py
index 41f9bc5..2c35a38 100644
--- a/dask_expr/_util.py
+++ b/dask_expr/_util.py
@@ -13,7 +13,6 @@ from dask.base import normalize_token, tokenize
 from dask.dataframe._compat import is_string_dtype
 from dask.utils import get_default_shuffle_method
 from packaging.version import Version
-from pandas.api.types import is_datetime64_dtype, is_numeric_dtype
 
 K = TypeVar("K", bound=Hashable)
 V = TypeVar("V")
@@ -95,10 +94,6 @@ def is_scalar(x):
     return not isinstance(x, Expr)
 
 
-def is_valid_nth_dtype(dtype):
-    return is_numeric_dtype(dtype) or is_datetime64_dtype(dtype)
-
-
 def _tokenize_deterministic(*args, **kwargs) -> str:
     # Utility to be strict about deterministic tokens
     with config.set({"tokenize.ensure-deterministic": True}):
diff --git a/dask_expr/io/parquet.py b/dask_expr/io/parquet.py
index a103e60..8cd9e8c 100644
--- a/dask_expr/io/parquet.py
+++ b/dask_expr/io/parquet.py
@@ -77,6 +77,33 @@ def normalize_pa_schema(schema):
     return schema.to_string()
 
 
+@normalize_token.register(pq.ParquetSchema)
+def normalize_pq_schema(schema):
+    try:
+        return hash(schema)
+    except TypeError:  # pyarrow version not supporting ParquetSchema hash
+        return hash(repr(schema))
+
+
+@normalize_token.register(pq.FileMetaData)
+def normalize_pq_filemetadata(meta):
+    try:
+        return hash(meta)
+    except TypeError:
+        # pyarrow version not implementing hash for FileMetaData
+        # use same logic as implemented in version that does support hashing
+        # https://github.com/apache/arrow/blob/bbe59b35de33a0534fc76c9617aa4746031ce16c/python/pyarrow/_parquet.pyx#L853
+        return hash(
+            (
+                repr(meta.schema),
+                meta.num_rows,
+                meta.num_row_groups,
+                meta.format_version,
+                meta.serialized_size,
+            )
+        )
+
+
 class ToParquet(Expr):
     _parameters = [
         "frame",


Test patch:
diff --git a/dask_expr/io/tests/test_io.py b/dask_expr/io/tests/test_io.py
index f6b8e98..3bea37f 100644
--- a/dask_expr/io/tests/test_io.py
+++ b/dask_expr/io/tests/test_io.py
@@ -4,6 +4,7 @@ import os
 import dask.array as da
 import dask.dataframe as dd
 import numpy as np
+import pyarrow.parquet as pq
 import pytest
 from dask.array.utils import assert_eq as array_assert_eq
 from dask.dataframe.utils import assert_eq
@@ -22,7 +23,7 @@ from dask_expr import (
 )
 from dask_expr._expr import Expr, Lengths, Literal, Replace
 from dask_expr._reductions import Len
-from dask_expr.io import FromArray, FromMap, ReadCSV, ReadParquet
+from dask_expr.io import FromArray, FromMap, ReadCSV, ReadParquet, parquet
 from dask_expr.tests._util import _backend_library
 
 # Set DataFrame backend for this module
@@ -501,3 +502,21 @@ def test_from_dict():
     expected = pd.DataFrame(data)
     assert_eq(result, expected)
     assert_eq(DataFrame.from_dict(data), expected)
+
+
+@pytest.mark.parametrize("normalizer", ("filemetadata", "schema"))
+def test_normalize_token_parquet_filemetadata_and_schema(tmpdir, normalizer):
+    df = pd.DataFrame({c: range(10) for c in "abcde"})
+    path = tmpdir / "df.parquet"
+    df.to_parquet(path)
+
+    meta = pq.read_metadata(path)
+
+    if normalizer == "filemetadata":
+        normalizer = parquet.normalize_pq_filemetadata
+        obj = meta
+    else:
+        normalizer = parquet.normalize_pq_schema
+        obj = meta.schema
+
+    assert normalizer(obj) == normalizer(obj)
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..fb4ecc2 100644
--- a/dask_expr/tests/test_shuffle.py
+++ b/dask_expr/tests/test_shuffle.py
@@ -367,81 +367,73 @@ def test_sort_values_descending(df, pdf):
 
 
 def test_sort_head_nlargest(df, pdf):
-    a = df.sort_values("x", ascending=False).head(10, compute=False).expr
-    b = df.nlargest(10, columns=["x"]).expr
+    a = df.sort_values("x", ascending=False).head(10, compute=False)
+    b = NFirst(df, 10, _columns=["x"], ascending=False)
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values("x", ascending=True).head(10, compute=False).expr
-    b = df.nsmallest(10, columns=["x"]).expr
+    a = df.sort_values("x", ascending=True).head(10, compute=False)
+    b = NFirst(df, 10, _columns=["x"], ascending=True)
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values("x", ascending=[False]).head(10, compute=False).expr
-    b = df.nlargest(10, columns=["x"]).expr
+    a = df.sort_values("x", ascending=[False]).head(10, compute=False)
+    b = NFirst(df, 10, _columns=["x"], ascending=[False])
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values("x", ascending=[True]).head(10, compute=False).expr
-    b = df.nsmallest(10, columns=["x"]).expr
+    a = df.sort_values("x", ascending=[True]).head(10, compute=False)
+    b = NFirst(df, 10, _columns=["x"], ascending=[True])
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values(["x"], ascending=[False]).head(10, compute=False).expr
-    b = df.nlargest(10, columns=["x"]).expr
+    a = df.sort_values(["x"], ascending=[False]).head(10, compute=False)
+    b = NFirst(df, 10, _columns=["x"], ascending=[False])
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values(["x"], ascending=[True]).head(10, compute=False).expr
-    b = df.nsmallest(10, columns=["x"]).expr
+    a = df.sort_values(["x"], ascending=[True]).head(10, compute=False)
+    b = NFirst(df, 10, _columns=["x"], ascending=[True])
     assert a.optimize()._name == b.optimize()._name
 
-    a = (
-        df.sort_values(["x", "y"], ascending=[False, False])
-        .head(10, compute=False)
-        .expr
-    )
-    b = df.nlargest(10, columns=["x", "y"]).expr
+    a = df.sort_values(["x", "y"], ascending=[False, False]).head(10, compute=False)
+    b = NFirst(df, 10, _columns=["x", "y"], ascending=[False, False])
     assert a.optimize()._name == b.optimize()._name
 
     a = df.sort_values(["x", "y"], ascending=[True, True]).head(10, compute=False).expr
-    b = df.nsmallest(10, columns=["x", "y"]).expr
+    b = NFirst(df, 10, _columns=["x", "y"], ascending=[True, True])
     assert a.optimize()._name == b.optimize()._name
 
 
 def test_sort_tail_nsmallest(df, pdf):
-    a = df.sort_values("x", ascending=False).tail(10, compute=False).expr
-    b = df.nsmallest(10, columns=["x"]).expr
+    a = df.sort_values("x", ascending=False).tail(10, compute=False)
+    b = NLast(df, 10, _columns=["x"], ascending=False)
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values("x", ascending=True).tail(10, compute=False).expr
-    b = df.nlargest(10, columns=["x"]).expr
+    a = df.sort_values("x", ascending=True).tail(10, compute=False)
+    b = NLast(df, 10, _columns=["x"], ascending=True)
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values("x", ascending=[False]).tail(10, compute=False).expr
-    b = df.nsmallest(10, columns=["x"]).expr
+    a = df.sort_values("x", ascending=[False]).tail(10, compute=False)
+    b = NLast(df, 10, _columns=["x"], ascending=[False])
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values("x", ascending=[True]).tail(10, compute=False).expr
-    b = df.nlargest(10, columns=["x"]).expr
+    a = df.sort_values("x", ascending=[True]).tail(10, compute=False)
+    b = NLast(df, 10, _columns=["x"], ascending=[True])
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values(["x"], ascending=[False]).tail(10, compute=False).expr
-    b = df.nsmallest(10, columns=["x"]).expr
+    a = df.sort_values(["x"], ascending=[False]).tail(10, compute=False)
+    b = NLast(df, 10, _columns=["x"], ascending=[False])
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.sort_values(["x"], ascending=[True]).tail(10, compute=False).expr
-    b = df.nlargest(10, columns=["x"]).expr
+    a = df.sort_values(["x"], ascending=[True]).tail(10, compute=False)
+    b = NLast(df, 10, _columns=["x"], ascending=[True])
     assert a.optimize()._name == b.optimize()._name
 
     with pytest.raises(ValueError, match="Length of ascending"):
         df.sort_values(["x", "y"], ascending=[False]).tail(10, compute=False)
 
-    a = df.sort_values(["x", "y"], ascending=[True, True]).tail(10, compute=False).expr
-    b = df.nlargest(10, columns=["x", "y"]).expr
+    a = df.sort_values(["x", "y"], ascending=[True, True]).tail(10, compute=False)
+    b = NLast(df, 10, _columns=["x", "y"], ascending=[True, True])
     assert a.optimize()._name == b.optimize()._name
 
-    a = (
-        df.sort_values(["x", "y"], ascending=[False, False])
-        .tail(10, compute=False)
-        .expr
-    )
-    b = df.nsmallest(10, columns=["x", "y"]).expr
+    a = df.sort_values(["x", "y"], ascending=[False, False]).tail(10, compute=False)
+    b = NLast(df, 10, _columns=["x", "y"], ascending=[False, False])
     assert a.optimize()._name == b.optimize()._name
 
 
@@ -495,12 +487,12 @@ def test_sort_head_nlargest_string(pdf):
 
 
 def test_set_index_head_nlargest(df, pdf):
-    a = df.set_index("x").head(10, compute=False).expr
-    b = df.nsmallest(10, columns="x").set_index("x").expr
+    a = df.set_index("x").head(10, compute=False)
+    b = new_collection(NFirst(df, 10, _columns="x", ascending=True)).set_index("x")
     assert a.optimize()._name == b.optimize()._name
 
-    a = df.set_index("x").tail(10, compute=False).expr
-    b = df.nlargest(10, columns="x").set_index("x").expr
+    a = df.set_index("x").tail(10, compute=False)
+    b = new_collection(NLast(df, 10, _columns="x", ascending=True)).set_index("x")
     assert a.optimize()._name == b.optimize()._name
 
     # These still work, even if we haven't optimized them yet
@@ -695,3 +687,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