Repository: kayak/pypika
Pull request #57 (2017-06-15 12:30:40)

Window frames in analytic queries
Added support for setting a window frame in analytic queries that support them.


Patch:
diff --git a/docs/3_advanced.rst b/docs/3_advanced.rst
index 12540bf..b905a2a 100644
--- a/docs/3_advanced.rst
+++ b/docs/3_advanced.rst
@@ -128,4 +128,35 @@ These functions take one or more arguments
 
 .. code-block:: sql
 
-    SELECT MEDIAN("annual_income") OVER(PARTITION BY "customer_state") "MEDIAN",AVG("annual_income") OVER(PARTITION BY "customer_state") "AVG",STDDEV("annual_income") OVER(PARTITION BY "customer_state") "STDDEV" FROM "customer_dimension" WHERE "customer_state" IN ('DC','WI') ORDER BY "customer_state"
\ No newline at end of file
+    SELECT MEDIAN("annual_income") OVER(PARTITION BY "customer_state") "MEDIAN",AVG("annual_income") OVER(PARTITION BY "customer_state") "AVG",STDDEV("annual_income") OVER(PARTITION BY "customer_state") "STDDEV" FROM "customer_dimension" WHERE "customer_state" IN ('DC','WI') ORDER BY "customer_state"
+
+Window Frames
+=============
+
+Functions which use window aggregation expose the functions ``rows()`` and ``range()`` with varying parameters to define
+the window.  Both of these functions take one or two parameters which specify the offset boundaries.  Boundaries can be
+set either as the current row with ``an.CURRENT_ROW`` or a value preceding or following the current row with
+``an.Preceding(constant_value)`` and ``an.Following(constant_value)``.  The ranges can be unbounded preceding or
+following the current row by omitting the ``constant_value`` parameter like ``an.Preceding()`` or ``an.Following()``.
+
+``FIRST_VALUE`` and ``LAST_VALUE`` also support window frames.
+
+.. code-block:: python
+
+    from pypika import Query, Table, analytics as an
+
+    t_transactions = Table('t_customers')
+
+    rolling_7_sum = an.Sum(t_transactions.total) \
+        .over(t_transactions.item_id) \
+        .orderby(t_transactions.day) \
+        .rows(an.Preceding(7), an.CURRENT_ROW)
+
+    query = Query.from_(t_transactions) \
+        .select(rolling_7_sum)
+
+.. code-block:: sql
+
+    SELECT SUM("total") OVER(PARTITION BY "item_id" ORDER BY "day" ROWS BETWEEN 7 PRECEDING AND CURRENT ROW) FROM "t_customers"
+
+
diff --git a/pypika/analytics.py b/pypika/analytics.py
index 69affde..ce38ba7 100644
--- a/pypika/analytics.py
+++ b/pypika/analytics.py
@@ -2,12 +2,25 @@
 """
 Package for SQL analytic functions wrappers
 """
-from pypika.terms import AnalyticFunction, IgnoreNullsAnalyticFunction
+from pypika.terms import (AnalyticFunction,
+                          WindowFrameAnalyticFunction,
+                          IgnoreNullsAnalyticFunction)
 
 __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
 
 
+class Preceding(WindowFrameAnalyticFunction.Edge):
+    modifier = 'PRECEDING'
+
+
+class Following(WindowFrameAnalyticFunction.Edge):
+    modifier = 'FOLLOWING'
+
+
+CURRENT_ROW = 'CURRENT ROW'
+
+
 class Rank(AnalyticFunction):
     def __init__(self, **kwargs):
         super(Rank, self).__init__('RANK', **kwargs)
@@ -18,12 +31,12 @@ class NTile(AnalyticFunction):
         super(NTile, self).__init__('NTILE', term, **kwargs)
 
 
-class FirstValue(IgnoreNullsAnalyticFunction):
+class FirstValue(WindowFrameAnalyticFunction, IgnoreNullsAnalyticFunction):
     def __init__(self, *terms, **kwargs):
         super(FirstValue, self).__init__('FIRST_VALUE', *terms, **kwargs)
 
 
-class LastValue(IgnoreNullsAnalyticFunction):
+class LastValue(WindowFrameAnalyticFunction, IgnoreNullsAnalyticFunction):
     def __init__(self, *terms, **kwargs):
         super(LastValue, self).__init__('LAST_VALUE', *terms, **kwargs)
 
