Repository: dask-contrib/dask-expr
Pull request #865 (2024-02-12 16:36:28)

Avoid passing unknown divisions into setindexblockwise
This just blows up the graph without any benefit


Patch:
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index 5a03c72..bdd821f 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -1812,6 +1812,13 @@ class Filter(Blockwise):
     operation = operator.getitem
 
     def _simplify_up(self, parent, dependents):
+        if isinstance(self.predicate, Or):
+            result = rewrite_filters(self.predicate)
+            if result._name != self.predicate._name:
+                return type(parent)(
+                    type(self)(self.frame, result), *parent.operands[1:]
+                )
+
         if isinstance(parent, Filter) and not isinstance(self.frame, Filter):
             if not self.frame._filter_passthrough_available(self, dependents):
                 # We want to collect filters again when we can't move them
@@ -3200,19 +3207,14 @@ class AssignAlign(MaybeAlignPartitions):
             if not isinstance(columns, list):
                 columns = [columns]
 
-            cols = set(columns) - set(self.keys)
+            cols = set(columns) - {self.column}
             if cols == set(self.frame.columns):
                 # Protect against pushing the same projection twice
                 return
 
-            diff = set(self.keys) - set(columns)
-            if len(diff) == len(self.keys):
+            diff = {self.column} - set(columns)
+            if len(diff) == 1:
                 return type(parent)(self.frame, *parent.operands[1:])
-            elif len(diff) > 0:
-                new_args = []
-                for k, v in zip(self.keys, self.vals):
-                    if k in columns:
-                        new_args.extend([k, v])
             else:
                 new_args = self.operands[1:]
 
@@ -3511,6 +3513,78 @@ def is_filter_pushdown_available(expr, parent, dependents, allow_reduction=True)
     )
 
 
+def rewrite_filters(predicate):
+    """Rewriting a filter to decompose OR clauses. If a predicate part is part of
+    all OR clauses, we can move it to the front so that we can push it down.
+    """
+    or_components = _get_predicate_components(predicate, [])
+    if len(or_components) == 1:
+        return predicate
+    result = _replace_common_or_components(or_components[0], or_components[1:])
+    if result is None:
+        return predicate
+    return result
+
+
+def _get_predicate_components(predicate, components, type_=Or):
+    if not isinstance(predicate, type_):
+        components.append(predicate)
+        return components
+    if isinstance(predicate.left, type_):
+        components = _get_predicate_components(predicate.left, components, type_)
+    else:
+        components.append(predicate.left)
+    if isinstance(predicate.right, type_):
+        components = _get_predicate_components(predicate.right, components, type_)
+    else:
+        components.append(predicate.right)
+    return components
+
+
+def _convert_mapping(components):
+    return dict(zip([e._name for e in components], components))
+
+
+def _replace_common_or_components(expr, or_components):
+    and_component = _get_predicate_components(expr, [], type_=And)
+    mapping = _convert_mapping(and_component)
+    and_components = [
+        _get_predicate_components(c, [], type_=And) for c in or_components
+    ]
+    and_components = list(map(_convert_mapping, and_components))
+
+    replacements = []
+    for c in mapping.keys():
+        if all(c in comp for comp in and_components):
+            # We can pull this component out if it's part of all or components
+            replacements.append(c)
+    if len(replacements) == 0:
+        # exit if we can't replace anything
+        return
+
+    outer_component = mapping[replacements[0]]
+    for r in replacements[1:]:
+        # construct the outer component
+        outer_component = outer_component & mapping[r]
+
+    #
+    result_components = []
+    for comp in [mapping] + and_components:
+        keep_components = [c for c in comp if c not in replacements]
+        if len(keep_components) == 0:
+            # Just return outer_component if we can replace a whole OR component
+            return outer_component
+        result_component = comp[keep_components[0]]
+        for c in keep_components[1:]:
+            result_component = result_component & comp[c]
+        result_components.append(result_component)
+
+    or_component = result_components[0]
+    for c in result_components[1:]:
+        or_component = or_component | c
+    return outer_component & or_component
+
+
 def _check_dependents_are_predicates(
     expr, other_names, parent: Expr, dependents, allow_reduction=True
 ):
diff --git a/dask_expr/_reductions.py b/dask_expr/_reductions.py
index 38f1359..f8072f9 100644
--- a/dask_expr/_reductions.py
+++ b/dask_expr/_reductions.py
@@ -232,8 +232,7 @@ class ShuffleReduce(Expr):
 
         # Reset the index if we we used it for shuffling
         if split_by_index:
