Overload df.rolling.max() (#476)

densmirn · AlexanderKalistratov · commit 3e72cc542409 · 2020-01-24T14:02:07.000+03:00
* Overload df.rolling.max()

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

* Make comment more clear in test for df.rolling.max
diff --git a/examples/dataframe/rolling/dataframe_rolling_max.py b/examples/dataframe/rolling/dataframe_rolling_max.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_max():
+    df = pd.DataFrame({'A': [4, 3, 5, 2, 6], 'B': [-4, -3, -5, -2, -6]})
+    out_df = df.rolling(3).max()
+
+    # Expect DataFrame of
+    # {'A': [NaN, NaN, 5.0, 5.0, 6.0], 'B': [NaN, NaN, -3.0, -2.0, -2.0]}
+    return out_df
+
+
+print(df_rolling_max())
diff --git a/sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py b/sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py
@@ -149,6 +149,15 @@ def sdc_pandas_dataframe_rolling_count(self):
     return gen_df_rolling_method_impl('count', self)
 
 
+@sdc_overload_method(DataFrameRollingType, 'max')
+def sdc_pandas_dataframe_rolling_max(self):
+
+    ty_checker = TypeChecker('Method rolling.max().')
+    ty_checker.check(self, DataFrameRollingType)
+
+    return gen_df_rolling_method_impl('max', self)
+
+
 @sdc_overload_method(DataFrameRollingType, 'min')
 def sdc_pandas_dataframe_rolling_min(self):
 
@@ -185,6 +194,13 @@ def sdc_pandas_dataframe_rolling_min(self):
     'extra_params': ''
 })
 
+sdc_pandas_dataframe_rolling_max.__doc__ = sdc_pandas_dataframe_rolling_docstring_tmpl.format(**{
+    'method_name': 'max',
+    'example_caption': 'Calculate the rolling maximum.',
+    'limitations_block': '',
+    'extra_params': ''
+})
+
 sdc_pandas_dataframe_rolling_min.__doc__ = sdc_pandas_dataframe_rolling_docstring_tmpl.format(**{
     'method_name': 'min',
     'example_caption': 'Calculate the rolling minimum.',
diff --git a/sdc/tests/test_rolling.py b/sdc/tests/test_rolling.py
@@ -583,14 +583,30 @@ def test_impl(obj, window, min_periods):
                     ref_result = test_impl(obj, window, min_periods)
                     assert_equal(jit_result, ref_result)
 
+    def _test_rolling_max(self, obj):
+        def test_impl(obj, window, min_periods):
+            return obj.rolling(window, min_periods).max()
+
+        hpat_func = self.jit(test_impl)
+        assert_equal = self._get_assert_equal(obj)
+
+        # python implementation crashes if window = 0, jit works correctly
+        for window in range(1, len(obj) + 2):
+            for min_periods in range(window + 1):
+                with self.subTest(obj=obj, window=window,
+                                  min_periods=min_periods):
+                    jit_result = hpat_func(obj, window, min_periods)
+                    ref_result = test_impl(obj, window, min_periods)
+                    assert_equal(jit_result, ref_result)
+
     def _test_rolling_min(self, obj):
         def test_impl(obj, window, min_periods):
             return obj.rolling(window, min_periods).min()
 
         hpat_func = self.jit(test_impl)
         assert_equal = self._get_assert_equal(obj)
 
-        # TODO: fix the issue when window = 0
+        # python implementation crashes if window = 0, jit works correctly
         for window in range(1, len(obj) + 2):
             for min_periods in range(window + 1):
                 with self.subTest(obj=obj, window=window,
@@ -661,6 +677,15 @@ def test_df_rolling_count(self):
 
         self._test_rolling_count(df)
 
+    @skip_sdc_jit('DataFrame.rolling.max() unsupported')
+    def test_df_rolling_max(self):
+        all_data = test_global_input_data_float64
+        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_max(df)
+
     @skip_sdc_jit('DataFrame.rolling.min() unsupported')
     def test_df_rolling_min(self):
         all_data = test_global_input_data_float64
@@ -893,23 +918,11 @@ def test_impl(series, window, min_periods):
 
     @skip_sdc_jit('Series.rolling.max() unsupported Series index')
     def test_series_rolling_max(self):
-        def test_impl(series, window, min_periods):
-            return series.rolling(window, min_periods).max()
-
-        hpat_func = self.jit(test_impl)
-
         all_data = test_global_input_data_float64
         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')
-            # TODO: fix the issue when window = 0
-            for window in range(1, len(series) + 2):
-                for min_periods in range(window + 1):
-                    with self.subTest(series=series, window=window,
-                                      min_periods=min_periods):
-                        jit_result = hpat_func(series, window, min_periods)
-                        ref_result = test_impl(series, window, min_periods)
-                        pd.testing.assert_series_equal(jit_result, ref_result)
+            self._test_rolling_max(series)
 
     @skip_sdc_jit('Series.rolling.mean() unsupported Series index')
     def test_series_rolling_mean(self):
diff --git a/sdc/tests/tests_perf/test_perf_df_rolling.py b/sdc/tests/tests_perf/test_perf_df_rolling.py
@@ -86,6 +86,7 @@ def setUpClass(cls):
         cls.total_data_length = {
             'apply': [2 * 10 ** 5],
             'count': [8 * 10 ** 5],
+            'max': [2 * 10 ** 5],
             'min': [2 * 10 ** 5],
         }
 
@@ -143,5 +144,8 @@ def test_df_rolling_apply_mean(self):
     def test_df_rolling_count(self):
         self._test_df_rolling_method('count')
 
+    def test_df_rolling_max(self):
+        self._test_df_rolling_method('max')
+
     def test_df_rolling_min(self):
         self._test_df_rolling_method('min')
