Expose spark_column from Series/Index to make it easier to work with Spark Columns. (#1438)

This PR is exposing `spark_column` property representing the Series/Index for users who are familiar with Spark functions to make it easier to work with them.

E.g.:

```py
>>> kdf = ks.DataFrame({'a': [1.0, 1.0, 1.0, 2.0, 2.0, 2.0], 'b': [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]})

>>> from pyspark.sql import functions as F
>>> kdf['greatest'] = F.greatest(kdf.a.spark_column, kdf.b.spark_column)
>>> kdf['least'] = F.least(kdf.a.spark_column, kdf.b.spark_column)
>>> kdf
     a    b  greatest  least
0  1.0  1.0       1.0    1.0
1  1.0  2.0       2.0    1.0
2  1.0  3.0       3.0    1.0
3  2.0  4.0       4.0    2.0
4  2.0  5.0       5.0    2.0
5  2.0  6.0       6.0    2.0
```

Author: ueshin

databricks/koalas/base.py

@@ -83,9 +83,9 @@ def wrapper(self, *args):
         cols = [arg for arg in args if isinstance(arg, IndexOpsMixin)]
         if all(self._kdf is col._kdf for col in cols):
             # Same DataFrame anchors
-            args = [arg._scol if isinstance(arg, IndexOpsMixin) else arg for arg in args]
-            scol = f(self._scol, *args)
-            scol = booleanize_null(self._scol, scol, f)
+            args = [arg.spark_column if isinstance(arg, IndexOpsMixin) else arg for arg in args]
+            scol = f(self.spark_column, *args)
+            scol = booleanize_null(self.spark_column, scol, f)
 
             return self._with_new_scol(scol)
         else:
@@ -154,7 +154,13 @@ def __init__(self, internal: _InternalFrame, kdf):
         self._kdf = kdf
 
     @property
-    def _scol(self):
+    def spark_column(self):
+        """
+        Spark Column object representing the Series/Index.
+
+        .. note:: This Spark Column object is strictly stick to its base DataFrame the Series/Index
+            was derived from.
+        """
         return self._internal.spark_column
 
     # arithmetic operators
@@ -202,7 +208,7 @@ def mod(left, right):
     def __radd__(self, other):
         # Handle 'literal' + df['col']
         if isinstance(self.spark_type, StringType) and isinstance(other, str):
-            return self._with_new_scol(F.concat(F.lit(other), self._scol))
+            return self._with_new_scol(F.concat(F.lit(other), self.spark_column))
         else:
             return _column_op(spark.Column.__radd__)(self, other)
 
@@ -336,8 +342,8 @@ def hasnans(self):
         >>> ks.Series([1, 2, 3]).rename("a").to_frame().set_index("a").index.hasnans
         False
         """
-        sdf = self._internal._sdf.select(self._scol)
-        col = self._scol
+        sdf = self._internal._sdf.select(self.spark_column)
+        col = self.spark_column
 
         ret = sdf.select(F.max(col.isNull() | F.isnan(col))).collect()[0][0]
         return ret
@@ -517,7 +523,7 @@ def _is_monotonic(self, order):
                     "__partition_id"
                 ),  # Make sure we use the same partition id in the whole job.
                 F.col(NATURAL_ORDER_COLUMN_NAME),
-                self._scol.alias("__origin"),
+                self.spark_column.alias("__origin"),
             )
             .select(
                 F.col("__partition_id"),
@@ -635,7 +641,7 @@ def astype(self, dtype):
         spark_type = as_spark_type(dtype)
         if not spark_type:
             raise ValueError("Type {} not understood".format(dtype))
-        return self._with_new_scol(self._scol.cast(spark_type))
+        return self._with_new_scol(self.spark_column.cast(spark_type))
 
     def isin(self, values):
         """
@@ -687,7 +693,7 @@ def isin(self, values):
                 " to isin(), you passed a [{values_type}]".format(values_type=type(values).__name__)
             )
 