@@ -33,11 +46,56 @@ class Median(AnalyticFunction):
         super(Median, self).__init__('MEDIAN', term, **kwargs)
 
 
-class Avg(AnalyticFunction):
+class Avg(WindowFrameAnalyticFunction):
     def __init__(self, term, **kwargs):
         super(Avg, self).__init__('AVG', term, **kwargs)
 
 
-class StdDev(AnalyticFunction):
+class StdDev(WindowFrameAnalyticFunction):
     def __init__(self, term, **kwargs):
         super(StdDev, self).__init__('STDDEV', term, **kwargs)
+
+
+class StdDevPop(WindowFrameAnalyticFunction):
+    def __init__(self, term, **kwargs):
+        super(StdDevPop, self).__init__('STDDEV_POP', term, **kwargs)
+
+
+class StdDevSamp(WindowFrameAnalyticFunction):
+    def __init__(self, term, **kwargs):
+        super(StdDevSamp, self).__init__('STDDEV_SAMP', term, **kwargs)
+
+
+class Variance(WindowFrameAnalyticFunction):
+    def __init__(self, term, **kwargs):
+        super(Variance, self).__init__('VARIANCE', term, **kwargs)
+
+
+class VarPop(WindowFrameAnalyticFunction):
+    def __init__(self, term, **kwargs):
+        super(VarPop, self).__init__('VAR_POP', term, **kwargs)
+
+
+class VarSamp(WindowFrameAnalyticFunction):
+    def __init__(self, term, **kwargs):
+        super(VarSamp, self).__init__('VAR_SAMP', term, **kwargs)
+
+
+class Count(WindowFrameAnalyticFunction):
+    def __init__(self, term, **kwargs):
+        super(Count, self).__init__('COUNT', term, **kwargs)
+
+
+class Sum(WindowFrameAnalyticFunction):
+    def __init__(self, term, **kwargs):
+        super(Sum, self).__init__('SUM', term, **kwargs)
+
+
+class Max(WindowFrameAnalyticFunction):
+    def __init__(self, term, **kwargs):
+        super(Max, self).__init__('MAX', term, **kwargs)
+
+
+class Min(WindowFrameAnalyticFunction):
+    def __init__(self, term, **kwargs):
+        super(Min, self).__init__('MIN', term, **kwargs)
diff --git a/pypika/terms.py b/pypika/terms.py
index cc30bc8..1cea8cd 100644
--- a/pypika/terms.py
+++ b/pypika/terms.py
@@ -645,15 +645,25 @@ class AnalyticFunction(Function):
     def __init__(self, name, *args, **kwargs):
         super(AnalyticFunction, self).__init__(name, *args, **kwargs)
         self._partition = []
-        self._order = []
+        self._orderbys = []
 
     @builder
     def over(self, *terms):
         self._partition += terms
 
     @builder
-    def orderby(self, *terms):
-        self._order += terms
+    def orderby(self, *terms, **kwargs):
+        self._orderbys += [(term, kwargs.get('order'))
+                           for term in terms]
+
+    def _orderby_field(self, field, orient, **kwargs):
+        if orient is None:
+            return field.get_sql(**kwargs)
+
+        return '{field} {orient}'.format(
+            field=field.get_sql(**kwargs),
+            orient=orient.value,
+        )
 
     def get_partition_sql(self, **kwargs):
         terms = []
@@ -664,12 +674,12 @@ class AnalyticFunction(Function):
                               else str(p)
                               for p in self._partition)))
 
-        if self._order:
-            terms.append('ORDER BY {args}'.format(
-                args=','.join(p.get_sql(**kwargs)
-                              if hasattr(p, 'get_sql')
-                              else str(p)
-                              for p in self._order)))
+        if self._orderbys:
+            terms.append('ORDER BY {orderby}'.format(
+                orderby=','.join(
+                    self._orderby_field(field, orient, **kwargs)
+                    for field, orient in self._orderbys
+                )))
 
         return ' '.join(terms)
 
@@ -686,6 +696,63 @@ class AnalyticFunction(Function):
         )
 
 
