Add window frame specificationpr/264

6 files changed, 151 insertions, 32 deletions

diff --git a/doc/build/core/tutorial.rst b/doc/build/core/tutorial.rst
index 0fd78abeb..7b28b1911 100644
--- a/doc/build/core/tutorial.rst
+++ b/doc/build/core/tutorial.rst
@@ -1423,7 +1423,22 @@ OVER clause, using the :meth:`.FunctionElement.over` method:
     ...         func.row_number().over(order_by=users.c.name)
     ...     ])
     >>> print(s)
-    SELECT users.id, row_number() OVER (ORDER BY users.name) AS anon_1
+    SELECT users.id, row_number() OVER (ORDER BY users.name RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS anon_1
+    FROM users
+
+.. versionadded:: 1.2
+
+   You can provide a window frame specification using either the ``rows`` or
+   ``range`` parameter:
+
+.. sourcecode:: pycon+sql
+
+    >>> s = select([
+    ...         users.c.id,
+    ...         func.row_number().over(order_by=users.c.name, rows={'preceding': 2, 'following': 3})
+    ...     ])
+    >>> print(s)
+    SELECT users.id, row_number() OVER (ORDER BY users.name ROWS BETWEEN 2 PRECEDING and 3 FOLLOWING) AS anon_1
     FROM users
 
 .. seealso::
diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index 3d2f02006..33de2424b 100644
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -812,12 +812,23 @@ class SQLCompiler(Compiled):
         return "%s OVER (%s)" % (
             over.element._compiler_dispatch(self, **kwargs),
             ' '.join(
-                '%s BY %s' % (word, clause._compiler_dispatch(self, **kwargs))
-                for word, clause in (
-                    ('PARTITION', over.partition_by),
-                    ('ORDER', over.order_by)
+                itertools.chain(
+                    (
+                        '%s BY %s' % (
+                            word, clause._compiler_dispatch(self, **kwargs)
+                        )
+                        for word, clause in (
+                            ('PARTITION', over.partition_by),
+                            ('ORDER', over.order_by)
+                        )
+                        if clause is not None and len(clause)
+                    ),
+                    [
+                        '%s BETWEEN %s AND %s' % (
+                            over.frame_kind, over.preceding, over.following
+                        )
+                    ]
                 )
-                if clause is not None and len(clause)
             )
         )
 
diff --git a/lib/sqlalchemy/sql/elements.py b/lib/sqlalchemy/sql/elements.py
index 00c2c37ba..52f46f295 100644
--- a/lib/sqlalchemy/sql/elements.py
+++ b/lib/sqlalchemy/sql/elements.py
@@ -3071,7 +3071,8 @@ class Over(ColumnElement):
     order_by = None
     partition_by = None
 
-    def __init__(self, element, partition_by=None, order_by=None):
+    def __init__(
+        self, element, partition_by=None, order_by=None, range=None, rows=None):
         """Produce an :class:`.Over` object against a function.
 
         Used against aggregate or so-called "window" functions,
@@ -3115,6 +3116,44 @@ class Over(ColumnElement):
                 *util.to_list(partition_by),
                 _literal_as_text=_literal_as_label_reference)
 
+        if range is not None and rows is not None:
+            raise ValueError(
+                'Must provide either RANGE specification or '
+                'ROWS specification, not both'
+            )
+
+        if range is None and rows is None:
+            range = {}
+
+        self.frame_kind = 'ROWS' if rows is not None else 'RANGE'
+        frame = rows if rows is not None else range
+
+        assert frame is not None
+
+        if frozenset(frame) - frozenset(('preceding', 'following')):
+            raise ValueError(
+                'Found additional keys in frame specification dict'
+            )
+        self.preceding = Over._format_frame_clause(
+            frame, 'preceding', 'UNBOUNDED PRECEDING'
+        )
+        self.following = Over._format_frame_clause(
+            frame, 'following', 'CURRENT ROW'
+        )
+
+    @staticmethod
+    def _format_frame_clause(frame, key, default):
+        value = frame.get(key, default)
+
+        if value == default:
+            return default
+        elif value is None:
+            return 'UNBOUNDED %s' % key.upper()
+        elif value == 0:
+            return 'CURRENT ROW'
+        else:
+            return '%d %s' % (value, key.upper())
+
     @property
     def func(self):
         """the element referred to by this :class:`.Over`
diff --git a/lib/sqlalchemy/sql/functions.py b/lib/sqlalchemy/sql/functions.py
index e369f5a61..a140544f0 100644
--- a/lib/sqlalchemy/sql/functions.py
+++ b/lib/sqlalchemy/sql/functions.py
@@ -94,7 +94,7 @@ class FunctionElement(Executable, ColumnElement, FromClause):
         """
         return self.clause_expr.element
 
-    def over(self, partition_by=None, order_by=None):
+    def over(self, partition_by=None, order_by=None, rows=None, range=None):
         """Produce an OVER clause against this function.
 
         Used against aggregate or so-called "window" functions,
@@ -114,7 +114,13 @@ class FunctionElement(Executable, ColumnElement, FromClause):
         .. versionadded:: 0.7
 
         """
-        return Over(self, partition_by=partition_by, order_by=order_by)
+        return Over(
+            self,
+            partition_by=partition_by,
+            order_by=order_by,
+            rows=rows,
+            range=range
+        )
 
     def within_group(self, *order_by):
         """Produce a WITHIN GROUP (ORDER BY expr) clause against this function.
diff --git a/test/sql/test_compiler.py b/test/sql/test_compiler.py
index dae178d31..2dffb427e 100644
--- a/test/sql/test_compiler.py
+++ b/test/sql/test_compiler.py
@@ -2196,20 +2196,23 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
     def test_over(self):
         self.assert_compile(
             func.row_number().over(),
-            "row_number() OVER ()"
+            "row_number() OVER "
+            "(RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)"
         )
         self.assert_compile(
             func.row_number().over(
                 order_by=[table1.c.name, table1.c.description]
             ),
-            "row_number() OVER (ORDER BY mytable.name, mytable.description)"
+            "row_number() OVER (ORDER BY mytable.name, mytable.description "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
+            ")"
         )
         self.assert_compile(
             func.row_number().over(
                 partition_by=[table1.c.name, table1.c.description]
             ),
-            "row_number() OVER (PARTITION BY mytable.name, "
-            "mytable.description)"
+            "row_number() OVER (PARTITION BY mytable.name, mytable.description"
+            " RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)"
         )
         self.assert_compile(
             func.row_number().over(
@@ -2217,7 +2220,8 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
                 order_by=[table1.c.description]
             ),
             "row_number() OVER (PARTITION BY mytable.name "
-            "ORDER BY mytable.description)"
+            "ORDER BY mytable.description "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)"
         )
         self.assert_compile(
             func.row_number().over(
@@ -2225,7 +2229,8 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
                 order_by=table1.c.description
             ),
             "row_number() OVER (PARTITION BY mytable.name "
-            "ORDER BY mytable.description)"
+            "ORDER BY mytable.description "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)"
         )
 
         self.assert_compile(
@@ -2234,7 +2239,8 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
                 order_by=[table1.c.name, table1.c.description]
             ),
             "row_number() OVER (PARTITION BY mytable.name "
