Introduce column_scols in InternalFrame substitude for data_columns. (#956)

Author: ueshin

databricks/koalas/frame.py

@@ -33,7 +33,7 @@
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
 from databricks.koalas.indexing import AtIndexer, ILocIndexer, LocIndexer
 from databricks.koalas.internal import _InternalFrame
-from databricks.koalas.utils import validate_arguments_and_invoke_function
+from databricks.koalas.utils import validate_arguments_and_invoke_function, scol_for
 from databricks.koalas.window import Rolling, Expanding
 
 
@@ -1416,15 +1416,18 @@ def median(self, accuracy=10000):
         # This code path cannot reuse `_reduce_for_stat_function` since there looks no proper way
         # to get a column name from Spark column but we need it to pass it through `expr`.
         kdf = kdf_or_kser
-        sdf = kdf._sdf
+        sdf = kdf._sdf.select(kdf._internal.scols)
         median = lambda name: F.expr("approx_percentile(`%s`, 0.5, %s)" % (name, accuracy))
         sdf = sdf.select([median(col).alias(col) for col in kdf._internal.data_columns])
 
         # Attach a dummy column for index to avoid default index.
         sdf = sdf.withColumn('__DUMMY__', F.monotonically_increasing_id())
 
         # This is expected to be small so it's fine to transpose.
-        return DataFrame(kdf._internal.copy(sdf=sdf, index_map=[('__DUMMY__', None)])) \
+        return DataFrame(kdf._internal.copy(
+            sdf=sdf,
+            index_map=[('__DUMMY__', None)],
+            column_scols=[scol_for(sdf, col) for col in kdf._internal.data_columns])) \
             ._to_internal_pandas().transpose().iloc[:, 0]
 
     # TODO: 'center', 'win_type', 'on', 'axis' parameter should be implemented.

databricks/koalas/generic.py

@@ -219,8 +219,8 @@ def _spark_groupby(kdf, func, groupkeys):
         else:
             index_map = None
         return _InternalFrame(sdf=sdf,
-                              data_columns=data_columns,
                               column_index=column_index,
+                              column_scols=[scol_for(sdf, col) for col in data_columns],
                               index_map=index_map)
 
     def count(self):
@@ -493,10 +493,10 @@ def size(self):
         else:
             name = 'count'
         internal = _InternalFrame(sdf=sdf,
-                                  data_columns=[name],
                                   index_map=[(SPARK_INDEX_NAME_FORMAT(i),
                                               s._internal.column_index[0])
-                                             for i, s in enumerate(groupkeys)])
+                                             for i, s in enumerate(groupkeys)],
+                                  column_scols=[scol_for(sdf, name)])
         return _col(DataFrame(internal))
 
     def diff(self, periods=1):
@@ -893,7 +893,9 @@ def apply(self, func):
 
         if should_infer_schema:
             # If schema is inferred, we can restore indexes too.
-            internal = kdf._internal.copy(sdf=sdf)
+            internal = kdf._internal.copy(sdf=sdf,
+                                          column_scols=[scol_for(sdf, col)
+                                                        for col in kdf._internal.data_columns])
         else:
             # Otherwise, it loses index.
             internal = _InternalFrame(sdf=sdf)
@@ -945,7 +947,9 @@ def pandas_filter(pdf):
 
         sdf = self._spark_group_map_apply(
             pandas_filter, data_schema, retain_index=True)
-        return DataFrame(self._kdf._internal.copy(sdf=sdf))
+        return DataFrame(self._kdf._internal.copy(
+            sdf=sdf,
+            column_scols=[scol_for(sdf, col) for col in self._kdf._internal.data_columns]))
 
     def _spark_group_map_apply(self, func, return_schema, retain_index):
         index_columns = self._kdf._internal.index_columns
@@ -1153,13 +1157,13 @@ def idxmax(self, skipna=True):
             stat_exprs.append(F.max(scol_for(sdf, name)).alias(name))
         sdf = sdf.groupby(*groupkey_cols).agg(*stat_exprs)
         internal = _InternalFrame(sdf=sdf,
-                                  data_columns=[ks._internal.data_columns[0]
-                                                for ks in self._agg_columns],
-                                  column_index=[ks._internal.column_index[0]
-                                                for ks in self._agg_columns],
                                   index_map=[(SPARK_INDEX_NAME_FORMAT(i),
                                               s._internal.column_index[0])
-                                             for i, s in enumerate(groupkeys)])
+                                             for i, s in enumerate(groupkeys)],
+                                  column_index=[ks._internal.column_index[0]
+                                                for ks in self._agg_columns],
+                                  column_scols=[scol_for(sdf, ks._internal.data_columns[0])
+                                                for ks in self._agg_columns])
         return DataFrame(internal)
 
     # TODO: add axis parameter
