Skip to content

Commit

Permalink
cache simplified
Browse files Browse the repository at this point in the history
  • Loading branch information
fjetter committed Jan 23, 2024
1 parent cae4df1 commit cc1470c
Showing 1 changed file with 8 additions and 2 deletions.
10 changes: 8 additions & 2 deletions dask_expr/_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ def __new__(cls, *args, **kwargs):
return Expr._instances[_name]

Expr._instances[_name] = inst
inst.fully_simplified = False
if inst._required_attribute:
dep = next(iter(_dependencies))._meta
if not hasattr(dep, inst._required_attribute):
Expand Down Expand Up @@ -299,9 +300,12 @@ def simplify_once(self, dependents: defaultdict):
expr:
output expression
"""
if self.fully_simplified:
return self
expr = self

n = 0
while True:
n += 1
out = expr._simplify_down()
if out is None:
out = expr
Expand Down Expand Up @@ -338,7 +342,8 @@ def simplify_once(self, dependents: defaultdict):
expr = type(expr)(*new_operands)

break

if expr is self:
expr.fully_simplified = True
return expr

def simplify(self) -> Expr:
Expand All @@ -349,6 +354,7 @@ def simplify(self) -> Expr:
if new._name == expr._name:
break
expr = new
# print(f"Cache has size {len(_seen)=}")
return expr

def _simplify_down(self):
Expand Down

0 comments on commit cc1470c

Please sign in to comment.