Skip to content

Commit

Permalink
Avoid recursion error if filter condition can optimize itself (#821)
Browse files Browse the repository at this point in the history
  • Loading branch information
phofl authored Jan 29, 2024
1 parent c88d6ce commit 6c6dc46
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 0 deletions.
3 changes: 3 additions & 0 deletions dask_expr/_expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -1183,6 +1183,9 @@ class Elemwise(Blockwise):

def _simplify_up(self, parent, dependents):
if self._filter_passthrough and isinstance(parent, Filter):
if self._name != parent.frame._name:
# We can't push the filter through the filter condition
return
parents = [x() for x in dependents[self._name] if x() is not None]
if not all(isinstance(p, Filter) for p in parents):
return
Expand Down
8 changes: 8 additions & 0 deletions dask_expr/tests/test_collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -2293,3 +2293,11 @@ def test_axes(df, pdf):
[assert_eq(d, p) for d, p in zip(df.axes, pdf.axes)]
assert len(df.x.axes) == len(pdf.x.axes)
assert_eq(df.x.axes[0], pdf.x.axes[0])


def test_filter_optimize_condition():
pdf = pd.DataFrame({"a": [1, 2, 3, 4], "b": [True, False, True, False]})
df = from_pandas(pdf, npartitions=2)
result = df[df.b.fillna(True)]
expected = pdf[pdf.b.fillna(True)]
assert_eq(result, expected)

0 comments on commit 6c6dc46

Please sign in to comment.