diff --git a/ibis/backends/dask/execution/reductions.py b/ibis/backends/dask/execution/reductions.py index fd15b3cc78dc..f18bc75437c9 100644 --- a/ibis/backends/dask/execution/reductions.py +++ b/ibis/backends/dask/execution/reductions.py @@ -176,3 +176,23 @@ def execute_standard_dev_series(op, data, mask, aggcontext=None, **kwargs): 'std', ddof=variance_ddof[op.how], ) + + +@execute_node.register( + ops.ArgMax, dd.Series, dd.Series, (dd.Series, type(None)) +) +def execute_argmax_series(op, data, key, mask, aggcontext=None, **kwargs): + idxmax = aggcontext.agg( + key[mask] if mask is not None else key, 'idxmax' + ).compute() + return data.loc[idxmax] + + +@execute_node.register( + ops.ArgMin, dd.Series, dd.Series, (dd.Series, type(None)) +) +def execute_argmin_series(op, data, key, mask, aggcontext=None, **kwargs): + idxmin = aggcontext.agg( + key[mask] if mask is not None else key, 'idxmin' + ).compute() + return data.loc[idxmin]