Repository: dask-contrib/dask-expr
Pull request #877 (2024-02-14 15:01:18)

Allow `columns` argument to be mandatory in `from_map`
`from_map` is performing some magic if the function that is being provided has a `columns` kwarg. However, if that columns `kwarg` is not optional, the magic breaks down. Status on main is that it raises this confusing error message

```
func() missing 1 required positional argument: 'columns'
```

this change adds functionality that allows the function to use a mandatory `columns` argument

Stumbled over this while trying to use it in https://github.com/dask-contrib/dask-deltatable/pull/69

Note: If the `meta` is not provided we have to raise because we don't have a way to generate the meta ourselves but in this case we're raising now a clean exception

fjetter: Lovely, another `RecursionError` 🔍 

Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index c20254a..ce39a1c 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -4857,6 +4857,7 @@ def from_map(
     NOTE: The underlying ``Expr`` object produced by this API
     will support column projection (via ``simplify``) if
     the ``func`` argument has "columns" in its signature.
+
     """
     from dask.dataframe.io.utils import DataFrameIOFunction
 
@@ -4887,8 +4888,16 @@ def from_map(
 
     # Check if `func` supports column projection
     allow_projection = False
-    if "columns" in inspect.signature(func).parameters:
+    columns_arg_required = False
+    if param := inspect.signature(func).parameters.get("columns", None):
         allow_projection = True
+        columns_arg_required = param.default is param.empty
+        if meta is no_default and columns_arg_required:
+            raise TypeError(
+                "Argument `func` of `from_map` has a required `columns` "
+                " parameter and not `meta` provided."
+                "Either provide `meta` yourself or make `columns` an optional argument."
+            )
     elif isinstance(func, DataFrameIOFunction):
         warnings.warn(
             "dask_expr does not support the DataFrameIOFunction "
@@ -4910,6 +4919,7 @@ def from_map(
                 columns,
                 args,
                 kwargs,
+                columns_arg_required,
                 meta,
                 enforce_metadata,
                 divisions,
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
diff --git a/dask_expr/io/io.py b/dask_expr/io/io.py
index 25ef1ab..1b6b34f 100644
--- a/dask_expr/io/io.py
+++ b/dask_expr/io/io.py
@@ -228,6 +228,7 @@ class FromMapProjectable(FromMap):
         "columns",
         "args",
         "kwargs",
+        "columns_arg_required",
         "user_meta",
         "enforce_metadata",
         "user_divisions",
@@ -265,9 +266,9 @@ class FromMapProjectable(FromMap):
     @functools.cached_property
     def kwargs(self):
         options = self.operand("kwargs")
-        if self.columns_operand:
+        if self.columns_arg_required or self.columns_operand:
             options = options.copy()
-            options["columns"] = self.columns_operand
+            options["columns"] = self.columns
         return options
 
     @functools.cached_property


Test patch:
diff --git a/dask_expr/io/tests/test_io.py b/dask_expr/io/tests/test_io.py
index 3bea37f..cca5d63 100644
--- a/dask_expr/io/tests/test_io.py
+++ b/dask_expr/io/tests/test_io.py
@@ -7,7 +7,7 @@ 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
+from dask.dataframe.utils import assert_eq, make_meta
 
 from dask_expr import (
     DataFrame,
@@ -463,6 +463,29 @@ def test_from_map(tmpdir, meta, label, allow_projection, enforce_metadata):
     assert_eq(result, pdf["a"], check_index=False)
 
 
+def func(path, columns):
+    raise NotImplementedError("This shouldn't ever be called")
+
+
+def func2(path, columns):
+    return pd.DataFrame({"a": range(10), "b": range(10)})
+
+
+def test_from_map_columns_required():
+    with pytest.raises(TypeError, match=r"columns.*optional"):
+        from_map(func, ["foo"])
+    meta = make_meta(func2("foo", None))
+    ddf = from_map(func2, ["foo"], meta=meta)
+    ddf_expected = from_pandas(func2("foo", None), npartitions=1)
+
+    assert_eq(ddf, ddf_expected, check_divisions=False)
+
+    actual = from_map(func2, ["foo"], meta=meta)[["a"]].optimize()
+    expected = from_map(func2, ["foo"], meta=meta, columns=["a"]).optimize()
+
+    assert actual._name == expected._name
+
+
 def test_from_array():
     arr = np.random.randint(1, 100, (100,))
     assert_eq(from_array(arr, chunksize=5), pd.Series(arr))
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