-            "ORDER BY mytable.name, mytable.description)"
+            "ORDER BY mytable.name, mytable.description "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)"
         )
 
         self.assert_compile(
@@ -2242,7 +2248,9 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
                 partition_by=[],
                 order_by=[table1.c.name, table1.c.description]
             ),
-            "row_number() OVER (ORDER BY mytable.name, mytable.description)"
+            "row_number() OVER (ORDER BY mytable.name, mytable.description "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
+            ")"
         )
 
         self.assert_compile(
@@ -2251,7 +2259,8 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
                 order_by=[]
             ),
             "row_number() OVER (PARTITION BY mytable.name, "
-            "mytable.description)"
+            "mytable.description "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)"
         )
 
         self.assert_compile(
@@ -2259,14 +2268,15 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
                 partition_by=[],
                 order_by=[]
             ),
-            "row_number() OVER ()"
+            "row_number() OVER (RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)"
         )
         self.assert_compile(
             select([func.row_number().over(
                 order_by=table1.c.description
             ).label('foo')]),
-            "SELECT row_number() OVER (ORDER BY mytable.description) "
-            "AS foo FROM mytable"
+            "SELECT row_number() OVER (ORDER BY mytable.description "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
+            ") AS foo FROM mytable"
         )
 
         # test from_obj generation.