-        return self._with_new_scol(self._scol.isin(list(values))).rename(self.name)
+        return self._with_new_scol(self.spark_column.isin(list(values))).rename(self.name)
 
     def isnull(self):
         """
@@ -721,9 +727,11 @@ def isnull(self):
         if isinstance(self, MultiIndex):
             raise NotImplementedError("isna is not defined for MultiIndex")
         if isinstance(self.spark_type, (FloatType, DoubleType)):
-            return self._with_new_scol(self._scol.isNull() | F.isnan(self._scol)).rename(self.name)
+            return self._with_new_scol(
+                self.spark_column.isNull() | F.isnan(self.spark_column)
+            ).rename(self.name)
         else:
-            return self._with_new_scol(self._scol.isNull()).rename(self.name)
+            return self._with_new_scol(self.spark_column.isNull()).rename(self.name)
 
     isna = isnull
 
@@ -819,7 +827,7 @@ def all(self, axis: Union[int, str] = 0) -> bool:
         if axis != 0:
             raise NotImplementedError('axis should be either 0 or "index" currently.')
 
-        sdf = self._internal._sdf.select(self._scol)
+        sdf = self._internal._sdf.select(self.spark_column)
         col = scol_for(sdf, sdf.columns[0])
 
         # Note that we're ignoring `None`s here for now.
@@ -882,7 +890,7 @@ def any(self, axis: Union[int, str] = 0) -> bool:
         if axis != 0:
             raise NotImplementedError('axis should be either 0 or "index" currently.')
 
-        sdf = self._internal._sdf.select(self._scol)
+        sdf = self._internal._sdf.select(self.spark_column)
         col = scol_for(sdf, sdf.columns[0])
 
         # Note that we're ignoring `None`s here for now.
@@ -949,7 +957,7 @@ def _shift(self, periods, fill_value, part_cols=()):
         if not isinstance(periods, int):
             raise ValueError("periods should be an int; however, got [%s]" % type(periods))
 
-        col = self._scol
+        col = self.spark_column
         window = (
             Window.partitionBy(*part_cols)
             .orderBy(NATURAL_ORDER_COLUMN_NAME)
@@ -1115,9 +1123,9 @@ def value_counts(self, normalize=False, sort=True, ascending=False, bins=None, d
             raise NotImplementedError("value_counts currently does not support bins")
 
         if dropna:
-            sdf_dropna = self._internal._sdf.select(self._scol).dropna()
+            sdf_dropna = self._internal._sdf.select(self.spark_column).dropna()
         else:
-            sdf_dropna = self._internal._sdf.select(self._scol)
+            sdf_dropna = self._internal._sdf.select(self.spark_column)
         index_name = SPARK_DEFAULT_INDEX_NAME
         column_name = self._internal.data_spark_column_names[0]
         sdf = sdf_dropna.groupby(scol_for(sdf_dropna, column_name).alias(index_name)).count()
@@ -1207,13 +1215,13 @@ def _nunique(self, dropna=True, approx=False, rsd=0.05):
         colname = self._internal.data_spark_column_names[0]
         count_fn = partial(F.approx_count_distinct, rsd=rsd) if approx else F.countDistinct
         if dropna:
-            return count_fn(self._scol).alias(colname)
+            return count_fn(self.spark_column).alias(colname)
         else:
             return (
-                count_fn(self._scol)
-                + F.when(F.count(F.when(self._scol.isNull(), 1).otherwise(None)) >= 1, 1).otherwise(
-                    0
-                )
+                count_fn(self.spark_column)
+                + F.when(
+                    F.count(F.when(self.spark_column.isNull(), 1).otherwise(None)) >= 1, 1
+                ).otherwise(0)
             ).alias(colname)
 
     def take(self, indices):

databricks/koalas/frame.py

@@ -2858,10 +2858,10 @@ def where(self, cond, other=np.nan):
         for label in self._internal.column_labels:
             data_spark_columns.append(
                 F.when(
-                    kdf[tmp_cond_col_name(name_like_string(label))]._scol,
+                    kdf[tmp_cond_col_name(name_like_string(label))].spark_column,
                     kdf._internal.spark_column_for(label),
                 )
-                .otherwise(kdf[tmp_other_col_name(name_like_string(label))]._scol)
+                .otherwise(kdf[tmp_other_col_name(name_like_string(label))].spark_column)
                 .alias(kdf._internal.spark_column_name_for(label))
             )
 
@@ -3715,7 +3715,7 @@ def round(self, decimals=0):
         def op(kser):
             label = kser._internal.column_labels[0]
             if label in decimals:
-                return F.round(kser._scol, decimals[label]).alias(
+                return F.round(kser.spark_column, decimals[label]).alias(
                     kser._internal.data_spark_column_names[0]
                 )
             else:
@@ -4541,7 +4541,11 @@ def _assign(self, kwargs):
 
         pairs = {
             (k if isinstance(k, tuple) else (k,)): (
-                v._scol if isinstance(v, Series) else v if isinstance(v, spark.Column) else F.lit(v)
+                v.spark_column
+                if isinstance(v, Series)
+                else v
+                if isinstance(v, spark.Column)
+                else F.lit(v)
             )
             for k, v in kwargs.items()
         }
@@ -4842,7 +4846,10 @@ def dropna(self, axis=0, how="any", thresh=None, subset=None, inplace=False):
 
             cnt = reduce(
                 lambda x, y: x + y,
-                [F.when(self._kser_for(label).notna()._scol, 1).otherwise(0) for label in labels],
+                [
+                    F.when(self._kser_for(label).notna().spark_column, 1).otherwise(0)
+                    for label in labels
+                ],
                 F.lit(0),
             )
             if thresh is not None:
@@ -5315,7 +5322,7 @@ def clip(self, lower: Union[float, int] = None, upper: Union[float, int] = None)
 
         def op(kser):
             if isinstance(kser.spark_type, numeric_types):
-                scol = kser._scol
+                scol = kser.spark_column
                 if lower is not None:
                     scol = F.when(scol < lower, lower).otherwise(scol)
                 if upper is not None:
@@ -6374,7 +6381,7 @@ def sort_values(
                     "The column %s is not unique. For a multi-index, the label must be a tuple "
                     "with elements corresponding to each level." % name_like_string(colname)
                 )
-            new_by.append(ser._scol)
+            new_by.append(ser.spark_column)
 
         return self._sort(by=new_by, ascending=ascending, inplace=inplace, na_position=na_position)
 
@@ -8036,7 +8043,7 @@ def _reindex_index(self, index):
         index_column = self._internal.index_spark_column_names[0]
 
         kser = ks.Series(list(index))
-        labels = kser._internal._sdf.select(kser._scol.alias(index_column))
+        labels = kser._internal._sdf.select(kser.spark_column.alias(index_column))
 
         joined_df = self._sdf.drop(NATURAL_ORDER_COLUMN_NAME).join(
             labels, on=index_column, how="right"
@@ -9275,8 +9282,8 @@ def pct_change(self, periods=1):
         window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween(-periods, -periods)
 
         def op(kser):
-            prev_row = F.lag(kser._scol, periods).over(window)
-            return ((kser._scol - prev_row) / prev_row).alias(
+            prev_row = F.lag(kser.spark_column, periods).over(window)
+            return ((kser.spark_column - prev_row) / prev_row).alias(
                 kser._internal.data_spark_column_names[0]
             )
 

databricks/koalas/generic.py

@@ -1395,7 +1395,7 @@ def abs(self):
         """
         # TODO: The first example above should not have "Name: 0".
         return self._apply_series_op(
-            lambda kser: kser._with_new_scol(F.abs(kser._scol)).rename(kser.name)
+            lambda kser: kser._with_new_scol(F.abs(kser.spark_column)).rename(kser.name)
         )
 
     # TODO: by argument only support the grouping name and as_index only for now. Documentation

