Thanks for pushing something up @madsbk . Having something concrete makes it a lot easier to think about this problem.

For the process of clearing out high level layers that aren't being used, this is clearly a win.

For lower level culling I still don't know exactly how to handle culling of the graphs within layers. Eventually, we're going to want to do this on the scheduler side. That means that we're going to want to cull the concrete low level graphs that we have here on the client side (the layers backed by dicts) but not explicitly expand the layers that have high level representations, like Blockwise, or a future Shuffle layer. We're going to want to keep those in abstract form so that we can pass them to the Scheduler.

With that in mind, I wonder if we maybe want the following two functions on every layer:

1. Given the keys that we are being asked to produce, generate the task graph that corresponds to those keys.

   This would be mildly more efficient than today's generate-then-cull even in the early case when everything is on the client side.

2. Given the keys that we are being asked to produce, generate the keys that we will need from each of our dependency layers.

   This will allow us to pass down enough information to continue the local culling process on dict-backed layers, but will not force us to go through the generation/culling step on the client side just yet.

As a test case we might create these methods on the ReadParquet and then the Blockwise mapping classes. The ReadParquet class is probably very easy and the Blockwise class is probably very hard. They might make good test cases though. To verify effectiveness we could do something like

df = dd.read_parquet("...")
df.head()

And then follow things through and see that we never actually generated the full graph. We could then expand this to ...

df = dd.read_parquet("...")
df["z"] = df.x + df.y
df.head()

And hopefully see the same thing, even now that there are now two layers and also a Blockwise in there.

This is just my impressions after looking at this code though. There may be better approaches here.

@madsbk I kicked CI to re-run the tests

Thanks @madsbk ! This is in.

Woohoo - thanks @madsbk! And thanks @mrocklin for reviewing :)

@madsbk and @mrocklin this PR seems to have broken some custom graph construction which was working for me prior to 2.28.0. If optimize_graph is True I now get errors like AttributeError: 'dict' object has no attribute 'discard'. I will likely look into raising a bug report tomorrow, but before I do, I was wondering if either of you might have an intuition regarding what could be going wrong? My instinct is that there is a dictionary where there should be a set, but I am unsure how these changes would have exposed this, given that it definitely worked prior to 2.28.0.