@@ -2277,8 +2287,9 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
                     partition_by=['description']
                 )
             ]),
-            "SELECT max(mytable.name) OVER (PARTITION BY mytable.description) "
-            "AS anon_1 FROM mytable"
+            "SELECT max(mytable.name) OVER (PARTITION BY mytable.description "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
+            ") AS anon_1 FROM mytable"
         )
         # from partition_by
         self.assert_compile(
@@ -2287,8 +2298,9 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
                     partition_by=[table1.c.name]
                 )
             ]),
-            "SELECT row_number() OVER (PARTITION BY mytable.name) "
-            "AS anon_1 FROM mytable"
+            "SELECT row_number() OVER (PARTITION BY mytable.name "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
+            ") AS anon_1 FROM mytable"
         )
         # from order_by
         self.assert_compile(
@@ -2297,15 +2309,17 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
                     order_by=table1.c.name
                 )
             ]),
-            "SELECT row_number() OVER (ORDER BY mytable.name) "
-            "AS anon_1 FROM mytable"
+            "SELECT row_number() OVER (ORDER BY mytable.name "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
+            ") AS anon_1 FROM mytable"
         )
 
         # this tests that _from_objects
         # concantenates OK
         self.assert_compile(
             select([column("x") + over(func.foo())]),
-            "SELECT x + foo() OVER () AS anon_1"
+            "SELECT x + foo() OVER "
+            "(RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS anon_1"
         )
 
         # test a reference to a label that in the referecned selectable;
@@ -2314,8 +2328,9 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
         stmt = select([expr]).alias()
         self.assert_compile(
             select([stmt.c.sum, func.row_number().over(order_by=stmt.c.sum)]),
-            "SELECT anon_1.sum, row_number() OVER (ORDER BY anon_1.sum) "
-            "AS anon_2 FROM (SELECT mytable.myid + :myid_1 AS sum "
+            "SELECT anon_1.sum, row_number() OVER (ORDER BY anon_1.sum "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
+            ") AS anon_2 FROM (SELECT mytable.myid + :myid_1 AS sum "
             "FROM mytable) AS anon_1"
         )
 
@@ -2326,7 +2341,39 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
             select([expr, func.row_number().over(order_by=expr)]),
             "SELECT mytable.myid + :myid_1 AS sum, "
             "row_number() OVER "
-            "(ORDER BY mytable.myid + :myid_1) AS anon_1 FROM mytable"
+            "(ORDER BY mytable.myid + :myid_1 "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
+            ") AS anon_1 FROM mytable"
+        )
+
+        # frame specification
+        expr = table1.c.myid
+        self.assert_compile(
+            select([func.row_number().over(order_by=expr, rows={'preceding': 0, 'following': None})]),
+            "SELECT row_number() OVER "
+            "(ORDER BY mytable.myid ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING)"
+            " AS anon_1 FROM mytable"
+        )
+
+        self.assert_compile(
+            select([func.row_number().over(order_by=expr, rows={'following': None})]),
+            "SELECT row_number() OVER "
+            "(ORDER BY mytable.myid ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)"
+            " AS anon_1 FROM mytable"
+        )
+
+        self.assert_compile(
+            select([func.row_number().over(order_by=expr, range={'preceding': None, 'following': 0})]),
+            "SELECT row_number() OVER "
+            "(ORDER BY mytable.myid RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)"
+            " AS anon_1 FROM mytable"
+        )
+
+        self.assert_compile(
+            select([func.row_number().over(order_by=expr, rows={})]),
+            "SELECT row_number() OVER "
+            "(ORDER BY mytable.myid ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)"
+            " AS anon_1 FROM mytable"
         )
 
     def test_date_between(self):
diff --git a/test/sql/test_text.py b/test/sql/test_text.py
index 20cb2a6fb..5bec5c79e 100644
--- a/test/sql/test_text.py
+++ b/test/sql/test_text.py
@@ -684,7 +684,8 @@ class OrderByLabelResolutionTest(fixtures.TestBase, AssertsCompiledSQL):
 
         self.assert_compile(
             stmt,
-            "SELECT row_number() OVER (PARTITION BY bar ORDER BY foo) "
+            "SELECT row_number() OVER (PARTITION BY bar ORDER BY foo "
+            "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) "
             "AS anon_1 FROM (SELECT foo, bar)"
         )