Dask optimization of [merge_asof + drop] graphs removes columns prematurely
- Dominant language
- Python
- Stars
- 13.9k
- Forks
- 2k
- PR merge metrics
- No merged PRs in 30d
Description
**Dask optimization of merge_asof + drop graphs removes columns prematurely**:
When computing a simple Dask graph of a `merge_asof` followed by a `drop` of a column used as key in the `merge_asof`, Dask seems to be overoptimizing the graph and dropping the column before the merge operation is carried out.
**Minimal Complete Verifiable Example**:
(jointly developed with [ramonruiz97](https://github.com/ramonruiz97))
```python
import dask
import dask.dataframe as dd
import pandas as pd
print("Dask:", dask.__version__)
left = dd.from_pandas(
pd.DataFrame(
{
"timestamp": pd.to_datetime(["2024-01-02", "2024-01-03", "2024-01-04"]),
"key": [1, 1, 2],
}
).set_index("timestamp"),
npartitions=1,
)
right = dd.from_pandas(
pd.DataFrame(
{
"timestamp_right": pd.to_datetime(["2024-01-01", "2024-01-02", "2024-01-03"]),
"key_right": [1, 1, 2],
"value": ["a", "b", "c"],
}
).set_index("timestamp_right"),
npartitions=1,
)
merged = dd.merge_asof(
left=left,
right=right,
left_index=True,
right_index=True,
left_by="key",
right_by="key_right",
direction="backward",
)
print(merged.compute())
dropped = merged.drop(columns=["key_right"])
print(dropped.compute())
```
Output
```
Dask: 2026.6.0
key key_right value
timestamp
2024-01-02 1 1 b
2024-01-03 1 1 b
2024-01-04 2 2 c
---------------------------------------------------------------------------
KeyError Traceback (most recent call last)
Cell In[2], [line 41](vscode-notebook-cell:?execution_count=2&line=41)
37
38 print(merged.compute())
39
40 dropped = merged.drop(columns=["key_right"])
---> [41](vscode-notebook-cell:?execution_count=2&line=41) print(dropped.compute())
File ~/$USER_PATH/lib/python3.12/site-packages/dask/base.py:377, in DaskMethodsMixin.compute(self, **kwargs)
353 def compute(self, **kwargs):
354 """Compute this dask collection
355
356 This turns a lazy Dask collection into its in-memory equivalent.
(...) 375 dask.compute
376 """
--> [377](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/base.py:377) (result,) = compute(self, traverse=False, **kwargs)
378 return result
File ~/$USER_PATH/lib/python3.12/site-packages/dask/base.py:682, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
663 expr = FinalizeCompute(expr)
665 with shorten_traceback():
666 # The high level optimize will have to be called client side (for now)
667 # The optimize can internally trigger already a computation
(...) 679 # change the graph submission to a handshake which introduces all sorts
680 # of concurrency control issues)
--> [682](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/base.py:682) expr = expr.optimize()
683 keys = list(flatten(expr.__dask_keys__()))
685 results = schedule(expr, keys, **kwargs)
File ~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:435, in Expr.optimize(self, fuse)
432 def optimize(self, fuse: bool = False) -> Expr:
433 stage: OptimizerStage = "fused" if fuse else "simplified-physical"
--> [435](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:435) return optimize_until(self, stage)
File ~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:929, in optimize_until(expr, stage)
926 return result
928 # Simplify
--> [929](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:929) expr = result.simplify()
930 if stage == "simplified-logical":
931 return expr
File ~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:445, in Expr.simplify(self)
443 while True:
444 dependents = collect_dependents(expr)
--> [445](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:445) new = expr.simplify_once(dependents=dependents, simplified={})
446 if new._name == expr._name:
447 break
File ~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:415, in Expr.simplify_once(self, dependents, simplified)
412 if isinstance(operand, Expr):
413 # Bandaid for now, waiting for Singleton
414 dependents[operand._name].append(weakref.ref(expr))
--> [415](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:415) new = operand.simplify_once(
416 dependents=dependents, simplified=simplified
417 )
418 simplified[operand._name] = new
419 if new._name != operand._name:
File ~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:415, in Expr.simplify_once(self, dependents, simplified)
412 if isinstance(operand, Expr):
413 # Bandaid for now, waiting for Singleton
414 dependents[operand._name].append(weakref.ref(expr))
--> [415](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:415) new = operand.simplify_once(
416 dependents=dependents, simplified=simplified
417 )
418 simplified[operand._name] = new
419 if new._name != operand._name:
File ~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:388, in Expr.simplify_once(self, dependents, simplified)
385 expr = self
387 while True:
--> [388](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/_expr.py:388) out = expr._simplify_down()
389 if out is None:
390 out = expr
File ~/$USER_PATH/lib/python3.12/site-packages/dask/dataframe/dask_expr/_expr.py:2227, in Projection._simplify_down(self)
2225 def _simplify_down(self):
2226 if (
-> [2227](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/dataframe/dask_expr/_expr.py:2227) str(self.frame.columns) == str(self.columns)
2228 and self._meta.ndim == self.frame._meta.ndim
2229 ):
2230 # TODO: we should get more precise around Expr.columns types
2231 return self.frame
2232 if isinstance(self.frame, Projection):
2233 # df[a][b]
File ~/$USER_PATH/lib/python3.12/site-packages/dask/dataframe/dask_expr/_expr.py:[451](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/dataframe/dask_expr/_expr.py:451), in Expr.columns(self)
448 @property
449 def columns(self) -> list:
450 try:
--> 451 return list(self._meta.columns)
452 except AttributeError:
453 if self.ndim == 1:
File ~/$USER_PATH/lib/python3.12/functools.py:[998](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/functools.py:998), in cached_property.__get__(self, instance, owner)
996 val = cache.get(self.attrname, _NOT_FOUND)
997 if val is _NOT_FOUND:
--> 998 val = self.func(instance)
999 try:
1000 cache[self.attrname] = val
File ~/$USER_PATH/lib/python3.12/site-packages/dask/dataframe/dask_expr/_merge_asof.py:[83](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/dask/dataframe/dask_expr/_merge_asof.py:83), in MergeAsof._meta(self)
80 @functools.cached_property
81 def _meta(self):
82 return make_meta(
---> 83 pd.merge_asof(
84 meta_nonempty(self.left._meta),
85 meta_nonempty(self.right._meta),
86 **self._kwargs,
87 )
88 )
File ~/$USER_PATH/lib/python3.12/site-packages/pandas/core/generic.py:1914, in NDFrame._get_label_or_level_values(self, key, axis)
1910 values = self.xs(key, axis=other_axes[0])._values
1911 elif self._is_level_reference(key, axis=axis):
1912 values = self.axes[axis].get_level_values(key)._values
1913 else:
-> [1914](https://vscode-remote+wsl-002bubuntu-002d26-002e04.vscode-resource.vscode-cdn.net/$USER_PATH/~/$USER_PATH/lib/python3.12/site-packages/pandas/core/generic.py:1914) raise KeyError(key)
1915
1916 # Check for duplicates
1917 if values.ndim > 1:
KeyError: 'key_right'
```
The problem vanishes when the Dask graph optimization is prevented by introducing a `persist` between the `merge_asof` and `drop` operations:
```python
dropped = merged.persist().drop(columns=["key_right"])
print(dropped.compute())
```
Output
```
key value
timestamp
2024-01-02 1 b
2024-01-03 1 b
2024-01-04 2 c
```
Another work-around is to run the drop operation at Pandas level using `map_partitions`, which I guess also prevents Dask from running its graph optimizations:
```python
meta = merged._meta.drop(columns=["key_right"])
dropped = merged.map_partitions(lambda part: part.drop(columns=["key_right"]), meta=meta)
print(dropped.compute())
```
Output
```
key value
timestamp
2024-01-02 1 b
2024-01-03 1 b
2024-01-04 2 c
```
I also tested this in an older dask version (`2024.12.1`) and everything worked as expected.
**Environment**:
- Dask version: `2026.6.0`
- Python version: `3.12.13`
- Operating System: `Ubuntu 26.04 LTS`
- Install method (conda, pip, source): `pip`
Contributor guide
Research direction
Start with the failing merge_asof/drop example and inspect dask/dataframe/dask_expr/_expr.py, especially Projection._simplify_down, then compare the metadata flow in dask/dataframe/dask_expr/_merge_asof.py. Add a regression test covering graph optimization without persist; done means dropping key_right computes successfully and returns the expected columns and values.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- pandas, python
- Domain
- data-engineering
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100