consolidate trailing window functions

ibis/expr/tests/test_window_functions.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import pytest
 
 import ibis
+from ibis.expr.window import _determine_how
 from ibis.tests.util import assert_equal
 
 
@@ -174,3 +176,35 @@ def test_preceding_following_validate(alltypes):
 @pytest.mark.xfail(raises=AssertionError, reason='NYT')
 def test_window_equals(alltypes):
     assert False
+
+
+def test_determine_how():
+    how = _determine_how((None, 5))
+    assert how == 'rows'
+
+    how = _determine_how((3, 1))
+    assert how == 'rows'
+
+    how = _determine_how(5)
+    assert how == 'rows'
+
+    how = _determine_how(np.int64(7))
+    assert how == 'rows'
+
+    how = _determine_how(ibis.interval(days=3))
+    assert how == 'range'
+
+    how = _determine_how(ibis.interval(months=5) + ibis.interval(days=10))
+    assert how == 'range'
+
+    with pytest.raises(TypeError):
+        _determine_how(8.9)
+
+    with pytest.raises(TypeError):
+        _determine_how('invalid preceding')
+
+    with pytest.raises(TypeError):
+        _determine_how({'start': 1, 'end': 2})
+
+    with pytest.raises(TypeError):
+        _determine_how([3, 5])

ibis/expr/window.py

@@ -1,5 +1,7 @@
 """Encapsulation of SQL window clauses."""
 
+import numpy as np
+
 import ibis.common as com
 import ibis.expr.operations as ops
 import ibis.expr.types as ir
@@ -10,6 +12,29 @@ def _sequence_to_tuple(x):
     return tuple(x) if util.is_iterable(x) else x
 
 
+def _determine_how(preceding):
+    if isinstance(preceding, tuple):
+        start, end = preceding
+        if start is None:
+            offset_type = type(end)
+        else:
+            offset_type = type(start)
+    else:
+        offset_type = type(preceding)
+
+    if issubclass(offset_type, (int, np.integer)):
+        how = 'rows'
+    elif issubclass(offset_type, ir.IntervalScalar):
+        how = 'range'
+    else:
+        raise TypeError(
+            'Type {} is not supported for row- or range- based trailing '
+            'window operations'.format(offset_type)
+        )
+
+    return how
+
+
 class Window:
     """Class to encapsulate the details of a window frame.
 
@@ -310,13 +335,16 @@ def cumulative_window(group_by=None, order_by=None):
     )
 
 
-def trailing_window(rows, group_by=None, order_by=None):
+def trailing_window(preceding, group_by=None, order_by=None):
     """Create a trailing window for use with aggregate window functions.
 
     Parameters
     ----------
-    rows : int
-        Number of trailing rows to include. 0 includes only the current row
+    preceding : int, float or expression of intervals, i.e.
+        ibis.interval(days=1) + ibis.interval(hours=5)
+        Int indicates number of trailing rows to include;
+        0 includes only the current row.
+        Interval indicates a trailing range window.
     group_by : expressions, default None
         Either specify here or with TableExpr.group_by
     order_by : expressions, default None
@@ -328,8 +356,13 @@ def trailing_window(rows, group_by=None, order_by=None):
     Window
 
     """
+    how = _determine_how(preceding)
     return Window(
-        preceding=rows, following=0, group_by=group_by, order_by=order_by
+        preceding=preceding,
+        following=0,
+        group_by=group_by,
+        order_by=order_by,
+        how=how
     )