After profiling this I am confused why this kind of cache would help us anything. Everything that lights up the profile is essentially a computation of _meta. Some of them can be optimized (see also #796 (comment)) but the fundamental problem is that we're generating a lot of intermediate expressions and we have to compute _meta for almost all of them which is unfortunately not always trivial.

How is this cache helping with this? I do indeed see fewer intermediates but I don't understand how this changes anything about the performance. I ran a test where I keep the cache but never access the element

diff --git a/dask_expr/_core.py b/dask_expr/_core.py
index 3388694..86ae404 100644
--- a/dask_expr/_core.py
+++ b/dask_expr/_core.py
@@ -282,8 +282,8 @@ class Expr:
         """
         # Check if we've already simplified for these dependents
         key = _tokenize_deterministic(sorted(dependents.keys()))
-        if key in self._simplified:
-            return self._simplified[key]
+        # if key in self._simplified:
+        #     return self._simplified[key]

         expr = self

@@ -315,6 +315,8 @@ class Expr:
                 if isinstance(operand, Expr):
                     new = operand.simplify_once(dependents=dependents)
                     if new._name != operand._name:
+                        if key in self._simplified:
+                            print("Already simplified but changed")
                         changed = True
                 else:
                     new = operand

and indeed... I get a lot of "Already simplified but changed" prints so even though the expression already passed once through this, there is still more work to do apparently. This cache prohibits that work which is why it's so much faster. However, will this yield the same results?

@fjetter - Thanks for looking into this.

and indeed... I get a lot of "Already simplified but changed" prints so even though the expression already passed once through this, there is still more work to do apparently. This cache prohibits that work which is why it's so much faster. However, will this yield the same results?

Right, my initial assumption here was that the new "dependents-informed" optimization loop was simply re-generating expressions that it didn't actually need to regenerate, and was therefore regenerating _meta that it had already generated before.

I'll admit that I didn't make a solid attempt to poke holes in my quick/simple "fix" yet. However, my assumption is that if our dependents graph has the same nodes, and our expression is the same, then we have already performed that "simplify" logic before. If something about the dependencies changed, then the node names would have changed.

I think this cache must only be populated if the expr is indeed identical in the end, so

diff --git a/dask_expr/_core.py b/dask_expr/_core.py
index 3388694..b3f70a7 100644
--- a/dask_expr/_core.py
+++ b/dask_expr/_core.py
@@ -324,8 +324,8 @@ class Expr:
                 expr = type(expr)(*new_operands)

             break
-
-        self._simplified[key] = expr  # Cache the result
+        if expr is self:
+            self._simplified[key] = expr  # Cache the result
         return expr

     def simplify(self) -> Expr:

with this I still get almost 3k cache hits

you could also check the name but you'd have to check if the expression is not modified. I think an equivalent would be

if expr._name == key:
    self._simplified[key] = expr

but we can't populate the cache immediately.

I tried the caching approach in #798 as well and all changes combined makes everything pleasantly fast.
Over there, I am using an attribute that's set on the expression that remembers whether an expression is fully simplified or not. This feels less invasive than a mapping on every instance

Over there, I am using an attribute that's set on the expression that remembers whether an expression is fully simplified or not. This feels less invasive than a mapping on every instance

I do like the feel of that better, but tests are failing. I think you really do need to keep track of whether the expression was simplified for the specific dependents in question (not just that the expression was simplified before and it didn't change).

I know this PR wasn't the best solution in its original form. However, I think it should be a reasonable solution for #796 in its revised form. The optimization time becomes negligible (<0.5s) for Q1 on my machine (a ~10x improvement).

Thanks for the review @phofl !

My other PR is still worthwhile btw, would appreciate if you could take a look there #842

Yes, I agree that's worthwhile as well. I'll definitely take a look.

thx