+class WindowFrameAnalyticFunction(AnalyticFunction):
+    class Edge:
+        def __init__(self, value=None):
+            self.value = value
+
+        def __str__(self):
+            return '{value} {modifier}'.format(
+                value=self.value or 'UNBOUNDED',
+                modifier=self.modifier,
+            )
+
+    def __init__(self, name, *args, **kwargs):
+        super(WindowFrameAnalyticFunction, self).__init__(name, *args, **kwargs)
+        self.frame = None
+        self.bound = None
+
+    def _set_frame_and_bounds(self, frame, bound, and_bound):
+        if self.frame or self.bound:
+            raise AttributeError()
+
+        self.frame = frame
+        self.bound = (bound, and_bound) if and_bound else bound
+
+    @builder
+    def rows(self, bound, and_bound=None):
+        self._set_frame_and_bounds('ROWS', bound, and_bound)
+
+    @builder
+    def range(self, bound, and_bound=None):
+        self._set_frame_and_bounds('RANGE', bound, and_bound)
+
+    def get_frame_sql(self):
+        if not isinstance(self.bound, tuple):
+            return '{frame} {bound}'.format(
+                frame=self.frame,
+                bound=self.bound
+            )
+
+        lower, upper = self.bound
+        return '{frame} BETWEEN {lower} AND {upper}'.format(
+            frame=self.frame,
+            lower=lower,
+            upper=upper,
+        )
+
+    def get_partition_sql(self, **kwargs):
+        partition_sql = super(WindowFrameAnalyticFunction, self).get_partition_sql(**kwargs)
+
+        if not self.frame and not self.bound:
+            return partition_sql
+
+        return '{over} {frame}'.format(
+            over=partition_sql,
+            frame=self.get_frame_sql()
+        )
+
+
 class IgnoreNullsAnalyticFunction(AnalyticFunction):
     def __init__(self, name, *args, **kwargs):
         super(IgnoreNullsAnalyticFunction, self).__init__(name, *args, **kwargs)
@@ -702,14 +769,6 @@ class IgnoreNullsAnalyticFunction(AnalyticFunction):
         # No special params unless ignoring nulls
         return None
 
-    def get_function_sql(self, with_namespace=False, quote_char=None):
-        function_sql = super(AnalyticFunction, self).get_function_sql(with_namespace=False, quote_char=quote_char)
-
-        return '{function_sql} OVER({partition_sql})'.format(
-            function_sql=function_sql,
-            partition_sql=self.get_partition_sql(with_alias=False, with_namespace=with_namespace, quote_char=quote_char)
-        )
-
 
 class Interval(object):
     units = ['years', 'months', 'days', 'hours', 'minutes', 'seconds', 'microseconds']


Test patch:
diff --git a/pypika/tests/test_analytic_queries.py b/pypika/tests/test_analytic_queries.py
index 4952f34..2fbf6f2 100644
--- a/pypika/tests/test_analytic_queries.py
+++ b/pypika/tests/test_analytic_queries.py
@@ -1,7 +1,7 @@
 # coding: utf-8
 import unittest
 
-from pypika import Query, Tables, analytics as an, JoinType
+from pypika import Query, Tables, analytics as an, JoinType, Order
 
 __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
@@ -22,6 +22,19 @@ class RankTests(unittest.TestCase):
                          'OVER(PARTITION BY "foo" ORDER BY "date") '
                          'FROM "abc"', str(q))
 
+    def test_rank_with_alias(self):
+        expr = an.Rank() \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .as_('rank')
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'RANK() '
+                         'OVER(PARTITION BY "foo" ORDER BY "date") "rank" '
+                         'FROM "abc"', str(q))
+
     def test_multiple_partitions(self):
         expr = an.Rank() \
             .orderby(self.table_abc.date) \
@@ -139,6 +152,30 @@ class RankTests(unittest.TestCase):
                          'OVER(PARTITION BY "foo" ORDER BY "date") '
                          'FROM "abc"', str(q))
 
+    def test_orderby_asc(self):
+        expr = an.LastValue(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date, order=Order.asc)
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'LAST_VALUE("fizz") '
+                         'OVER(PARTITION BY "foo" ORDER BY "date" ASC) '
+                         'FROM "abc"', str(q))
+
+    def test_orderby_desc(self):
+        expr = an.LastValue(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date, order=Order.desc)
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'LAST_VALUE("fizz") '
+                         'OVER(PARTITION BY "foo" ORDER BY "date" DESC) '
+                         'FROM "abc"', str(q))
+
     def test_last_value_multi_argument(self):
         expr = an.LastValue(self.table_abc.fizz, self.table_abc.buzz) \
             .over(self.table_abc.foo) \
