Missing index columns after upgrade to dask-expr

[aimran-adroll]

Not sure if this qualifies as bug report or not. But we recently bumped dask (2023.12.0 -> 2024.9.1). We are now running into quite a bit of crashes involving missing index keys (usually preceded by read parquet + a bunch of groupby agg/apply).

Here is an example error msg:
KeyError: "None of [Index(['advertisable_eid', 'user', 'subset_id'], dtype='object')] are in the [columns]"

If I step through the code, I can see where its failing. But if I dump out the columns, the dataframe spits out the right columns. Debugging is complicated by the fact that the job only crashes on distributed workload (for now).

Curiously, adding a few strategic df = df.optimize() right after the groupby operation mysteriously resolves the issue. I only accidentally came across the solution and have no idea why it helps.

I am filing it here because the crash logs are filled with calls to this module (esp the fuse step). If there are any tips or tricks to get more debugging traces out of the system, kindly share 🙏🏽