databricks/koalas/groupby.py

@@ -1912,7 +1912,7 @@ def __init__(
     ):
         self._kdf = kdf
         self._groupkeys = by
-        self._groupkeys_scols = [s._scol for s in self._groupkeys]
+        self._groupkeys_scols = [s.spark_column for s in self._groupkeys]
         self._as_index = as_index
         self._should_drop_index = should_drop_index
         self._have_agg_columns = True
@@ -1925,7 +1925,7 @@ def __init__(
             ]
             self._have_agg_columns = False
         self._agg_columns = [kdf[label] for label in agg_columns]
-        self._agg_columns_scols = [s._scol for s in self._agg_columns]
+        self._agg_columns_scols = [s.spark_column for s in self._agg_columns]
 
     def __getattr__(self, item: str) -> Any:
         if hasattr(_MissingPandasLikeDataFrameGroupBy, item):

databricks/koalas/indexes.py

@@ -142,7 +142,7 @@ def _summary(self, name=None):
         String with a summarized representation of the index
         """
         head, tail, total_count = self._kdf._sdf.select(
-            F.first(self._scol), F.last(self._scol), F.count(F.expr("*"))
+            F.first(self.spark_column), F.last(self.spark_column), F.count(F.expr("*"))
         ).first()
 
         if total_count > 0:
@@ -440,7 +440,7 @@ def has_duplicates(self) -> bool:
         >>> kdf.index.has_duplicates
         True
         """
-        df = self._kdf._sdf.select(self._scol)
+        df = self._kdf._sdf.select(self.spark_column)
         col = df.columns[0]
 
         return df.select(F.count(col) != F.countDistinct(col)).first()[0]
