Skip to content

Commit

Permalink
feat(ux): expose between, rows and range keyword arguments in `…
Browse files Browse the repository at this point in the history
…value.over()`
  • Loading branch information
kszucs authored and cpcloud committed Feb 6, 2023
1 parent d85fb28 commit 5763063
Show file tree
Hide file tree
Showing 3 changed files with 67 additions and 5 deletions.
30 changes: 26 additions & 4 deletions ibis/expr/types/generic.py
Expand Up @@ -319,13 +319,29 @@ def substitute(

return expr.else_(else_ if else_ is not None else self).end()

def over(self, window) -> Value:
def over(
self,
window=None,
*,
rows=None,
range=None,
group_by=None,
order_by=None,
) -> Value:
"""Construct a window expression.
Parameters
----------
window
Window specification
rows
Whether to use the `ROWS` window clause
range
Whether to use the `RANGE` window clause
group_by
Grouping key
order_by
Ordering key
Returns
-------
Expand All @@ -336,13 +352,19 @@ def over(self, window) -> Value:
import ibis.expr.builders as bl
import ibis.expr.deferred as de

op = self.op()
if window is None:
window = ibis.window(
rows=rows,
range=range,
group_by=group_by,
order_by=order_by,
)

def bind(table):
frame = window.bind(table)
node = ops.WindowFunction(self, frame)
return node.to_expr()
return ops.WindowFunction(self, frame).to_expr()

op = self.op()
if isinstance(op, ops.Alias):
return op.arg.to_expr().over(window).name(op.name)
elif isinstance(op, ops.WindowFunction):
Expand Down
27 changes: 26 additions & 1 deletion ibis/expr/types/groupby.py
Expand Up @@ -19,6 +19,7 @@
import types
from typing import Iterable, Sequence

import ibis
import ibis.expr.analysis as an
import ibis.expr.operations as ops
import ibis.expr.types as ir
Expand Down Expand Up @@ -222,19 +223,43 @@ def _get_window(self):
order_by=self._window.order_by + self._order_by,
)

def over(self, window) -> GroupedTable:
def over(
self,
window=None,
*,
rows=None,
range=None,
group_by=None,
order_by=None,
) -> GroupedTable:
"""Apply a window over the input expressions.
Parameters
----------
window
Window to add to the input
rows
Whether to use the `ROWS` window clause
range
Whether to use the `RANGE` window clause
group_by
Grouping key
order_by
Ordering key
Returns
-------
GroupedTable
A new grouped table expression
"""
if window is None:
window = ibis.window(
rows=rows,
range=range,
group_by=group_by,
order_by=order_by,
)

return self.__class__(
self.table,
self.by,
Expand Down
15 changes: 15 additions & 0 deletions ibis/tests/expr/test_window_functions.py
Expand Up @@ -38,3 +38,18 @@ def test_mutate_with_analytic_functions(alltypes):
for field in proj.op().selections[1:]:
assert isinstance(field, ops.Alias)
assert isinstance(field.arg, ops.WindowFunction)


def test_value_over_api(alltypes):
t = alltypes

w1 = ibis.window(rows=(0, 1), group_by=t.g, order_by=[t.f, t.h])
w2 = ibis.window(range=(-1, 1), group_by=[t.g, t.a], order_by=[t.f])

expr = t.f.cumsum().over(rows=(0, 1), group_by=t.g, order_by=[t.f, t.h])
expected = t.f.cumsum().over(w1)
assert expr.equals(expected)

expr = t.f.cumsum().over(range=(-1, 1), group_by=[t.g, t.a], order_by=[t.f])
expected = t.f.cumsum().over(w2)
assert expr.equals(expected)

0 comments on commit 5763063

Please sign in to comment.