Skip to content

Commit

Permalink
feat(api): add argmax and argmin column methods
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Aug 20, 2022
1 parent ed91344 commit b52216a
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 0 deletions.
14 changes: 14 additions & 0 deletions ibis/expr/operations/reductions.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,20 @@ class Min(Filterable, Reduction):
output_dtype = rlz.dtype_like('arg')


@public
class ArgMax(Filterable, Reduction):
arg = rlz.column(rlz.any)
key = rlz.column(rlz.any)
output_dtype = rlz.dtype_like("arg")


@public
class ArgMin(Filterable, Reduction):
arg = rlz.column(rlz.any)
key = rlz.column(rlz.any)
output_dtype = rlz.dtype_like("arg")


@public
class ApproxCountDistinct(Filterable, Reduction):
"""Approximate number of unique values using HyperLogLog algorithm.
Expand Down
18 changes: 18 additions & 0 deletions ibis/expr/types/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -652,15 +652,33 @@ def approx_median(
return ops.ApproxMedian(self, where).to_expr().name("approx_median")

def max(self, where: ir.BooleanValue | None = None) -> Scalar:
"""Return the maximum of a column."""
import ibis.expr.operations as ops

return ops.Max(self, where).to_expr().name("max")

def min(self, where: ir.BooleanValue | None = None) -> Scalar:
"""Return the minimum of a column."""
import ibis.expr.operations as ops

return ops.Min(self, where).to_expr().name("min")

def argmax(
self, key: ir.Value, where: ir.BooleanValue | None = None
) -> Scalar:
"""Return the value of `self` that maximizes `key`."""
import ibis.expr.operations as ops

return ops.ArgMax(self, key=key, where=where).to_expr()

def argmin(
self, key: ir.Value, where: ir.BooleanValue | None = None
) -> Scalar:
"""Return the value of `self` that minimizes `key`."""
import ibis.expr.operations as ops

return ops.ArgMin(self, key=key, where=where).to_expr()

def nunique(
self, where: ir.BooleanValue | None = None
) -> ir.IntegerScalar:
Expand Down

0 comments on commit b52216a

Please sign in to comment.