diff --git a/ibis/expr/types/generic.py b/ibis/expr/types/generic.py index cef32748b470..8fc103a2e024 100644 --- a/ibis/expr/types/generic.py +++ b/ibis/expr/types/generic.py @@ -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 ------- @@ -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): diff --git a/ibis/expr/types/groupby.py b/ibis/expr/types/groupby.py index 4e61a71949e5..78d5a7a0b3d4 100644 --- a/ibis/expr/types/groupby.py +++ b/ibis/expr/types/groupby.py @@ -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 @@ -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, diff --git a/ibis/tests/expr/test_window_functions.py b/ibis/tests/expr/test_window_functions.py index 092b37d2e124..f1e2c9ab8e3c 100644 --- a/ibis/tests/expr/test_window_functions.py +++ b/ibis/tests/expr/test_window_functions.py @@ -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)