Overload df.rolling.var() (#483)

densmirn · web-flow · commit 333e8438e417 · 2020-02-02T00:13:02.000+03:00
* Overload df.rolling.var()

* Fix a code style issue

* Add perf.test for df.rolling.var

* Minor changes in a test for df.rolling.var
diff --git a/examples/dataframe/rolling/dataframe_rolling_var.py b/examples/dataframe/rolling/dataframe_rolling_var.py
@@ -0,0 +1,42 @@
+# *****************************************************************************
+# 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_var():
+    df = pd.DataFrame({'A': [4, 3, 5, 2, 6], 'B': [-4, -3, -5, -2, -6]})
+    out_df = df.rolling(3).var()
+
+    # Expect DataFrame of
+    # {'A': [NaN, NaN, 1.000000, 2.333333, 4.333333],
+    #  'B': [NaN, NaN, 1.000000, 2.333333, 4.333333]}
+    return out_df
+
+
+print(df_rolling_var())
diff --git a/sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py b/sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py
@@ -400,6 +400,18 @@ def sdc_pandas_dataframe_rolling_sum(self):
     return gen_df_rolling_method_impl('sum', self)
 
 
+@sdc_overload_method(DataFrameRollingType, 'var')
+def sdc_pandas_dataframe_rolling_var(self, ddof=1):
+
+    ty_checker = TypeChecker('Method rolling.var().')
+    ty_checker.check(self, DataFrameRollingType)
+
+    if not isinstance(ddof, (int, Integer, Omitted)):
+        ty_checker.raise_exc(ddof, 'int', 'ddof')
+
+    return gen_df_rolling_method_impl('var', self, kws={'ddof': '1'})
+
+
 sdc_pandas_dataframe_rolling_apply.__doc__ = sdc_pandas_dataframe_rolling_docstring_tmpl.format(**{
     'method_name': 'apply',
     'example_caption': 'Calculate the rolling apply.',
@@ -538,3 +550,19 @@ def sdc_pandas_dataframe_rolling_sum(self):
     """,
     'extra_params': ''
 })
+
+sdc_pandas_dataframe_rolling_var.__doc__ = sdc_pandas_dataframe_rolling_docstring_tmpl.format(**{
+    'method_name': 'var',
+    'example_caption': 'Calculate unbiased rolling variance.',
+    'limitations_block':
+    """
+    Limitations
+    -----------
+    DataFrame elements cannot be max/min float/integer. Otherwise SDC and Pandas results are different.
+    """,
+    'extra_params':
+    """
+    ddof: :obj:`int`
+        Delta Degrees of Freedom.
+    """
+})
diff --git a/sdc/tests/test_rolling.py b/sdc/tests/test_rolling.py
@@ -56,8 +56,8 @@ def rolling_std_usecase(obj, window, min_periods, ddof):
     return obj.rolling(window, min_periods).std(ddof)
 
 
-def series_rolling_var_usecase(series, window, min_periods, ddof):
-    return series.rolling(window, min_periods).var(ddof)
+def rolling_var_usecase(obj, window, min_periods, ddof):
+    return obj.rolling(window, min_periods).var(ddof)
 
 
 class TestRolling(TestCase):
@@ -807,6 +807,29 @@ def test_impl(obj, window, min_periods):
                     ref_result = test_impl(obj, window, min_periods)
                     assert_equal(jit_result, ref_result)
 
+    def _test_rolling_var(self, obj):
+        test_impl = rolling_var_usecase
+        hpat_func = self.jit(test_impl)
+        assert_equal = self._get_assert_equal(obj)
+
+        for window in range(0, len(obj) + 3, 2):
+            for min_periods, ddof in product(range(0, window, 2), [0, 1]):
+                with self.subTest(obj=obj, window=window,
+                                  min_periods=min_periods, ddof=ddof):
+                    jit_result = hpat_func(obj, window, min_periods, ddof)
+                    ref_result = test_impl(obj, window, min_periods, ddof)
+                    assert_equal(jit_result, ref_result)
+
+    def _test_rolling_var_exception_unsupported_ddof(self, obj):
+        test_impl = rolling_var_usecase
+        hpat_func = self.jit(test_impl)
+
+        window, min_periods, invalid_ddof = 3, 2, '1'
+        with self.assertRaises(TypingError) as raises:
+            hpat_func(obj, window, min_periods, invalid_ddof)
+        msg = 'Method rolling.var(). The object ddof\n given: unicode_type\n expected: int'
+        self.assertIn(msg, str(raises.exception))
+
     @skip_sdc_jit('DataFrame.rolling.min() unsupported exceptions')
     def test_df_rolling_unsupported_values(self):
         all_data = test_global_input_data_float64
@@ -1075,6 +1098,28 @@ def test_df_rolling_sum(self):
 
         self._test_rolling_sum(df)
 
+    @skip_sdc_jit('DataFrame.rolling.var() unsupported')
+    def test_df_rolling_var(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_var(df)
+
+    @skip_sdc_jit('DataFrame.rolling.var() unsupported exceptions')
+    def test_df_rolling_var_exception_unsupported_ddof(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_var_exception_unsupported_ddof(df)
+
     @skip_sdc_jit('Series.rolling.min() unsupported exceptions')
     def test_series_rolling_unsupported_values(self):
         series = pd.Series(test_global_input_data_float64[0])
@@ -1335,9 +1380,6 @@ def test_series_rolling_sum(self):
 
     @skip_sdc_jit('Series.rolling.var() unsupported Series index')
     def test_series_rolling_var(self):
-        test_impl = series_rolling_var_usecase
-        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],
@@ -1346,24 +1388,12 @@ def test_series_rolling_var(self):
         indices = [list(range(len(data)))[::-1] for data in all_data]
         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, ddof in product(range(0, window, 2), [0, 1]):
-                    with self.subTest(series=series, window=window,
-                                      min_periods=min_periods, ddof=ddof):
-                        jit_result = hpat_func(series, window, min_periods, ddof)
-                        ref_result = test_impl(series, window, min_periods, ddof)
-                        pd.testing.assert_series_equal(jit_result, ref_result)
+            self._test_rolling_var(series)
 
     @skip_sdc_jit('Series.rolling.var() unsupported exceptions')
     def test_series_rolling_var_exception_unsupported_ddof(self):
-        test_impl = series_rolling_var_usecase
-        hpat_func = self.jit(test_impl)
-
         series = pd.Series([1., -1., 0., 0.1, -0.1])
-        with self.assertRaises(TypingError) as raises:
-            hpat_func(series, 3, 2, '1')
-        msg = 'Method rolling.var(). The object ddof\n given: unicode_type\n expected: int'
-        self.assertIn(msg, str(raises.exception))
+        self._test_rolling_var_exception_unsupported_ddof(series)
 
 
 if __name__ == "__main__":
diff --git a/sdc/tests/tests_perf/test_perf_df_rolling.py b/sdc/tests/tests_perf/test_perf_df_rolling.py
@@ -98,6 +98,7 @@ def setUpClass(cls):
             'skew': [2 * 10 ** 5],
             'std': [2 * 10 ** 5],
             'sum': [2 * 10 ** 5],
+            'var': [2 * 10 ** 5],
         }
 
     def _test_jitted(self, pyfunc, record, *args, **kwargs):
@@ -209,3 +210,6 @@ def test_df_rolling_std(self):
 
     def test_df_rolling_sum(self):
         self._test_df_rolling_method('sum')
+
+    def test_df_rolling_var(self):
+        self._test_df_rolling_method('var')