@@ -554,7 +554,7 @@ def rename(
         )
 
         idx = kdf.index
-        idx._internal._scol = self._scol
+        idx._internal = idx._internal.copy(spark_column=self.spark_column)
         if inplace:
             self._internal = idx._internal
         else:
@@ -664,7 +664,7 @@ def to_series(self, name: Union[str, Tuple[str, ...]] = None) -> Series:
         Name: 0, dtype: object
         """
         kdf = self._kdf
-        scol = self._scol
+        scol = self.spark_column
         if name is not None:
             scol = scol.alias(name_like_string(name))
         column_labels = [None] if len(kdf._internal.index_map) > 1 else kdf._internal.index_names
@@ -731,7 +731,7 @@ def to_frame(self, index=True, name=None) -> DataFrame:
                 name = self._internal.index_names[0]
         elif isinstance(name, str):
             name = (name,)
-        scol = self._scol.alias(name_like_string(name))
+        scol = self.spark_column.alias(name_like_string(name))
 
         sdf = self._internal.spark_frame.select(scol, NATURAL_ORDER_COLUMN_NAME)
 
@@ -1370,7 +1370,7 @@ def argmax(self):
         >>> kidx.argmax()
         4
         """
-        sdf = self._internal.spark_frame.select(self._scol)
+        sdf = self._internal.spark_frame.select(self.spark_column)
         sequence_col = verify_temp_column_name(sdf, "__distributed_sequence_column__")
         sdf = _InternalFrame.attach_distributed_sequence_column(sdf, column_name=sequence_col)
         # spark_frame here looks like below
@@ -1388,7 +1388,7 @@ def argmax(self):
         # |                1|              9|
         # +-----------------+---------------+
 
-        return sdf.orderBy(self._scol.desc(), F.col(sequence_col).asc()).first()[0]
+        return sdf.orderBy(self.spark_column.desc(), F.col(sequence_col).asc()).first()[0]
 
     def argmin(self):
         """
@@ -1411,11 +1411,11 @@ def argmin(self):
         >>> kidx.argmin()
         7
         """
-        sdf = self._internal.spark_frame.select(self._scol)
+        sdf = self._internal.spark_frame.select(self.spark_column)
         sequence_col = verify_temp_column_name(sdf, "__distributed_sequence_column__")
         sdf = _InternalFrame.attach_distributed_sequence_column(sdf, column_name=sequence_col)
 
-        return sdf.orderBy(self._scol.asc(), F.col(sequence_col).asc()).first()[0]
+        return sdf.orderBy(self.spark_column.asc(), F.col(sequence_col).asc()).first()[0]
 
     def set_names(self, names, level=None, inplace=False):
         """
@@ -1689,9 +1689,9 @@ def asof(self, label):
         """
         sdf = self._internal._sdf
         if self.is_monotonic_increasing:
-            sdf = sdf.where(self._scol <= label).select(F.max(self._scol))
+            sdf = sdf.where(self.spark_column <= label).select(F.max(self.spark_column))
         elif self.is_monotonic_decreasing:
-            sdf = sdf.where(self._scol >= label).select(F.min(self._scol))
+            sdf = sdf.where(self.spark_column >= label).select(F.min(self.spark_column))
         else:
             raise ValueError("index must be monotonic increasing or decreasing")
         result = sdf.head()[0]
@@ -1780,7 +1780,11 @@ def __repr__(self):
         if max_display_count is None:
             return repr(self.to_pandas())
 
-        pindex = self._kdf.head(max_display_count + 1).index._with_new_scol(self._scol).to_pandas()
+        pindex = (
+            self._kdf.head(max_display_count + 1)
+            .index._with_new_scol(self.spark_column)
+            .to_pandas()
+        )
 
         pindex_length = len(pindex)
         repr_string = repr(pindex[:max_display_count])
@@ -2072,7 +2076,7 @@ def _is_monotonic(self, order):
             return self._is_monotonic_decreasing().all()
 
     def _is_monotonic_increasing(self):
-        scol = self._scol
+        scol = self.spark_column
         window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween(-1, -1)
         prev = F.lag(scol, 1).over(window)
 
@@ -2108,7 +2112,7 @@ def _comparator_for_monotonic_decreasing(data_type):
             return compare_null_first
 
     def _is_monotonic_decreasing(self):
-        scol = self._scol
+        scol = self.spark_column
         window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween(-1, -1)
         prev = F.lag(scol, 1).over(window)