@@ -1223,13 +1227,13 @@ def idxmin(self, skipna=True):
             stat_exprs.append(F.max(scol_for(sdf, name)).alias(name))
         sdf = sdf.groupby(*groupkey_cols).agg(*stat_exprs)
         internal = _InternalFrame(sdf=sdf,
-                                  data_columns=[ks._internal.data_columns[0]
-                                                for ks in self._agg_columns],
-                                  column_index=[ks._internal.column_index[0]
-                                                for ks in self._agg_columns],
                                   index_map=[(SPARK_INDEX_NAME_FORMAT(i),
                                               s._internal.column_index[0])
-                                             for i, s in enumerate(groupkeys)])
+                                             for i, s in enumerate(groupkeys)],
+                                  column_index=[ks._internal.column_index[0]
+                                                for ks in self._agg_columns],
+                                  column_scols=[scol_for(sdf, ks._internal.data_columns[0])
+                                                for ks in self._agg_columns])
         return DataFrame(internal)
 
     def fillna(self, value=None, method=None, axis=None, inplace=False, limit=None):
@@ -1581,7 +1585,9 @@ def pandas_transform(pdf):
             sdf = self._spark_group_map_apply(
                 pandas_transform, return_schema, retain_index=True)
             # If schema is inferred, we can restore indexes too.
-            internal = kdf._internal.copy(sdf=sdf)
+            internal = kdf._internal.copy(sdf=sdf,
+                                          column_scols=[scol_for(sdf, col)
+                                                        for col in kdf._internal.data_columns])
         else:
             return_type = _infer_return_type(func).tpe
             data_columns = self._kdf._internal.data_columns
@@ -1708,8 +1714,8 @@ def _reduce_for_stat_function(self, sfun, only_numeric):
                                   index_map=[(SPARK_INDEX_NAME_FORMAT(i),
                                               s._internal.column_index[0])
                                              for i, s in enumerate(groupkeys)],
-                                  data_columns=data_columns,
                                   column_index=column_index,
+                                  column_scols=[scol_for(sdf, col) for col in data_columns],
                                   column_index_names=self._kdf._internal.column_index_names)
         kdf = DataFrame(internal)
         if not self._as_index:
@@ -1767,8 +1773,9 @@ def _diff(self, *args, **kwargs):
 
         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])
+                                      column_index=[c._internal.column_index[0] for c in applied],
+                                      column_scols=[scol_for(sdf, c._internal.data_columns[0])
+                                                    for c in applied])
         return DataFrame(internal)
 
     def _rank(self, *args, **kwargs):
@@ -1781,8 +1788,9 @@ def _rank(self, *args, **kwargs):
 
         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])
+                                      column_index=[c._internal.column_index[0] for c in applied],
+                                      column_scols=[scol_for(sdf, c._internal.data_columns[0])
+                                                    for c in applied])
         return DataFrame(internal)
 
     def _cum(self, func):
@@ -1806,8 +1814,9 @@ def _cum(self, func):
         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])
+                                      column_index=[c._internal.column_index[0] for c in applied],
+                                      column_scols=[scol_for(sdf, c._internal.data_columns[0])
+                                                    for c in applied])
         return DataFrame(internal)
 
     def _fillna(self, *args, **kwargs):
@@ -1820,8 +1829,9 @@ def _fillna(self, *args, **kwargs):
 
         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])
+                                      column_index=[c._internal.column_index[0] for c in applied],
+                                      column_scols=[scol_for(sdf, c._internal.data_columns[0])
+                                                    for c in applied])
         return DataFrame(internal)
 
     def _shift(self, periods, fill_value):
@@ -1833,8 +1843,9 @@ def _shift(self, periods, fill_value):
 
         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])
+                                      column_index=[c._internal.column_index[0] for c in applied],
+                                      column_scols=[scol_for(sdf, c._internal.data_columns[0])
+                                                    for c in applied])
         return DataFrame(internal)
 
 
@@ -1956,11 +1967,11 @@ def nsmallest(self, n=5):
         window = Window.partitionBy([s._scol for s in groupkeys]).orderBy(F.col(name))
         sdf = sdf.withColumn('rank', F.row_number().over(window)).filter(F.col('rank') <= n)
         internal = _InternalFrame(sdf=sdf,
-                                  data_columns=[name],
                                   index_map=([(s._internal.data_columns[0],
                                                s._internal.column_index[0])
                                               for s in self._groupkeys]
-                                             + self._kdf._internal.index_map))
+                                             + self._kdf._internal.index_map),
+                                  column_scols=[scol_for(sdf, name)])
         return _col(DataFrame(internal))
 
     # TODO: add keep parameter
@@ -2002,11 +2013,11 @@ def nlargest(self, n=5):
         window = Window.partitionBy([s._scol for s in groupkeys]).orderBy(F.col(name).desc())
         sdf = sdf.withColumn('rank', F.row_number().over(window)).filter(F.col('rank') <= n)
         internal = _InternalFrame(sdf=sdf,
-                                  data_columns=[name],
                                   index_map=([(s._internal.data_columns[0],
                                                s._internal.column_index[0])
                                               for s in self._groupkeys]
-                                             + self._kdf._internal.index_map))
+                                             + self._kdf._internal.index_map),
+                                  column_scols=[scol_for(sdf, name)])
         return _col(DataFrame(internal))
 
     # TODO: add bins, normalize parameter
