Implement rolling().count() in Series and Frame (#990)

This PR implements rolling().count():

```python
>>> import databricks.koalas as ks
>>> s = ks.Series([2, 3, float("nan"), 10])
>>> s.rolling(1).count()
0    1.0
1    1.0
2    0.0
3    1.0
Name: 0, dtype: float64

>>> s.to_frame().rolling(1).count()
     0
0  1.0
1  1.0
2  0.0
3  1.0
```

Relates to #977

Author: itholic

databricks/koalas/generic.py

@@ -1427,8 +1427,9 @@ def median(self, accuracy=10000):
         return DataFrame(kdf._internal.copy(sdf=sdf, index_map=[('__DUMMY__', None)])) \
             ._to_internal_pandas().transpose().iloc[:, 0]
 
-    def rolling(self, *args, **kwargs):
-        return Rolling(self)
+    # TODO: 'center', 'win_type', 'on', 'axis' parameter should be implemented.
+    def rolling(self, window, min_periods=None):
+        return Rolling(self, window=window, min_periods=min_periods)
 
     # TODO: 'center' and 'axis' parameter should be implemented.
     #   'axis' implementation, refer https://github.com/databricks/koalas/pull/607

databricks/koalas/groupby.py

@@ -1648,8 +1648,8 @@ def nunique(self, dropna=True):
                  F.when(F.count(F.when(col.isNull(), 1).otherwise(None)) >= 1, 1).otherwise(0))
         return self._reduce_for_stat_function(stat_function, only_numeric=False)
 