@@ -212,3 +249,228 @@ class RankTests(unittest.TestCase):
 
         self.assertEqual('SELECT *,RANK() OVER(PARTITION BY "abc"."foo" ORDER BY "efg"."date") '
                          'FROM "abc" LEFT JOIN "efg" ON "abc"."foo"="efg"."bar"', str(query))
+
+    def test_sum_rows_unbounded_preceeding(self):
+        expr = an.Sum(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .rows(an.Preceding())
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'SUM("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'ROWS UNBOUNDED PRECEDING'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_max_rows_x_preceeding(self):
+        expr = an.Max(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .rows(an.Preceding(5))
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'MAX("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'ROWS 5 PRECEDING'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_min_rows_current_row(self):
+        expr = an.Min(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .rows(an.CURRENT_ROW)
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'MIN("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'ROWS CURRENT ROW'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_varpop_range_unbounded_preceeding(self):
+        expr = an.VarPop(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .range(an.Preceding())
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'VAR_POP("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'RANGE UNBOUNDED PRECEDING'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_max_range_x_preceeding(self):
+        expr = an.Max(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .range(an.Preceding(5))
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'MAX("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'RANGE 5 PRECEDING'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_min_range_current_row(self):
+        expr = an.Min(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .range(an.CURRENT_ROW)
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'MIN("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'RANGE CURRENT ROW'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_variance_rows_between_unbounded_preceeding_unbounded_following(self):
+        expr = an.Variance(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .rows(an.Preceding(), an.Following())
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'VARIANCE("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_first_value_range_between_x_preceeding_unbounded_following(self):
+        expr = an.FirstValue(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .range(an.Preceding(3), an.Following())
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'FIRST_VALUE("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'RANGE BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_varpop_rows_between_unbounded_preceeding_x_following(self):
+        expr = an.VarPop(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .rows(an.Preceding(), an.Following(6))
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'VAR_POP("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'ROWS BETWEEN UNBOUNDED PRECEDING AND 6 FOLLOWING'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_count_range_between_unbounded_preceeding_current_row(self):
+        expr = an.Count(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .range(an.Preceding(), an.CURRENT_ROW)
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'COUNT("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_last_value_rows_between_current_row_unbounded_following(self):
+        expr = an.LastValue(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .rows(an.CURRENT_ROW, an.Following())
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'LAST_VALUE("fizz") '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_last_value_rows_between_current_row_unbounded_following_ignore_nulls(self):
+        expr = an.LastValue(self.table_abc.fizz) \
+            .over(self.table_abc.foo) \
+            .orderby(self.table_abc.date) \
+            .ignore_nulls() \
+            .rows(an.CURRENT_ROW, an.Following(8))
+
+        q = Query.from_(self.table_abc).select(expr)
+
+        self.assertEqual('SELECT '
+                         'LAST_VALUE("fizz" IGNORE NULLS) '
+                         'OVER('
+                         'PARTITION BY "foo" ORDER BY "date" '
+                         'ROWS BETWEEN CURRENT ROW AND 8 FOLLOWING'
+                         ') '
+                         'FROM "abc"', str(q))
+
+    def test_rows_called_twice_raises_attribute_error(self):
+        with self.assertRaises(AttributeError):
+            an.Sum(self.table_abc.fizz) \
+                .over(self.table_abc.foo) \
+                .orderby(self.table_abc.date) \
+                .rows(an.Preceding()) \
+                .rows(an.Preceding())
+
+    def test_range_called_twice_raises_attribute_error(self):
+        with self.assertRaises(AttributeError):
+            an.Sum(self.table_abc.fizz) \
+                .over(self.table_abc.foo) \
+                .orderby(self.table_abc.date) \
+                .range(an.Preceding()) \
+                .range(an.Preceding())
+
+    def test_rows_then_range_raises_attribute_error(self):
+        with self.assertRaises(AttributeError):
+            an.Sum(self.table_abc.fizz) \
+                .over(self.table_abc.foo) \
+                .orderby(self.table_abc.date) \
+                .rows(an.Preceding()) \
+                .range(an.Preceding())
+
+    def test_range_then_rows_raises_attribute_error(self):
+        with self.assertRaises(AttributeError):
+            an.Sum(self.table_abc.fizz) \
+                .over(self.table_abc.foo) \
+                .orderby(self.table_abc.date) \
+                .range(an.Preceding()) \
+                .rows(an.Preceding())