Overload df.rolling.quantile() (#481)

densmirn · AlexanderKalistratov · web-flow · commit 5b14c7012318 · 2020-01-31T02:15:21.000+03:00
* Overload df.rolling.quantile()

* Add perf.test for df.rolling.quantile()

Co-authored-by: Alexander Kalistratov &lt;alexander.kalistratov@intel.com&gt;
diff --git a/examples/dataframe/rolling/dataframe_rolling_quantile.py b/examples/dataframe/rolling/dataframe_rolling_quantile.py
@@ -0,0 +1,41 @@
+# *****************************************************************************
+# Copyright (c) 2020, Intel Corporation All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#     Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#     Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
+# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+# *****************************************************************************
+
+import pandas as pd
+from numba import njit
+
+
+@njit
+def df_rolling_quantile():
+    df = pd.DataFrame({'A': [4, 3, 5, 2, 6], 'B': [-4, -3, -5, -2, -6]})
+    out_df = df.rolling(3).quantile(0.25)
+
+    # Expect DataFrame of
+    # {'A': [NaN, NaN, 3.5, 2.5, 3.5], 'B': [NaN, NaN, -4.5, -4.0, -5.5]}
+    return out_df
+
+
+print(df_rolling_quantile())
diff --git a/sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py b/sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py
@@ -26,7 +26,8 @@
 import numpy
 import pandas
 
-from numba.types import float64, Boolean, Omitted, NoneType
+from numba.types import (float64, Boolean, Number, Omitted,
+                         NoneType, StringLiteral, UnicodeType)
 from sdc.datatypes.common_functions import TypeChecker, params2list
 from sdc.datatypes.hpat_pandas_dataframe_rolling_types import DataFrameRollingType
 from sdc.hiframes.pd_dataframe_ext import get_dataframe_data
@@ -352,6 +353,23 @@ def sdc_pandas_dataframe_rolling_min(self):
     return gen_df_rolling_method_impl('min', self)
 
 
+@sdc_overload_method(DataFrameRollingType, 'quantile')
+def sdc_pandas_dataframe_rolling_quantile(self, quantile, interpolation='linear'):
+
+    ty_checker = TypeChecker('Method rolling.quantile().')
+    ty_checker.check(self, DataFrameRollingType)
+
+    if not isinstance(quantile, Number):
+        ty_checker.raise_exc(quantile, 'float', 'quantile')
+
+    str_types = (Omitted, StringLiteral, UnicodeType)
+    if not isinstance(interpolation, str_types) and interpolation != 'linear':
+        ty_checker.raise_exc(interpolation, 'str', 'interpolation')
+
+    return gen_df_rolling_method_impl('quantile', self, args=['quantile'],
+                                      kws={'interpolation': '"linear"'})
+
+
 @sdc_overload_method(DataFrameRollingType, 'skew')
 def sdc_pandas_dataframe_rolling_skew(self):
 
@@ -455,6 +473,25 @@ def sdc_pandas_dataframe_rolling_sum(self):
     'extra_params': ''
 })
 
