Add branch_id to distinguish between reusable branches and pipeline breakers

dask_expr/_core.py

@@ -5,6 +5,7 @@
 import weakref
 from collections import defaultdict
 from collections.abc import Generator
+from typing import NamedTuple
 
 import dask
 import pandas as pd
@@ -15,6 +16,10 @@
 from dask_expr._util import _BackendData, _tokenize_deterministic
 
 
+class BranchId(NamedTuple):
+    branch_id: int
+
+
 def _unpack_collections(o):
     if isinstance(o, Expr):
         return o
@@ -30,23 +35,36 @@ class Expr:
     _defaults = {}
     _instances = weakref.WeakValueDictionary()
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args, _branch_id=None, **kwargs):
         operands = list(args)
+        if _branch_id is None and len(operands) and isinstance(operands[-1], BranchId):
+            _branch_id = operands.pop(-1)
+        elif _branch_id is None:
+            _branch_id = BranchId(0)
+
         for parameter in cls._parameters[len(operands) :]:
             try:
                 operands.append(kwargs.pop(parameter))
             except KeyError:
                 operands.append(cls._defaults[parameter])
         assert not kwargs, kwargs
         inst = object.__new__(cls)
-        inst.operands = [_unpack_collections(o) for o in operands]
+        inst.operands = [_unpack_collections(o) for o in operands] + [_branch_id]
         _name = inst._name
         if _name in Expr._instances:
             return Expr._instances[_name]
 
         Expr._instances[_name] = inst
         return inst
 
+    @functools.cached_property
+    def argument_operands(self):
+        return self.operands[:-1]
+
+    @functools.cached_property
+    def _branch_id(self):
+        return self.operands[-1]
+
     def _tune_down(self):
         return None
 
@@ -107,6 +125,10 @@ def _tree_repr_lines(self, indent=0, recursive=True):
                     op = "<series>"
                 elif is_arraylike(op):
                     op = "<array>"
+                elif isinstance(op, BranchId):
+                    if op.branch_id == 0:
+                        continue
+                    op = f" branch_id={op.branch_id}"
                 header = self._tree_repr_argument_construction(i, op, header)
 
         lines = [header] + lines
@@ -203,7 +225,7 @@ def _layer(self) -> dict:
 
         return {(self._name, i): self._task(i) for i in range(self.npartitions)}
 
-    def rewrite(self, kind: str):
+    def rewrite(self, kind: str, cache):
         """Rewrite an expression
 
         This leverages the ``._{kind}_down`` and ``._{kind}_up``
@@ -216,6 +238,9 @@ def rewrite(self, kind: str):
         changed:
             whether or not any change occured
         """
+        if self._name in cache:
+            return cache[self._name]
+
         expr = self
         down_name = f"_{kind}_down"
         up_name = f"_{kind}_up"
@@ -252,7 +277,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, cache=cache)
+                    cache[operand._name] = new
                     if new._name != operand._name:
                         changed = True
                 else:
@@ -267,6 +293,28 @@ def rewrite(self, kind: str):
 
         return expr
 