@@ -2064,10 +2075,10 @@ def value_counts(self, sort=None, ascending=None, dropna=True):
                 sdf = sdf.orderBy(F.col(agg_column).desc())
 
         internal = _InternalFrame(sdf=sdf,
-                                  data_columns=[agg_column],
                                   index_map=[(SPARK_INDEX_NAME_FORMAT(i),
                                               s._internal.column_index[0])
-                                             for i, s in enumerate(groupkeys)])
+                                             for i, s in enumerate(groupkeys)],
+                                  column_scols=[scol_for(sdf, agg_column)])
         return _col(DataFrame(internal))
 
 

databricks/koalas/groupby.py

@@ -88,8 +88,8 @@ def __init__(self, data: Union[DataFrame, list], dtype=None, name=None,
         if scol is None:
             scol = kdf._internal.index_scols[0]
         internal = kdf._internal.copy(scol=scol,
-                                      data_columns=kdf._internal.index_columns,
                                       column_index=kdf._internal.index_names,
+                                      column_scols=kdf._internal.index_scols,
                                       column_index_names=None)
         IndexOpsMixin.__init__(self, internal, kdf)
 
@@ -139,7 +139,7 @@ def to_pandas(self) -> pd.Index:
         internal = self._kdf._internal.copy(
             sdf=sdf,
             index_map=[(sdf.schema[0].name, self._kdf._internal.index_names[0])],
-            data_columns=[], column_index=[], column_index_names=None)
+            column_index=[], column_scols=[], column_index_names=None)
         return DataFrame(internal)._to_internal_pandas().index
 
     toPandas = to_pandas

databricks/koalas/indexes.py

@@ -28,6 +28,7 @@
 
 from databricks.koalas.internal import _InternalFrame
 from databricks.koalas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError
+from databricks.koalas.utils import scol_for
 
 
 def _make_col(c):
@@ -437,7 +438,7 @@ def raiseNotImplemented(description):
             cols_sel = None
 
         if cols_sel is None:
-            columns = self._kdf._internal.data_scols
+            columns = self._kdf._internal.column_scols
         elif isinstance(cols_sel, spark.Column):
             columns = [cols_sel]
             column_index = None
@@ -475,9 +476,9 @@ def raiseNotImplemented(description):
             sdf = sdf.select(self._kdf._internal.index_scols + columns)
             index_columns = self._kdf._internal.index_columns
             data_columns = [column for column in sdf.columns if column not in index_columns]
-            internal = _InternalFrame(
-                sdf=sdf, data_columns=data_columns,
-                index_map=self._kdf._internal.index_map, column_index=column_index)
+            column_scols = [scol_for(sdf, col) for col in data_columns]
+            internal = _InternalFrame(sdf=sdf, index_map=self._kdf._internal.index_map,
+                                      column_index=column_index, column_scols=column_scols)
             kdf = DataFrame(internal)
         except AnalysisException:
             raise KeyError('[{}] don\'t exist in columns'
@@ -710,13 +711,13 @@ def raiseNotImplemented(description):
         if isinstance(cols_sel, Series):
             columns = [cols_sel._scol]
         elif isinstance(cols_sel, int):
-            columns = [self._kdf._internal.data_scols[cols_sel]]
+            columns = [self._kdf._internal.column_scols[cols_sel]]
         elif cols_sel is None or cols_sel == slice(None):
-            columns = self._kdf._internal.data_scols
+            columns = self._kdf._internal.column_scols
         elif isinstance(cols_sel, slice):
             if all(s is None or isinstance(s, int)
                    for s in (cols_sel.start, cols_sel.stop, cols_sel.step)):
-                columns = self._kdf._internal.data_scols[cols_sel]
+                columns = self._kdf._internal.column_scols[cols_sel]
             else:
                 not_none = cols_sel.start if cols_sel.start is not None \
                     else cols_sel.stop if cols_sel.stop is not None else cols_sel.step
@@ -733,10 +734,7 @@ def raiseNotImplemented(description):
 
         try:
             sdf = sdf.select(self._kdf._internal.index_scols + columns)
-            index_columns = self._kdf._internal.index_columns
-            data_columns = [column for column in sdf.columns if column not in index_columns]
-            internal = _InternalFrame(
-                sdf=sdf, data_columns=data_columns, index_map=self._kdf._internal.index_map)
+            internal = _InternalFrame(sdf=sdf, index_map=self._kdf._internal.index_map)
             kdf = DataFrame(internal)
         except AnalysisException:
             raise KeyError('[{}] don\'t exist in columns'