Cleanup

Author: HyukjinKwon

databricks/koalas/frame.py

@@ -26,7 +26,7 @@
 from functools import partial, reduce
 import sys
 from itertools import zip_longest
-from typing import Any, Optional, List, Tuple, Union, Generic, TypeVar, Iterable
+from typing import Any, Optional, List, Tuple, Union, Generic, TypeVar, Iterable, Dict
 
 import numpy as np
 import pandas as pd
@@ -46,7 +46,6 @@
 from pyspark.sql.functions import pandas_udf, PandasUDFType
 
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
-from databricks.koalas.config import get_option
 from databricks.koalas.utils import validate_arguments_and_invoke_function, align_diff_frames
 from databricks.koalas.generic import _Frame
 from databricks.koalas.internal import _InternalFrame, IndexMap
@@ -874,7 +873,7 @@ def applymap(self, func):
 
     # TODO: Series support is not implemented yet.
     # TODO: not all arguments are implemented comparing to Pandas' for now.
-    def aggregate(self, func_or_funcs):
+    def aggregate(self, func: Union[List[str], Dict[str, List[str]]]):
         """Aggregate using one or more operations over the specified axis.
 
         Parameters
@@ -889,12 +888,6 @@ def aggregate(self, func_or_funcs):
         -------
         DataFrame
 
-            The return can be:
-
-            * DataFrame : when DataFrame.agg is called with several functions
-
-            Return a DataFrame.
-
         Notes
         -----
         `agg` is an alias for `aggregate`. Use the alias.
@@ -934,46 +927,48 @@ def aggregate(self, func_or_funcs):
         min   1.0  2.0
         sum  12.0  NaN
         """
-        if isinstance(func_or_funcs, list):
-            func_or_funcs = dict([
-                (column, func_or_funcs) for column in self.columns])
+        from databricks.koalas.groupby import GroupBy
 
-        if not isinstance(func_or_funcs, dict) or \
+        if isinstance(func, list):
+            if all((isinstance(f, str) for f in func)):
+                func = dict([
+                    (column, func) for column in self.columns])
+            else:
+                raise ValueError("If the given function is a list, it "
+                                 "should only contains function names as strings.")
+
+        if not isinstance(func, dict) or \
                 not all(isinstance(key, str) and
                         (isinstance(value, str) or
                          isinstance(value, list) and all(isinstance(v, str) for v in value))
-                        for key, value in func_or_funcs.items()):
+                        for key, value in func.items()):
             raise ValueError("aggs must be a dict mapping from column name (string) to aggregate "
-                             "functions (string or list of strings).")
+                             "functions (list of strings).")
 
-        sdf = self._sdf
-        multi_aggs = any(isinstance(v, list) for v in func_or_funcs.values())
-        reordered = []
-        data_columns = []
-        column_index = []
-        for key, value in func_or_funcs.items():
-            for aggfunc in [value] if isinstance(value, str) else value:
-                data_col = "('{0}', '{1}')".format(key, aggfunc) if multi_aggs else key
-                data_columns.append(data_col)
-                column_index.append((key, aggfunc))
-                if aggfunc == "nunique":
-                    reordered.append(F.expr('count(DISTINCT `{0}`) as `{1}`'.format(key, data_col)))
-                else:
-                    reordered.append(F.expr('{1}(`{0}`) as `{2}`'.format(key, aggfunc, data_col)))
-        sdf = sdf.groupby().agg(*reordered)
-        internal = _InternalFrame(sdf=sdf,
-                                  data_columns=data_columns,
-                                  column_index=column_index if multi_aggs else None)
+        kdf = DataFrame(GroupBy._spark_groupby(self, func, ()))  # type: DataFrame
 
-        kdf = DataFrame(internal)
+        # The codes below basically converts:
+        #
+        #           A         B
+        #         sum  min  min  max
+        #     0  12.0  1.0  2.0  8.0
+        #
+        # to:
+        #             A    B
+        #     max   NaN  8.0
+        #     min   1.0  2.0
+        #     sum  12.0  NaN
+        #
+        # Aggregated output is usually pretty much small. So it is fine to directly use pandas API.
         pdf = kdf.to_pandas().transpose().reset_index()
         pdf = pdf.groupby(['level_1']).apply(
             lambda gpdf: gpdf.drop('level_1', 1).set_index('level_0').transpose()
         ).reset_index(level=1)
         pdf = pdf.drop(columns='level_1')
         pdf.columns.names = [None]
         pdf.index.names = [None]
-        return DataFrame(pdf[func_or_funcs.keys()])
+
+        return DataFrame(pdf[list(func.keys())])
 
     agg = aggregate
 

databricks/koalas/groupby.py

@@ -120,15 +120,23 @@ def aggregate(self, func_or_funcs, *args, **kwargs):
             raise ValueError("aggs must be a dict mapping from column name (string) to aggregate "
                              "functions (string or list of strings).")
 
-        sdf = self._kdf._sdf
-        groupkeys = self._groupkeys
+        kdf = DataFrame(GroupBy._spark_groupby(self._kdf, func_or_funcs, self._groupkeys))
+        if not self._as_index:
+            kdf = kdf.reset_index()
+        return kdf
+
+    agg = aggregate
+
+    @staticmethod
+    def _spark_groupby(kdf, func, groupkeys):
+        sdf = kdf._sdf
         groupkey_cols = [s._scol.alias('__index_level_{}__'.format(i))
                          for i, s in enumerate(groupkeys)]
-        multi_aggs = any(isinstance(v, list) for v in func_or_funcs.values())
+        multi_aggs = any(isinstance(v, list) for v in func.values())
         reordered = []
         data_columns = []
         column_index = []
-        for key, value in func_or_funcs.items():
+        for key, value in func.items():
             for aggfunc in [value] if isinstance(value, str) else value:
                 data_col = "('{0}', '{1}')".format(key, aggfunc) if multi_aggs else key
                 data_columns.append(data_col)
@@ -138,18 +146,18 @@ def aggregate(self, func_or_funcs, *args, **kwargs):
                 else:
                     reordered.append(F.expr('{1}(`{0}`) as `{2}`'.format(key, aggfunc, data_col)))
         sdf = sdf.groupby(*groupkey_cols).agg(*reordered)
-        internal = _InternalFrame(sdf=sdf,
+        if len(groupkeys) > 0:
+            index_map = [('__index_level_{}__'.format(i),
+                          s._internal.column_index[0])
+                         for i, s in enumerate(groupkeys)]
+            return _InternalFrame(sdf=sdf,
                                   data_columns=data_columns,
                                   column_index=column_index if multi_aggs else None,
-                                  index_map=[('__index_level_{}__'.format(i),
-                                              s._internal.column_index[0])
-                                             for i, s in enumerate(groupkeys)])
-        kdf = DataFrame(internal)
-        if not self._as_index:
-            kdf = kdf.reset_index()
-        return kdf
-
-    agg = aggregate
+                                  index_map=index_map)
+        else:
+            return _InternalFrame(sdf=sdf,
+                                  data_columns=data_columns,
+                                  column_index=column_index if multi_aggs else None)
 
     def count(self):
         """