+sdc_pandas_dataframe_rolling_quantile.__doc__ = sdc_pandas_dataframe_rolling_docstring_tmpl.format(**{
+    'method_name': 'quantile',
+    'example_caption': 'Calculate the rolling quantile.',
+    'limitations_block':
+    """
+    Limitations
+    -----------
+    Supported ``interpolation`` only can be `'linear'`.
+    DataFrame elements cannot be max/min float/integer. Otherwise SDC and Pandas results are different.
+    """,
+    'extra_params':
+    """
+    quantile: :obj:`float`
+        Quantile to compute. 0 <= quantile <= 1.
+    interpolation: :obj:`str`
+        This optional parameter specifies the interpolation method to use.
+    """
+})
+
 sdc_pandas_dataframe_rolling_skew.__doc__ = sdc_pandas_dataframe_rolling_docstring_tmpl.format(**{
     'method_name': 'skew',
     'example_caption': 'Unbiased rolling skewness.',
diff --git a/sdc/tests/test_rolling.py b/sdc/tests/test_rolling.py
@@ -706,6 +706,54 @@ def test_impl(obj, window, min_periods):
                     ref_result = test_impl(obj, window, min_periods)
                     assert_equal(jit_result, ref_result)
 
+    def _test_rolling_quantile(self, obj):
+        def test_impl(obj, window, min_periods, quantile):
+            return obj.rolling(window, min_periods).quantile(quantile)
+
+        hpat_func = self.jit(test_impl)
+        assert_equal = self._get_assert_equal(obj)
+        quantiles = [0, 0.25, 0.5, 0.75, 1]
+
+        for window in range(0, len(obj) + 3, 2):
+            for min_periods, q in product(range(0, window, 2), quantiles):
+                with self.subTest(obj=obj, window=window,
+                                  min_periods=min_periods, quantiles=q):
+                    jit_result = hpat_func(obj, window, min_periods, q)
+                    ref_result = test_impl(obj, window, min_periods, q)
+                    assert_equal(jit_result, ref_result)
+
+    def _test_rolling_quantile_exception_unsupported_types(self, obj):
+        def test_impl(obj, quantile, interpolation):
+            return obj.rolling(3, 2).quantile(quantile, interpolation)
+
+        hpat_func = self.jit(test_impl)
+
+        msg_tmpl = 'Method rolling.quantile(). The object {}\n given: {}\n expected: {}'
+
+        with self.assertRaises(TypingError) as raises:
+            hpat_func(obj, '0.5', 'linear')
+        msg = msg_tmpl.format('quantile', 'unicode_type', 'float')
+        self.assertIn(msg, str(raises.exception))
+
+        with self.assertRaises(TypingError) as raises:
+            hpat_func(obj, 0.5, None)
+        msg = msg_tmpl.format('interpolation', 'none', 'str')
+        self.assertIn(msg, str(raises.exception))
+
+    def _test_rolling_quantile_exception_unsupported_values(self, obj):
+        def test_impl(obj, quantile, interpolation):
+            return obj.rolling(3, 2).quantile(quantile, interpolation)
+
+        hpat_func = self.jit(test_impl)
+
+        with self.assertRaises(ValueError) as raises:
+            hpat_func(obj, 2, 'linear')
+        self.assertIn('quantile value not in [0, 1]', str(raises.exception))
+
+        with self.assertRaises(ValueError) as raises:
+            hpat_func(obj, 0.5, 'lower')
+        self.assertIn('interpolation value not "linear"', str(raises.exception))
+
     def _test_rolling_skew(self, obj):
         def test_impl(obj, window, min_periods):
             return obj.rolling(window, min_periods).skew()
@@ -929,6 +977,37 @@ def test_impl(df):
 
         pd.testing.assert_frame_equal(hpat_func(df), test_impl(df))
 
+    @skip_sdc_jit('DataFrame.rolling.quantile() unsupported')
+    def test_df_rolling_quantile(self):
+        all_data = [
+            list(range(10)), [1., -1., 0., 0.1, -0.1],
+            [1., np.inf, np.inf, -1., 0., np.inf, np.NINF, np.NINF],
+            [np.nan, np.inf, np.inf, np.nan, np.nan, np.nan, np.NINF, np.NZERO]
+        ]
+        length = min(len(d) for d in all_data)
+        data = {n: d[:length] for n, d in zip(string.ascii_uppercase, all_data)}
+        df = pd.DataFrame(data)
+
+        self._test_rolling_quantile(df)
+
+    @skip_sdc_jit('DataFrame.rolling.quantile() unsupported exceptions')
+    def test_df_rolling_quantile_exception_unsupported_types(self):
+        all_data = [[1., -1., 0., 0.1, -0.1], [-1., 1., 0., -0.1, 0.1]]
+        length = min(len(d) for d in all_data)
+        data = {n: d[:length] for n, d in zip(string.ascii_uppercase, all_data)}
+        df = pd.DataFrame(data)
+
+        self._test_rolling_quantile_exception_unsupported_types(df)
+
+    @skip_sdc_jit('DataFrame.rolling.quantile() unsupported exceptions')
+    def test_df_rolling_quantile_exception_unsupported_values(self):
+        all_data = [[1., -1., 0., 0.1, -0.1], [-1., 1., 0., -0.1, 0.1]]
+        length = min(len(d) for d in all_data)
+        data = {n: d[:length] for n, d in zip(string.ascii_uppercase, all_data)}
+        df = pd.DataFrame(data)
+
+        self._test_rolling_quantile_exception_unsupported_values(df)
+
     @skip_sdc_jit('DataFrame.rolling.skew() unsupported')
     def test_df_rolling_skew(self):
         all_data = test_global_input_data_float64
@@ -1152,63 +1231,25 @@ def test_series_rolling_min(self):
 
     @skip_sdc_jit('Series.rolling.quantile() unsupported Series index')
     def test_series_rolling_quantile(self):
-        def test_impl(series, window, min_periods, quantile):
-            return series.rolling(window, min_periods).quantile(quantile)
-
-        hpat_func = self.jit(test_impl)
-
         all_data = [
             list(range(10)), [1., -1., 0., 0.1, -0.1],
             [1., np.inf, np.inf, -1., 0., np.inf, np.NINF, np.NINF],
             [np.nan, np.inf, np.inf, np.nan, np.nan, np.nan, np.NINF, np.NZERO]
         ]
         indices = [list(range(len(data)))[::-1] for data in all_data]
-        quantiles = [0, 0.25, 0.5, 0.75, 1]
         for data, index in zip(all_data, indices):
             series = pd.Series(data, index, name='A')
-            for window in range(0, len(series) + 3, 2):
-                for min_periods, q in product(range(0, window, 2), quantiles):
-                    with self.subTest(series=series, window=window,
-                                      min_periods=min_periods, quantiles=q):
-                        jit_result = hpat_func(series, window, min_periods, q)
-                        ref_result = test_impl(series, window, min_periods, q)
-                        pd.testing.assert_series_equal(jit_result, ref_result)
+            self._test_rolling_quantile(series)
 
     @skip_sdc_jit('Series.rolling.quantile() unsupported exceptions')
     def test_series_rolling_quantile_exception_unsupported_types(self):
-        def test_impl(quantile, interpolation):
-            series = pd.Series([1., -1., 0., 0.1, -0.1])
-            return series.rolling(3, 2).quantile(quantile, interpolation)
-
-        hpat_func = self.jit(test_impl)
-
-        msg_tmpl = 'Method rolling.quantile(). The object {}\n given: {}\n expected: {}'
-
-        with self.assertRaises(TypingError) as raises:
-            hpat_func('0.5', 'linear')
-        msg = msg_tmpl.format('quantile', 'unicode_type', 'float')
-        self.assertIn(msg, str(raises.exception))
-
-        with self.assertRaises(TypingError) as raises:
-            hpat_func(0.5, None)
-        msg = msg_tmpl.format('interpolation', 'none', 'str')
-        self.assertIn(msg, str(raises.exception))
+        series = pd.Series([1., -1., 0., 0.1, -0.1])
+        self._test_rolling_quantile_exception_unsupported_types(series)
 
     @skip_sdc_jit('Series.rolling.quantile() unsupported exceptions')
     def test_series_rolling_quantile_exception_unsupported_values(self):
-        def test_impl(quantile, interpolation):
-            series = pd.Series([1., -1., 0., 0.1, -0.1])
-            return series.rolling(3, 2).quantile(quantile, interpolation)
-
-        hpat_func = self.jit(test_impl)
-
-        with self.assertRaises(ValueError) as raises:
-            hpat_func(2, 'linear')
-        self.assertIn('quantile value not in [0, 1]', str(raises.exception))
-
-        with self.assertRaises(ValueError) as raises:
-            hpat_func(0.5, 'lower')
-        self.assertIn('interpolation value not "linear"', str(raises.exception))
+        series = pd.Series([1., -1., 0., 0.1, -0.1])
+        self._test_rolling_quantile_exception_unsupported_values(series)
 
     @skip_sdc_jit('Series.rolling.skew() unsupported Series index')
     def test_series_rolling_skew(self):
diff --git a/sdc/tests/tests_perf/test_perf_df_rolling.py b/sdc/tests/tests_perf/test_perf_df_rolling.py
@@ -94,6 +94,7 @@ def setUpClass(cls):
             'mean': [2 * 10 ** 5],
             'median': [2 * 10 ** 5],
             'min': [2 * 10 ** 5],
+            'quantile': [2 * 10 ** 5],
             'skew': [2 * 10 ** 5],
             'sum': [2 * 10 ** 5],
         }
@@ -196,6 +197,9 @@ def test_df_rolling_median(self):
     def test_df_rolling_min(self):
         self._test_df_rolling_method('min')
 
+    def test_df_rolling_quantile(self):
+        self._test_df_rolling_method('quantile', method_params='0.25')
+
     def test_df_rolling_skew(self):
         self._test_df_rolling_method('skew')
 
