Skip to content

Commit

Permalink
feat(dask): implement ungrouped argmin and argmax
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Aug 20, 2022
1 parent bf9b948 commit 854aea7
Showing 1 changed file with 20 additions and 0 deletions.
20 changes: 20 additions & 0 deletions ibis/backends/dask/execution/reductions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

0 comments on commit 854aea7

Please sign in to comment.