-            divisions = (None,) * (shuffle_npartitions + 1)
-            shuffled = SetIndexBlockwise(shuffled, split_by, True, divisions)
+            shuffled = SetIndexBlockwise(shuffled, split_by, True, None)
 
         # Convert back to Series if necessary
         if self.shuffle_by_index is not False:
@@ -241,9 +240,8 @@ class ShuffleReduce(Expr):
                 shuffled = shuffled[shuffled.columns[0]]
             elif is_index_like(self._meta):
                 column = shuffled.columns[0]
-                shuffled = Index(
-                    SetIndexBlockwise(shuffled, column, True, shuffled.divisions)
-                )
+                divs = None if shuffled.divisions[0] is None else shuffled.divisions
+                shuffled = Index(SetIndexBlockwise(shuffled, column, True, divs))
                 if column == "__index__":
                     shuffled = RenameSeries(shuffled, self.frame._meta.name)
 


Test patch:
diff --git a/dask_expr/tests/test_align_partitions.py b/dask_expr/tests/test_align_partitions.py
index 00f2911..a9c8bf3 100644
--- a/dask_expr/tests/test_align_partitions.py
+++ b/dask_expr/tests/test_align_partitions.py
@@ -97,3 +97,13 @@ def test_series_map(sorted_index, sorted_map_index):
     dask_map = from_pandas(mapper, npartitions=1, sort=False)
     result = dask_base.map(dask_map)
     assert_eq(expected, result)
+
+
+def test_assign_align_partitions():
+    pdf = pd.DataFrame({"x": [0] * 20, "y": range(20)})
+    df = from_pandas(pdf, npartitions=2)
+    s = pd.Series(range(10, 30))
+    ds = from_pandas(s, npartitions=df.npartitions)
+    result = df.assign(z=ds)[["y", "z"]]
+    expected = pdf.assign(z=s)[["y", "z"]]
+    assert_eq(result, expected)
diff --git a/dask_expr/tests/test_predicate_pushdown.py b/dask_expr/tests/test_predicate_pushdown.py
new file mode 100644
index 0000000..276d90b
--- /dev/null
+++ b/dask_expr/tests/test_predicate_pushdown.py
@@ -0,0 +1,59 @@
+import pytest
+
+from dask_expr import from_pandas
+from dask_expr._expr import rewrite_filters
+from dask_expr.tests._util import _backend_library, assert_eq
+
+pd = _backend_library()
+
+
+@pytest.fixture
+def pdf():
+    pdf = pd.DataFrame({"x": range(100), "a": 1, "b": 2})
+    pdf["y"] = pdf.x // 7  # Not unique; duplicates span different partitions
+    yield pdf
+
+
+@pytest.fixture
+def df(pdf):
+    yield from_pandas(pdf, npartitions=10)
+
+
+def test_rewrite_filters(df):
+    predicate = (df.x == 1) | (df.x == 1)
+    expected = df.x == 1
+    assert rewrite_filters(predicate.expr)._name == expected._name
+
+    predicate = (df.x == 1) | ((df.x == 1) & (df.y == 2))
+    expected = df.x == 1
+    assert rewrite_filters(predicate.expr)._name == expected._name
+
+    predicate = ((df.x == 1) & (df.y == 3)) | ((df.x == 1) & (df.y == 2))
+    expected = (df.x == 1) & ((df.y == 3) | (df.y == 2))
+    assert rewrite_filters(predicate.expr)._name == expected._name
+
+    predicate = ((df.x == 1) & (df.y == 3) & (df.a == 1)) | (
+        (df.x == 1) & (df.y == 2) & (df.a == 1)
+    )
+    expected = ((df.x == 1) & (df.a == 1)) & ((df.y == 3) | (df.y == 2))
+    assert rewrite_filters(predicate.expr)._name == expected._name
+
+    predicate = (df.x == 1) | (df.y == 1)
+    assert rewrite_filters(predicate.expr)._name == predicate._name
+
+    predicate = df.x == 1
+    assert rewrite_filters(predicate.expr)._name == predicate._name
+
+    predicate = (df.x.isin([1, 2, 3]) & (df.y == 3)) | (
+        df.x.isin([1, 2, 3]) & (df.y == 2)
+    )
+    expected = df.x.isin([1, 2, 3]) & ((df.y == 3) | (df.y == 2))
+    assert rewrite_filters(predicate.expr)._name == expected._name
+
+
+def test_rewrite_filters_query(df, pdf):
+    result = df[((df.x == 1) & (df.y > 1)) | ((df.x == 1) & (df.y > 2))]
+    result = result[["x"]]
+    expected = pdf[((pdf.x == 1) & (pdf.y > 1)) | ((pdf.x == 1) & (pdf.y > 2))]
+    expected = expected[["x"]]
+    assert_eq(result, expected)