+    def _reuse_up(self, parent):
+        return
+
+    def _reuse_down(self):
+        if not self.dependencies():
+            return
+        return self._bubble_branch_id_down()
+
+    def _bubble_branch_id_down(self):
+        b_id = self._branch_id
+        if any(b_id.branch_id != d._branch_id.branch_id for d in self.dependencies()):
+            ops = [
+                op._substitute_branch_id(b_id) if isinstance(op, Expr) else op
+                for op in self.argument_operands
+            ]
+            return type(self)(*ops)
+
+    def _substitute_branch_id(self, branch_id):
+        if self._branch_id.branch_id != 0:
+            return self
+        return type(self)(*self.argument_operands, branch_id)
+
     def simplify_once(self, dependents: defaultdict, simplified: dict):
         """Simplify an expression
 

dask_expr/_cumulative.py

@@ -47,7 +47,7 @@ def operation(self):
 
     @functools.cached_property
     def _args(self) -> list:
-        return self.operands[:-1]
+        return self.argument_operands[:-1]
 
 
 class TakeLast(Blockwise):

dask_expr/_expr.py

@@ -53,6 +53,7 @@
 from tlz import merge_sorted, partition, unique
 
 from dask_expr import _core as core
+from dask_expr._core import BranchId
 from dask_expr._util import (
     _calc_maybe_new_divisions,
     _convert_to_list,
@@ -476,9 +477,9 @@ def _args(self) -> list:
         if self._keyword_only:
             args = [
                 self.operand(p) for p in self._parameters if p not in self._keyword_only
-            ] + self.operands[len(self._parameters) :]
+            ] + self.argument_operands[len(self._parameters) :]
             return args
-        return self.operands
+        return self.argument_operands
 
     def _broadcast_dep(self, dep: Expr):
         # Checks if a dependency should be broadcasted to
@@ -562,7 +563,7 @@ def _broadcast_dep(self, dep: Expr):
 
     @property
     def args(self):
-        return [self.frame] + self.operands[len(self._parameters) :]
+        return [self.frame] + self.argument_operands[len(self._parameters) :]
 
     @functools.cached_property
     def _meta(self):
@@ -658,7 +659,7 @@ def __str__(self):
 
     @functools.cached_property
     def args(self):
-        return self.operands[len(self._parameters) :]
+        return self.argument_operands[len(self._parameters) :]
 
     @functools.cached_property
     def _dfs(self):
@@ -725,7 +726,7 @@ def _meta(self):
         meta = self.operand("meta")
         args = [self.frame._meta] + [
             arg._meta if isinstance(arg, Expr) else arg
-            for arg in self.operands[len(self._parameters) :]
+            for arg in self.argument_operands[len(self._parameters) :]
         ]
         return _get_meta_map_partitions(
             args,
@@ -737,7 +738,7 @@ def _meta(self):
         )
 
     def _divisions(self):
-        args = [self.frame] + self.operands[len(self._parameters) :]
+        args = [self.frame] + self.argument_operands[len(self._parameters) :]
         return calc_divisions_for_align(*args)
 
     def _lower(self):
@@ -792,15 +793,15 @@ def args(self):
         return (
             [self.frame]
             + [self.func, self.before, self.after]
-            + self.operands[len(self._parameters) :]
+            + self.argument_operands[len(self._parameters) :]
         )
 
     @functools.cached_property
     def _meta(self):
         meta = self.operand("meta")
         args = [self.frame._meta] + [
             arg._meta if isinstance(arg, Expr) else arg
-            for arg in self.operands[len(self._parameters) :]
+            for arg in self.argument_operands[len(self._parameters) :]
         ]
         return _get_meta_map_partitions(
             args,
@@ -1094,7 +1095,11 @@ class Sample(Blockwise):
 
     @functools.cached_property
     def _meta(self):
-        args = [self.operands[0]._meta] + [self.operands[1][0]] + self.operands[2:]
+        args = (
+            [self.operands[0]._meta]
+            + [self.operands[1][0]]
+            + self.argument_operands[2:]
+        )
         return self.operation(*args)
 
     def _task(self, index: int):
@@ -1696,11 +1701,11 @@ class Assign(Elemwise):
 
     @functools.cached_property
     def keys(self):
-        return self.operands[1::2]
+        return self.argument_operands[1::2]
 
     @functools.cached_property
     def vals(self):
-        return self.operands[2::2]
+        return self.argument_operands[2::2]
 
     @functools.cached_property
     def _meta(self):
@@ -1725,7 +1730,7 @@ def _simplify_down(self):
             if self._check_for_previously_created_column(self.frame):
                 # don't squash if we are using a column that was previously created
                 return
-            return Assign(*self.frame.operands, *self.operands[1:])
+            return Assign(*self.frame.argument_operands, *self.operands[1:])
 
     def _check_for_previously_created_column(self, child):
         input_columns = []
@@ -1753,7 +1758,7 @@ def _simplify_up(self, parent, dependents):
                     if k in columns:
                         new_args.extend([k, v])
             else:
-                new_args = self.operands[1:]
+                new_args = self.argument_operands[1:]
 
             columns = [col for col in self.frame.columns if col in cols]
             return type(parent)(
@@ -1778,12 +1783,12 @@ class CaseWhen(Elemwise):
 
     @functools.cached_property
     def caselist(self):
-        c = self.operands[1:]
+        c = self.argument_operands[1:]
         return [(c[i], c[i + 1]) for i in range(0, len(c), 2)]
 
     @functools.cached_property
     def _meta(self):
-        c = self.operands[1:]
+        c = self.argument_operands[1:]
         caselist = [
             (
                 meta_nonempty(c[i]._meta) if isinstance(c[i], Expr) else c[i],
@@ -2714,9 +2719,11 @@ class _DelayedExpr(Expr):
     # TODO
     _parameters = ["obj"]
 
-    def __init__(self, obj):
+    def __init__(self, obj, _branch_id=None):
         self.obj = obj
-        self.operands = [obj]
+        if _branch_id is None:
+            _branch_id = BranchId(0)
+        self.operands = [obj, _branch_id]
 
     def __str__(self):
         return f"{type(self).__name__}({str(self.obj)})"
@@ -2744,7 +2751,9 @@ def normalize_expression(expr):
     return expr._name
 
 
-def optimize(expr: Expr, fuse: bool = True) -> Expr:
+def optimize(
+    expr: Expr, fuse: bool = True, common_subplan_elimination: bool = True
+) -> Expr:
     """High level query optimization
 
     This leverages three optimization passes:
@@ -2758,24 +2767,37 @@ def optimize(expr: Expr, fuse: bool = True) -> Expr:
         Input expression to optimize
     fuse:
         whether or not to turn on blockwise fusion
+    common_subplan_elimination : bool
+        whether we want to reuse common subplans that are found in the graph and
+        are used in self-joins or similar which require all data be held in memory
+        at some point. Only set this to false if your dataset fits into memory.
 
     See Also
     --------
     simplify
     optimize_blockwise_fusion
     """
+    result = expr
+    while True:
+        if common_subplan_elimination:
+            out = result.rewrite("reuse", cache={})
+        else:
+            out = result
+        out = out.simplify()
+        if out._name == result._name or not common_subplan_elimination:
+            break
+        result = out
 
-    # Simplify
-    result = expr.simplify()
+    result = out
 
     # Manipulate Expression to make it more efficient
-    result = result.rewrite(kind="tune")
+    result = result.rewrite(kind="tune", cache={})
 
     # Lower
     result = result.lower_completely()
 
     # Cull
-    result = result.rewrite(kind="cull")
+    result = result.rewrite(kind="cull", cache={})
 
     # Final graph-specific optimizations
     if fuse:
@@ -3307,7 +3329,7 @@ def __str__(self):
 
     @functools.cached_property
     def args(self):
-        return self.operands[len(self._parameters) :]
+        return self.argument_operands[len(self._parameters) :]
 
     @functools.cached_property
     def _dfs(self):

dask_expr/_groupby.py

@@ -106,7 +106,7 @@ def split_by(self):
 
     @functools.cached_property
     def by(self):
-        return self.operands[len(self._parameters) :]
+        return self.argument_operands[len(self._parameters) :]
 
     @functools.cached_property
     def levels(self):
@@ -741,6 +741,7 @@ def _lower(self):
                 for param in self._parameters
             ],
             *self.by,
+            self._branch_id,
         )
         if is_dataframe_like(s._meta):
             c = c[s.columns]