Repository: databricks/koalas

Issue:
[WIP] Cleanups to groupy
This simplifies somewhat the calls to `GroupBy`.




Discussion:

Resolving patch:
diff --git a/databricks/koalas/generic.py b/databricks/koalas/generic.py
--- a/databricks/koalas/generic.py
+++ b/databricks/koalas/generic.py
@@ -17,6 +17,8 @@
 """
 A base class to be monkey-patched to DataFrame/Column to behave similar to pandas DataFrame/Series.
 """
+from collections.abc import Iterable
+
 import pandas as pd
 
 from pyspark import sql as spark
@@ -134,17 +136,32 @@ def groupby(self, by):
         Falcon      375.0
         Parrot       25.0
         """
-        from databricks.koalas.groupby import GroupBy
+        from databricks.koalas.frame import DataFrame
         from databricks.koalas.series import Series
+        from databricks.koalas.groupby import DataFrameGroupBy, SeriesGroupBy
+
+        df_or_s = self
         if isinstance(by, str):
             by = [by]
         elif isinstance(by, Series):
             by = [by]
-        else:
+        elif isinstance(by, Iterable):
             by = list(by)
-        if len(by) == 0:
+        else:
+            raise ValueError('Not a valid index: TODO')
+        if not len(by):
             raise ValueError('No group keys passed!')
-        return GroupBy(self, by=by)
+        if isinstance(df_or_s, DataFrame):
+            df = df_or_s  # type: DataFrame
+            col_by = [_resolve_col(df, col_or_s) for col_or_s in by]
+            return DataFrameGroupBy(df_or_s, col_by)
+        if isinstance(df_or_s, Series):
+            col = df_or_s  # type: Series
+            anchor = df_or_s._kdf
+            col_by = [_resolve_col(anchor, col_or_s) for col_or_s in by]
+            return SeriesGroupBy(col, col_by)
+        raise TypeError('Constructor expects DataFrame or Series; however, '
+                        'got [%s]' % (df_or_s,))
 
     def compute(self):
         """Alias of `toPandas()` to mimic dask for easily porting tests."""
@@ -160,6 +177,17 @@ def _count_expr(col: spark.Column, spark_type: DataType) -> spark.Column:
             return F.count(col)
 
 
+def _resolve_col(kdf, col_like):
+    if isinstance(col_like, ks.Series):
+        assert kdf == col_like._kdf, \
+            "Cannot combine column argument because it comes from a different dataframe"
+        return col_like
+    elif isinstance(col_like, str):
+        return kdf[col_like]
+    else:
+        raise ValueError(col_like)
+
+
 def _spark_col_apply(kdf_or_ks, sfun):
     """
     Performs a function to all cells on a dataframe, the function being a known sql function.
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -19,7 +19,7 @@
 """
 
 from functools import partial
-from typing import Any, List, Union
+from typing import Any, List
 import numpy as np
 
 from pyspark.sql import functions as F
@@ -27,7 +27,6 @@
 
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
 from databricks.koalas.dask.compatibility import string_types
-from databricks.koalas.generic import _Frame
 from databricks.koalas.frame import DataFrame
 from databricks.koalas.metadata import Metadata
 from databricks.koalas.missing.groupby import _MissingPandasLikeDataFrameGroupBy, \
@@ -35,19 +34,13 @@
 from databricks.koalas.series import Series, _col
 
 
-ColumnLike = Union[str, Series]
-
-
 class GroupBy(object):
-
-    def __new__(cls, obj: _Frame, *args, **kwargs):
-        if isinstance(obj, DataFrame):
-            return super(GroupBy, cls).__new__(DataFrameGroupBy)
-        elif isinstance(obj, Series):
-            return super(GroupBy, cls).__new__(SeriesGroupBy)
-        else:
-            raise TypeError('Constructor expects DataFrame or Series; however, '
-                            'got [%s]' % (obj,))
+    """
+    :ivar _kdf: The parent dataframe that is used to perform the groupby
+    :type _kdf: DataFrame
+    :ivar _groupkeys: The list of keys that will be used to perform the grouping
+    :type _groupkeys: List[Series]
+    """
 
     # TODO: Series support is not implemented yet.
     # TODO: not all arguments are implemented comparing to Pandas' for now.
@@ -245,6 +238,8 @@ def _reduce_for_stat_function(self, sfun, only_numeric):
             stat_exprs = []
             for ks in self._agg_columns:
                 spark_type = ks.spark_type
+                # TODO: we should have a function that takes dataframes and converts the numeric
+                # types. Converting the NaNs is used in a few places, it should be in utils.
                 # Special handle floating point types because Spark's count treats nan as a valid
                 # value, whereas Pandas count doesn't include nan.
                 if isinstance(spark_type, DoubleType) or isinstance(spark_type, FloatType):
@@ -265,9 +260,9 @@ def _reduce_for_stat_function(self, sfun, only_numeric):
 
 class DataFrameGroupBy(GroupBy):
 
-    def __init__(self, kdf: DataFrame, by: List[ColumnLike], agg_columns=None):
+    def __init__(self, kdf: DataFrame, by: List[Series], agg_columns: List[str] = None):
         self._kdf = kdf
-        self._groupkeys = [_resolve_col(kdf, col) for col in by]
+        self._groupkeys = by
 
         if agg_columns is None:
             groupkey_names = set(s.name for s in self._groupkeys)
@@ -284,14 +279,15 @@ def __getitem__(self, item):
         if isinstance(item, str):
             return SeriesGroupBy(self._kdf[item], self._groupkeys)
         else:
+            # TODO: check that item is a list of strings
             return DataFrameGroupBy(self._kdf, self._groupkeys, item)
 
 
 class SeriesGroupBy(GroupBy):
 
-    def __init__(self, ks: Series, by: List[ColumnLike]):
+    def __init__(self, ks: Series, by: List[Series]):
         self._ks = ks
-        self._groupkeys = [_resolve_col(ks._kdf, col) for col in by]
+        self._groupkeys = by
 
     def __getattr__(self, item: str) -> Any:
         if hasattr(_MissingPandasLikeSeriesGroupBy, item):
@@ -299,7 +295,7 @@ def __getattr__(self, item: str) -> Any:
         raise AttributeError(item)
 
     @property
-    def _kdf(self):
+    def _kdf(self) -> DataFrame:
         return self._ks._kdf
 
     @property
@@ -308,14 +304,3 @@ def _agg_columns(self):
 
     def _reduce_for_stat_function(self, sfun, only_numeric):
         return _col(super(SeriesGroupBy, self)._reduce_for_stat_function(sfun, only_numeric))
-
-
-def _resolve_col(kdf: DataFrame, col_like: Union[ColumnLike]) -> Series:
-    if isinstance(col_like, Series):
-        assert kdf == col_like._kdf, \
-            "Cannot combine column argument because it comes from a different dataframe"
-        return col_like
-    elif isinstance(col_like, str):
-        return kdf[col_like]
-    else:
-        raise ValueError(col_like)
diff --git a/databricks/koalas/series.py b/databricks/koalas/series.py
--- a/databricks/koalas/series.py
+++ b/databricks/koalas/series.py
@@ -82,7 +82,9 @@ class Series(_Frame):
     internally.
 
     :ivar _scol: Spark Column instance
+    :type _scol: pyspark.Column
     :ivar _kdf: Parent's Koalas DataFrame
+    :type _kdf: ks.DataFrame
     :ivar _index_info: Each pair holds the index field name which exists in Spark fields,
       and the index name.
     """
@@ -116,7 +118,7 @@ def _init_from_spark(self, scol, kdf, index_info, *args):
         """
         assert index_info is not None
         self._scol = scol
-        self._kdf = kdf
+        self._kdf: ks.DataFrame = kdf
         self._index_info = index_info
 
     # arithmetic operators


Test patch: