Add caching to recursive simplify_once calls

dask_expr/_core.py

@@ -30,6 +30,7 @@ class Expr:
     _defaults = {}
 
     def __init__(self, *args, **kwargs):
+        self._simplified = {}
         operands = list(args)
         for parameter in type(self)._parameters[len(operands) :]:
             try:
@@ -279,6 +280,11 @@ def simplify_once(self, dependents: defaultdict):
         expr:
             output expression
         """
+        # Check if we've already simplified for these dependents
+        key = _tokenize_deterministic(sorted(dependents.keys()))
+        if key in self._simplified:
+            return self._simplified[key]
+
         expr = self
 
         while True:
@@ -319,6 +325,7 @@ def simplify_once(self, dependents: defaultdict):
 
             break
 
+        self._simplified[key] = expr  # Cache the result
         return expr
 
     def simplify(self) -> Expr:

dask_expr/_expr.py

@@ -1768,9 +1768,7 @@ def vals(self):
 
     @functools.cached_property
     def _meta(self):
-        args = [
-            meta_nonempty(op._meta) if isinstance(op, Expr) else op for op in self._args
-        ]
+        args = [op._meta if isinstance(op, Expr) else op for op in self._args]
         return make_meta(self.operation(*args, **self._kwargs))
 
     def _tree_repr_argument_construction(self, i, op, header):

dask_expr/io/_delayed.py

@@ -22,6 +22,7 @@ class _DelayedExpr(Expr):
     # TODO
 
     def __init__(self, obj):
+        self._simplified = {}
         self.obj = obj
         self.operands = [obj]