-    def rolling(self, *args, **kwargs):
-        return RollingGroupby(self)
+    def rolling(self, window, *args, **kwargs):
+        return RollingGroupby(self, window)
 
     def expanding(self, min_periods=1):
         """

databricks/koalas/missing/window.py

@@ -64,7 +64,6 @@ class _MissingPandasLikeRolling(object):
     aggregate = unsupported_function_rolling("aggregate")
     apply = unsupported_function_rolling("apply")
     corr = unsupported_function_rolling("corr")
-    count = unsupported_function_rolling("count")
     cov = unsupported_function_rolling("cov")
     kurt = unsupported_function_rolling("kurt")
     median = unsupported_function_rolling("median")
@@ -105,7 +104,6 @@ class _MissingPandasLikeRollingGroupby(object):
     aggregate = unsupported_function_rolling("aggregate")
     apply = unsupported_function_rolling("apply")
     corr = unsupported_function_rolling("corr")
-    count = unsupported_function_rolling("count")
     cov = unsupported_function_rolling("cov")
     kurt = unsupported_function_rolling("kurt")
     median = unsupported_function_rolling("median")

databricks/koalas/tests/test_rolling.py

@@ -0,0 +1,46 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import databricks.koalas as ks
+from databricks.koalas.testing.utils import ReusedSQLTestCase, TestUtils
+from databricks.koalas.window import Rolling
+
+
+class RollingTests(ReusedSQLTestCase, TestUtils):
+
+    def test_rolling_error(self):
+        with self.assertRaisesRegex(ValueError, "window must be >= 0"):
+            ks.range(10).rolling(window=-1)
+        with self.assertRaisesRegex(ValueError, "min_periods must be >= 0"):
+            ks.range(10).rolling(window=1, min_periods=-1)
+
+        with self.assertRaisesRegex(
+                TypeError,
+                "kdf_or_kser must be a series or dataframe; however, got:.*int"):
+            Rolling(1, 2)
+
+    def _test_rolling_func(self, f):
+        kser = ks.Series([1, 2, 3])
+        pser = kser.to_pandas()
+        self.assert_eq(repr(getattr(kser.rolling(2), f)()), repr(getattr(pser.rolling(2), f)()))
+
+        kdf = ks.DataFrame({'a': [1, 2, 3, 2], 'b': [4.0, 2.0, 3.0, 1.0]})
+        pdf = kdf.to_pandas()
+        self.assert_eq(repr(getattr(kdf.rolling(2), f)()), repr(getattr(pdf.rolling(2), f)()))
+
+    def test_rolling_count(self):
+        self._test_rolling_func("count")

databricks/koalas/window.py

@@ -33,8 +33,26 @@ class _RollingAndExpanding(object):
 
 
 class Rolling(_RollingAndExpanding):
-    def __init__(self, obj):
-        self.obj = obj
+    def __init__(self, kdf_or_kser, window, min_periods=None):
+        from databricks.koalas import DataFrame, Series
+        from databricks.koalas.groupby import SeriesGroupBy, DataFrameGroupBy
+        window = window - 1
+        min_periods = min_periods if min_periods is not None else 0
+
+        if window < 0:
+            raise ValueError("window must be >= 0")
+        if (min_periods is not None) and (min_periods < 0):
+            raise ValueError("min_periods must be >= 0")
+        self._window_val = window
+        self._min_periods = min_periods
+        self.kdf_or_kser = kdf_or_kser
+        if not isinstance(kdf_or_kser, (DataFrame, Series, DataFrameGroupBy, SeriesGroupBy)):
+            raise TypeError(
+                "kdf_or_kser must be a series or dataframe; however, got: %s" % type(kdf_or_kser))
+        if isinstance(kdf_or_kser, (DataFrame, Series)):
+            self._index_scols = kdf_or_kser._internal.index_scols
+            self._window = Window.orderBy(self._index_scols).rowsBetween(
+                Window.currentRow - window, Window.currentRow)
 
     def __getattr__(self, item: str) -> Any:
         if hasattr(_MissingPandasLikeRolling, item):
@@ -45,6 +63,91 @@ def __getattr__(self, item: str) -> Any:
                 return partial(property_or_func, self)
         raise AttributeError(item)
 
+    def _apply_as_series_or_frame(self, func):
+        """
+        Decorator that can wraps a function that handles Spark column in order
+        to support it in both Koalas Series and DataFrame.
+        Note that the given `func` name should be same as the API's method name.
+        """
+        from databricks.koalas import DataFrame, Series
+
+        if isinstance(self.kdf_or_kser, Series):
+            kser = self.kdf_or_kser
+            return kser._with_new_scol(
+                func(kser._scol)).rename(kser.name)
+        elif isinstance(self.kdf_or_kser, DataFrame):
+            kdf = self.kdf_or_kser
+            applied = []
+            for column in kdf.columns:
+                applied.append(
+                    getattr(kdf[column].rolling(self._window_val + 1,
+                            self._min_periods), func.__name__)())
+
+            sdf = kdf._sdf.select(
+                kdf._internal.index_scols + [c._scol for c in applied])
+            internal = kdf._internal.copy(
+                sdf=sdf,
+                data_columns=[c._internal.data_columns[0] for c in applied],
+                column_index=[c._internal.column_index[0] for c in applied])
+            return DataFrame(internal)
+
+    def count(self):
+        """
+        The rolling count of any non-NaN observations inside the window.
+
+        .. note:: the current implementation of this API uses Spark's Window without
+            specifying partition specification. This leads to move all data into
+            single partition in single machine and could cause serious
+            performance degradation. Avoid this method against very large dataset.
+
+        Returns
+        -------
+        Series or DataFrame
+            Returned object type is determined by the caller of the rolling
+            calculation.
+
+        See Also
+        --------
+        Series.rolling : Calling object with Series data.
+        DataFrame.rolling : Calling object with DataFrames.
+        DataFrame.count : Count of the full DataFrame.
+
+        Examples
+        --------
+        >>> s = ks.Series([2, 3, float("nan"), 10])
+        >>> s.rolling(1).count()
+        0    1.0
+        1    1.0
+        2    0.0
+        3    1.0
+        Name: 0, dtype: float64
+
+        >>> s.rolling(3).count()
+        0    1.0
+        1    2.0
+        2    2.0
+        3    2.0
+        Name: 0, dtype: float64
+
+        >>> s.to_frame().rolling(1).count()
+             0
+        0  1.0
+        1  1.0
+        2  0.0
+        3  1.0
+
+        >>> s.to_frame().rolling(3).count()
+             0
+        0  1.0
+        1  2.0
+        2  2.0
+        3  2.0
+        """
+        def count(scol):
+            return F.count(scol).over(self._window)
+
+        return self._apply_as_series_or_frame(count).astype('float64')
+
 
 class RollingGroupby(Rolling):
     def __getattr__(self, item: str) -> Any:
@@ -56,6 +159,9 @@ def __getattr__(self, item: str) -> Any:
                 return partial(property_or_func, self)
         raise AttributeError(item)
 
+    def count(self):
+        raise NotImplementedError("groupby.rolling().count() is currently not implemented yet.")
+
 
 class Expanding(_RollingAndExpanding):
     def __init__(self, kdf_or_kser, min_periods=1):

docs/source/reference/window.rst

@@ -1,13 +1,21 @@
 ======
 Window
 ======
+.. currentmodule:: databricks.koalas.window
 
 Rolling objects are returned by ``.rolling`` calls: :func:`koalas.DataFrame.rolling`, :func:`koalas.Series.rolling`, etc.
 Expanding objects are returned by ``.expanding`` calls: :func:`koalas.DataFrame.expanding`, :func:`koalas.Series.expanding`, etc.
 
+Standard moving window functions
+--------------------------------
+
+.. autosummary::
+   :toctree: api/
+
+   Rolling.count
+
 Standard expanding window functions
 -----------------------------------
-.. currentmodule:: databricks.koalas.window
 
 .. autosummary::
    :toctree: api/