Skip to content

Commit

Permalink
Introduce more caching when walking the expression (#1165)
Browse files Browse the repository at this point in the history
  • Loading branch information
phofl authored Nov 14, 2024
1 parent bb0aef0 commit 095b665
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 5 deletions.
21 changes: 17 additions & 4 deletions dask_expr/_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,17 +160,26 @@ def __reduce__(self):
raise RuntimeError(f"Serializing a {type(self)} object")
return type(self), tuple(self.operands)

def _depth(self):
def _depth(self, cache=None):
"""Depth of the expression tree
Returns
-------
depth: int
"""
if cache is None:
cache = {}
if not self.dependencies():
return 1
else:
return max(expr._depth() for expr in self.dependencies()) + 1
result = []
for expr in self.dependencies():
if expr._name in cache:
result.append(cache[expr._name])
else:
result.append(expr._depth(cache) + 1)
cache[expr._name] = result[-1]
return max(result)

def operand(self, key):
# Access an operand unambiguously
Expand Down Expand Up @@ -242,7 +251,7 @@ def _layer(self) -> dict:
for i in range(self.npartitions)
}

def rewrite(self, kind: str):
def rewrite(self, kind: str, rewritten):
"""Rewrite an expression
This leverages the ``._{kind}_down`` and ``._{kind}_up``
Expand All @@ -255,6 +264,9 @@ def rewrite(self, kind: str):
changed:
whether or not any change occured
"""
if self._name in rewritten:
return rewritten[self._name]

expr = self
down_name = f"_{kind}_down"
up_name = f"_{kind}_up"
Expand Down Expand Up @@ -291,7 +303,8 @@ def rewrite(self, kind: str):
changed = False
for operand in expr.operands:
if isinstance(operand, Expr):
new = operand.rewrite(kind=kind)
new = operand.rewrite(kind=kind, rewritten=rewritten)
rewritten[operand._name] = new
if new._name != operand._name:
changed = True
else:
Expand Down
2 changes: 1 addition & 1 deletion dask_expr/_expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -3053,7 +3053,7 @@ def optimize_until(expr: Expr, stage: core.OptimizerStage) -> Expr:
return expr

# Manipulate Expression to make it more efficient
expr = expr.rewrite(kind="tune")
expr = expr.rewrite(kind="tune", rewritten={})
if stage == "tuned-logical":
return expr

Expand Down

0 comments on commit 095b665

Please sign in to comment.