Repository: databricks/koalas

Issue:
DataFrame.sum( axis=1 ) - axis parameter for df operations 
I'm missing axis parameter in DataFrame functions.
Sum function is just an example, I'm missing it among many others, like `apply()`, `mean()` etc.:
```python
import pandas as pd
import databricks.koalas as ks

with open("input_file2.csv", 'r') as content_file:
    print("Content of the CSV file")
    print(content_file.read())
    
print("pandas behavior:")
print(pd.read_csv("input_file2.csv").sum(axis=1))

print("koalas behavior:")
print(ks.read_csv("input_file2.csv").sum(axis=1))
```
prints:
```bash
Content of the CSV file
A,B,C,D
item1,1,1,1
item2,2,2,2

pandas behavior:
0    3
1    6
dtype: int64
koalas behavior:
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-24-1b42310256f7> in <module>
     13 
     14 print("koalas behavior:")
---> 15 print(ks.read_csv("input_file2.csv").sum(axis=1))

TypeError: sum() got an unexpected keyword argument 'axis'
```
Workaround I'm using for now (very inefficient):
```python
print("koalas workaround:")
kdf = ks.read_csv("input_file2.csv")
kdf = kdf.set_index("A")
kdf.transpose().sum()
```
prints:
```bash
koalas workaround:
item1    3
item2    6
dtype: int64
```



Discussion:

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -39,6 +39,7 @@
                                IntegerType, LongType, NumericType, ShortType, StructType)
 from pyspark.sql.utils import AnalysisException
 from pyspark.sql.window import Window
+from pyspark.sql.functions import pandas_udf
 
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
 from databricks.koalas.utils import validate_arguments_and_invoke_function
@@ -47,6 +48,7 @@
 from databricks.koalas.missing.frame import _MissingPandasLikeDataFrame
 from databricks.koalas.ml import corr
 from databricks.koalas.utils import scol_for
+from databricks.koalas.typedef import as_spark_type
 
 
 # These regular expression patterns are complied and defined here to avoid to compile the same
@@ -306,7 +308,7 @@ def __init__(self, data=None, index=None, columns=None, dtype=None, copy=False):
     def _sdf(self) -> spark.DataFrame:
         return self._internal.sdf
 
-    def _reduce_for_stat_function(self, sfun, numeric_only=False):
+    def _reduce_for_stat_function(self, sfun, name, axis=None, numeric_only=False):
         """
         Applies sfun to each column and returns a pd.Series where the number of rows equal the
         number of columns.
@@ -314,42 +316,68 @@ def _reduce_for_stat_function(self, sfun, numeric_only=False):
         Parameters
         ----------
         sfun : either an 1-arg function that takes a Column and returns a Column, or
-        a 2-arg function that takes a Column and its DataType and returns a Column.
+            a 2-arg function that takes a Column and its DataType and returns a Column.
+            axis: used only for sanity check because series only support index axis.
+        name : original pandas API name.
+        axis : axis to apply. 0 or 1, or 'index' or 'columns.
         numeric_only : boolean, default False
             If True, sfun is applied on numeric columns (including booleans) only.
         """
         from inspect import signature
-        exprs = []
-        num_args = len(signature(sfun).parameters)
-        for col in self.columns:
-            col_sdf = self._internal.scol_for(col)
-            col_type = self._internal.spark_type_for(col)
-
-            is_numeric_or_boolean = isinstance(col_type, (NumericType, BooleanType))
-            min_or_max = sfun.__name__ in ('min', 'max')
-            keep_column = not numeric_only or is_numeric_or_boolean or min_or_max
-
-            if keep_column:
-                if isinstance(col_type, BooleanType) and not min_or_max:
-                    # Stat functions cannot be used with boolean values by default
-                    # Thus, cast to integer (true to 1 and false to 0)
-                    # Exclude the min and max methods though since those work with booleans
-                    col_sdf = col_sdf.cast('integer')
-                if num_args == 1:
-                    # Only pass in the column if sfun accepts only one arg
-                    col_sdf = sfun(col_sdf)
-                else:  # must be 2
-                    assert num_args == 2
-                    # Pass in both the column and its data type if sfun accepts two args
-                    col_sdf = sfun(col_sdf, col_type)
-                exprs.append(col_sdf.alias(col))
-
-        sdf = self._sdf.select(*exprs)
-        pdf = sdf.toPandas()
-        assert len(pdf) == 1, (sdf, pdf)
-        row = pdf.iloc[0]
-        row.name = None
-        return row  # Return first row as a Series
+        from databricks.koalas import Series
+
+        if axis in ('index', 0, None):
+            exprs = []
+            num_args = len(signature(sfun).parameters)
+            for col in self.columns:
+                col_sdf = self._internal.scol_for(col)
+                col_type = self._internal.spark_type_for(col)
+
+                is_numeric_or_boolean = isinstance(col_type, (NumericType, BooleanType))
+                min_or_max = sfun.__name__ in ('min', 'max')
+                keep_column = not numeric_only or is_numeric_or_boolean or min_or_max
+
+                if keep_column:
+                    if isinstance(col_type, BooleanType) and not min_or_max:
+                        # Stat functions cannot be used with boolean values by default
+                        # Thus, cast to integer (true to 1 and false to 0)
+                        # Exclude the min and max methods though since those work with booleans
+                        col_sdf = col_sdf.cast('integer')
+                    if num_args == 1:
+                        # Only pass in the column if sfun accepts only one arg
+                        col_sdf = sfun(col_sdf)
+                    else:  # must be 2
+                        assert num_args == 2
+                        # Pass in both the column and its data type if sfun accepts two args
+                        col_sdf = sfun(col_sdf, col_type)
+                    exprs.append(col_sdf.alias(col))
+
+            sdf = self._sdf.select(*exprs)
+            pdf = sdf.toPandas()
+            assert len(pdf) == 1, (sdf, pdf)
+            row = pdf.iloc[0]
+            row.name = None
+            # TODO: return Koalas series.
+            return row  # Return first row as a Series
+
+        elif axis in ('columns', 1):
+            # Here we execute with the first 1000 to get the return type.
+            # If the records were less than 1000, it uses pandas API directly for a shortcut.
+            limit = 1000
+            pdf = self.head(limit + 1).to_pandas()
+            pser = getattr(pdf, name)(axis=axis, numeric_only=numeric_only)
+            if len(pdf) <= limit:
+                return Series(pser)
+
+            @pandas_udf(returnType=as_spark_type(pser.dtype.type))
+            def calculate_columns_axis(*cols):
+                return getattr(pd.concat(cols, axis=1), name)(axis=axis, numeric_only=numeric_only)
+
+            df = self._sdf.select(calculate_columns_axis(*self._internal.data_scols).alias("0"))
+            return DataFrame(df)["0"]
+
+        else:
+            raise ValueError("No axis named %s for object type %s." % (axis, type(axis)))
 
     # Arithmetic Operators
     def _map_series_op(self, op, other):
@@ -3886,12 +3914,18 @@ def select_dtypes(self, include=None, exclude=None):
                                  [scol_for(self._sdf, col) for col in columns]),
             data_columns=columns))
 
-    def count(self):
+    def count(self, axis=None):
         """
         Count non-NA cells for each column.
 
         The values `None`, `NaN` are considered NA.
 
+        Parameters
+        ----------
+        axis : {0 or ‘index’, 1 or ‘columns’}, default 0
+            If 0 or ‘index’ counts are generated for each column. If 1 or ‘columns’ counts are
+            generated for each row.
+
         Returns
         -------
         pandas.Series
@@ -3928,8 +3962,17 @@ def count(self):
         Age       4
         Single    5
         dtype: int64
+
+        >>> df.count(axis=1)
+        0    3
+        1    2
+        2    3
+        3    3
+        4    3
+        Name: 0, dtype: int64
         """
-        return self._reduce_for_stat_function(_Frame._count_expr, numeric_only=False)
+        return self._reduce_for_stat_function(
+            _Frame._count_expr, name="count", axis=axis, numeric_only=False)
 
     def drop(self, labels=None, axis=1, columns: Union[str, List[str]] = None):
         """
diff --git a/databricks/koalas/generic.py b/databricks/koalas/generic.py
--- a/databricks/koalas/generic.py
+++ b/databricks/koalas/generic.py
@@ -839,12 +839,14 @@ def to_excel(self, excel_writer, sheet_name="Sheet1", na_rep="", float_format=No
         return validate_arguments_and_invoke_function(
             kdf.to_pandas(), self.to_excel, f, args)
 
-    def mean(self, numeric_only=True):
+    def mean(self, axis=None, numeric_only=True):
         """
         Return the mean of the values.
 
         Parameters
         ----------
+        axis : {index (0), columns (1)}
+            Axis for the function to be applied on.
         numeric_only : bool, default None
             Include only float, int, boolean columns. If None, will attempt to use
             everything, then use only numeric data. Not implemented for Series.
@@ -866,19 +868,29 @@ def mean(self, numeric_only=True):
         b    0.2
         dtype: float64
 
+        >>> df.mean(axis=1)
+        0    0.55
+        1    1.10
+        2    1.65
+        3     NaN
+        Name: 0, dtype: float64
+
         On a Series:
 
         >>> df['a'].mean()
         2.0
         """
-        return self._reduce_for_stat_function(F.mean, numeric_only=numeric_only)
+        return self._reduce_for_stat_function(
+            F.mean, name="mean", numeric_only=numeric_only, axis=axis)
 
-    def sum(self, numeric_only=True):
+    def sum(self, axis=None, numeric_only=True):
         """
         Return the sum of the values.
 
         Parameters
         ----------
+        axis : {index (0), columns (1)}
+            Axis for the function to be applied on.
         numeric_only : bool, default None
             Include only float, int, boolean columns. If None, will attempt to use
             everything, then use only numeric data. Not implemented for Series.
@@ -900,19 +912,29 @@ def sum(self, numeric_only=True):
         b    0.6
         dtype: float64
 
+        >>> df.sum(axis=1)
+        0    1.1
+        1    2.2
+        2    3.3
+        3    0.0
+        Name: 0, dtype: float64
+
         On a Series:
 
         >>> df['a'].sum()
         6.0
         """
-        return self._reduce_for_stat_function(F.sum, numeric_only=numeric_only)
+        return self._reduce_for_stat_function(
+            F.sum, name="sum", numeric_only=numeric_only, axis=axis)
 
-    def skew(self, numeric_only=True):
+    def skew(self, axis=None, numeric_only=True):
         """
         Return unbiased skew normalized by N-1.
 
         Parameters
         ----------
+        axis : {index (0), columns (1)}
+            Axis for the function to be applied on.
         numeric_only : bool, default None
             Include only float, int, boolean columns. If None, will attempt to use
             everything, then use only numeric data. Not implemented for Series.
@@ -939,15 +961,18 @@ def skew(self, numeric_only=True):
         >>> df['a'].skew()
         0.0
         """
-        return self._reduce_for_stat_function(F.skewness, numeric_only=numeric_only)
+        return self._reduce_for_stat_function(
+            F.skewness, name="skew", numeric_only=numeric_only, axis=axis)
 
-    def kurtosis(self, numeric_only=True):
+    def kurtosis(self, axis=None, numeric_only=True):
         """
         Return unbiased kurtosis using Fisher’s definition of kurtosis (kurtosis of normal == 0.0).
         Normalized by N-1.
 
         Parameters
         ----------
+        axis : {index (0), columns (1)}
+            Axis for the function to be applied on.
         numeric_only : bool, default None
             Include only float, int, boolean columns. If None, will attempt to use
             everything, then use only numeric data. Not implemented for Series.
@@ -974,16 +999,19 @@ def kurtosis(self, numeric_only=True):
         >>> df['a'].kurtosis()
         -1.5
         """
-        return self._reduce_for_stat_function(F.kurtosis, numeric_only=numeric_only)
+        return self._reduce_for_stat_function(
+            F.kurtosis, name="kurtosis", numeric_only=numeric_only, axis=axis)
 
     kurt = kurtosis
 
-    def min(self, numeric_only=False):
+    def min(self, axis=None, numeric_only=False):
         """
         Return the minimum of the values.
 
         Parameters
         ----------
+        axis : {index (0), columns (1)}
+            Axis for the function to be applied on.
         numeric_only : bool, default None
             Include only float, int, boolean columns. If None, will attempt to use
             everything, then use only numeric data. Not implemented for Series.
@@ -1005,19 +1033,29 @@ def min(self, numeric_only=False):
         b    0.1
         dtype: float64
 
+        >>> df.min(axis=1)
+        0    0.1
+        1    0.2
+        2    0.3
+        3    NaN
+        Name: 0, dtype: float64
+
         On a Series:
 
         >>> df['a'].min()
         1.0
         """
-        return self._reduce_for_stat_function(F.min, numeric_only=numeric_only)
+        return self._reduce_for_stat_function(
+            F.min, name="min", numeric_only=numeric_only, axis=axis)
 
-    def max(self, numeric_only=False):
+    def max(self, axis=None, numeric_only=False):
         """
         Return the maximum of the values.
 
         Parameters
         ----------
+        axis : {index (0), columns (1)}
+            Axis for the function to be applied on.
         numeric_only : bool, default None
             Include only float, int, boolean columns. If None, will attempt to use
             everything, then use only numeric data. Not implemented for Series.
@@ -1039,19 +1077,29 @@ def max(self, numeric_only=False):
         b    0.3
         dtype: float64
 
+        >>> df.max(axis=1)
+        0    1.0
+        1    2.0
+        2    3.0
+        3    NaN
+        Name: 0, dtype: float64
+
         On a Series:
 
         >>> df['a'].max()
         3.0
         """
-        return self._reduce_for_stat_function(F.max, numeric_only=numeric_only)
+        return self._reduce_for_stat_function(
+            F.max, name="max", numeric_only=numeric_only, axis=axis)
 
-    def std(self, numeric_only=True):
+    def std(self, axis=None, numeric_only=True):
         """
         Return sample standard deviation.
 
         Parameters
         ----------
+        axis : {index (0), columns (1)}
+            Axis for the function to be applied on.
         numeric_only : bool, default None
             Include only float, int, boolean columns. If None, will attempt to use
             everything, then use only numeric data. Not implemented for Series.
@@ -1073,19 +1121,29 @@ def std(self, numeric_only=True):
         b    0.1
         dtype: float64
 
+        >>> df.std(axis=1)
+        0    0.636396
+        1    1.272792
+        2    1.909188
+        3         NaN
+        Name: 0, dtype: float64
+
         On a Series:
 
         >>> df['a'].std()
         1.0
         """
-        return self._reduce_for_stat_function(F.stddev, numeric_only=numeric_only)
+        return self._reduce_for_stat_function(
+            F.stddev, name="std", numeric_only=numeric_only, axis=axis)
 
-    def var(self, numeric_only=True):
+    def var(self, axis=None, numeric_only=True):
         """
         Return unbiased variance.
 
         Parameters
         ----------
+        axis : {index (0), columns (1)}
+            Axis for the function to be applied on.
         numeric_only : bool, default None
             Include only float, int, boolean columns. If None, will attempt to use
             everything, then use only numeric data. Not implemented for Series.
@@ -1107,12 +1165,20 @@ def var(self, numeric_only=True):
         b    0.01
         dtype: float64
 
+        >>> df.var(axis=1)
+        0    0.405
+        1    1.620
+        2    3.645
+        3      NaN
+        Name: 0, dtype: float64
+
         On a Series:
 
         >>> df['a'].var()
         1.0
         """
-        return self._reduce_for_stat_function(F.variance, numeric_only=numeric_only)
+        return self._reduce_for_stat_function(
+            F.variance, name="var", numeric_only=numeric_only, axis=axis)
 
     @property
     def size(self) -> int:
@@ -1342,7 +1408,8 @@ def median(self, accuracy=10000):
         if isinstance(kdf_or_ks, Series):
             ks = kdf_or_ks
             return self._reduce_for_stat_function(
-                lambda _: F.expr("approx_percentile(`%s`, 0.5, %s)" % (ks.name, accuracy)))
+                lambda _: F.expr(
+                    "approx_percentile(`%s`, 0.5, %s)" % (ks.name, accuracy)), name="median")
         assert isinstance(kdf_or_ks, DataFrame)
 
         # This code path cannot reuse `_reduce_for_stat_function` since there looks no proper way
diff --git a/databricks/koalas/series.py b/databricks/koalas/series.py
--- a/databricks/koalas/series.py
+++ b/databricks/koalas/series.py
@@ -1972,7 +1972,7 @@ def count(self):
         >>> df['Age'].count()
         4
         """
-        return self._reduce_for_stat_function(_Frame._count_expr)
+        return self._reduce_for_stat_function(_Frame._count_expr, name="count")
 
     def append(self, to_append: 'Series', ignore_index: bool = False,
                verify_integrity: bool = False) -> 'Series':
@@ -2369,7 +2369,8 @@ def quantile(self, q=0.5, accuracy=10000):
             return ser
         else:
             return self._reduce_for_stat_function(
-                lambda _: F.expr("approx_percentile(`%s`, %s, %s)" % (self.name, q, accuracy)))
+                lambda _: F.expr("approx_percentile(`%s`, %s, %s)" % (self.name, q, accuracy)),
+                name="median")
 
     # TODO: add axis, numeric_only, pct, na_option parameter
     def rank(self, method='average', ascending=True):
@@ -2833,12 +2834,20 @@ def _cum(self, func, skipna):
 
     # ----------------------------------------------------------------------
 
-    def _reduce_for_stat_function(self, sfun, numeric_only=None):
+    def _reduce_for_stat_function(self, sfun, name, axis=None, numeric_only=None):
         """
-        :param sfun: the stats function to be used for aggregation
-        :param numeric_only: not used by this implementation, but passed down by stats functions
+        Applies sfun to the column and returns a scalar
+
+        Parameters
+        ----------
+        sfun : the stats function to be used for aggregation
+        name : original pandas API name.
+        axis : used only for sanity check because series only support index axis.
+        numeric_only : not used by this implementation, but passed down by stats functions
         """
         from inspect import signature
+        if axis in ('columns', 1):
+            raise ValueError("Series does not support columns axis.")
         num_args = len(signature(sfun).parameters)
         col_sdf = self._scol
         col_type = self.spark_type


Test patch:
diff --git a/databricks/koalas/tests/test_stats.py b/databricks/koalas/tests/test_stats.py
--- a/databricks/koalas/tests/test_stats.py
+++ b/databricks/koalas/tests/test_stats.py
@@ -56,6 +56,24 @@ def test_abs(self):
         self.assert_eq(kdf[['B', 'C']].abs(), pdf[['B', 'C']].abs())
         # self.assert_eq(kdf.select('A', 'B').abs(), pdf[['A', 'B']].abs())
 
+    def test_axis_on_dataframe(self):
+        # The number of each count is intentionally big
+        # because when data is small, it executes a shortcut.
+        pdf = pd.DataFrame({'A': [1, -2, 3, -4, 5] * 300,
+                            'B': [1., -2, 3, -4, 5] * 300,
+                            'C': [-6., -7, -8, -9, 10] * 300,
+                            'D': [True, False, True, False, False] * 300})
+        kdf = koalas.from_pandas(pdf)
+        self.assert_eq(kdf.count(axis=1), pdf.count(axis=1))
+        self.assert_eq(kdf.var(axis=1), pdf.var(axis=1))
+        self.assert_eq(kdf.std(axis=1), pdf.std(axis=1))
+        self.assert_eq(kdf.max(axis=1), pdf.max(axis=1))
+        self.assert_eq(kdf.min(axis=1), pdf.min(axis=1))
+        self.assert_eq(kdf.sum(axis=1), pdf.sum(axis=1))
+        self.assert_eq(kdf.kurtosis(axis=1), pdf.kurtosis(axis=1))
+        self.assert_eq(kdf.skew(axis=1), pdf.skew(axis=1))
+        self.assert_eq(kdf.mean(axis=1), pdf.mean(axis=1))
+
     def test_corr(self):
         # Disable arrow execution since corr() is using UDT internally which is not supported.
         with self.sql_conf({'spark.sql.execution.arrow.enabled': False}):