Exclude Index columns for exposed Spark DataFrame and disallow Koalas DataFrame with no index

Author: HyukjinKwon

databricks/koalas/frame.py

@@ -1598,9 +1598,7 @@ def index(self):
         Index
         """
         from databricks.koalas.indexes import Index, MultiIndex
-        if len(self._internal.index_map) == 0:
-            return None
-        elif len(self._internal.index_map) == 1:
+        if len(self._internal.index_map) == 1:
             return Index(self)
         else:
             return MultiIndex(self)
@@ -1860,9 +1858,6 @@ class    max    type
         lion           mammal   80.5     run
         monkey         mammal    NaN    jump
         """
-        if len(self._internal.index_map) == 0:
-            raise NotImplementedError('Can\'t reset index because there is no index.')
-
         multi_index = len(self._internal.index_map) > 1
 
         def rename(index):
@@ -1877,7 +1872,10 @@ def rename(index):
         if level is None:
             new_index_map = [(column, name if name is not None else rename(i))
                              for i, (column, name) in enumerate(self._internal.index_map)]
-            index_map = []
+            # TODO: this will end up with multiple columns that point the same index column.
+            #     For instance, if we do df.reset_index(), the new column `index` points the
+            #     same '__index_level_0__' internally. We should have new column.
+            index_map = [('__index_level_0__', None)]
         else:
             if isinstance(level, (int, str)):
                 level = [level]
@@ -2382,13 +2380,13 @@ def to_koalas(self):
 
         >>> spark_df = df.to_spark()
         >>> spark_df
-        DataFrame[__index_level_0__: bigint, col1: bigint, col2: bigint]
+        DataFrame[col1: bigint, col2: bigint]
 
         >>> kdf = spark_df.to_koalas()
         >>> kdf
-           __index_level_0__  col1  col2
-        0                  0     1     3
-        1                  1     2     4
+           col1  col2
+        0     1     3
+        1     2     4
 
         Calling to_koalas on a Koalas DataFrame simply returns itself.
 
@@ -2493,8 +2491,8 @@ def to_table(self, name: str, format: Optional[str] = None, mode: str = 'error',
 
         >>> df.to_table('%s.my_table' % db, partition_cols='date')
         """
-        self._sdf.write.saveAsTable(name=name, format=format, mode=mode,
-                                    partitionBy=partition_cols, options=options)
+        self.to_spark().write.saveAsTable(name=name, format=format, mode=mode,
+                                          partitionBy=partition_cols, options=options)
 
     def to_delta(self, path: str, mode: str = 'error',
                  partition_cols: Union[str, List[str], None] = None, **options):
@@ -2604,8 +2602,8 @@ def to_parquet(self, path: str, mode: str = 'error',
         ...     mode = 'overwrite',
         ...     partition_cols=['date', 'country'])
         """
-        self._sdf.write.parquet(path=path, mode=mode, partitionBy=partition_cols,
-                                compression=compression)
+        self.to_spark().write.parquet(
+            path=path, mode=mode, partitionBy=partition_cols, compression=compression)
 
     def to_spark_io(self, path: Optional[str] = None, format: Optional[str] = None,
                     mode: str = 'error', partition_cols: Union[str, List[str], None] = None,
@@ -2657,13 +2655,16 @@ def to_spark_io(self, path: Optional[str] = None, format: Optional[str] = None,
 
         >>> df.to_spark_io(path='%s/to_spark_io/foo.json' % path, format='json')
         """
-        self._sdf.write.save(path=path, format=format, mode=mode, partitionBy=partition_cols,
-                             options=options)
+        self.to_spark().write.save(
+            path=path, format=format, mode=mode, partitionBy=partition_cols, options=options)
 
     def to_spark(self):
         """
         Return the current DataFrame as a Spark DataFrame.
 
+        .. note:: Index information is lost. So, if the index columns are not present in
+            actual columns, they are lost.
+
         See Also
         --------
         DataFrame.to_koalas
@@ -3653,14 +3654,21 @@ def pivot_table(self, values=None, index=None, columns=None,
             sdf = sdf.fillna(fill_value)
 
         if index is not None:
-            return DataFrame(sdf).set_index(index)
+            data_columns = [column for column in sdf.columns if column not in index]
+            index_map = [(column, column) for column in index]
+            internal = _InternalFrame(sdf=sdf, data_columns=data_columns, index_map=index_map)
+            return DataFrame(internal)
         else:
             if isinstance(values, list):
                 index_values = values[-1]
             else:
                 index_values = values
 
-            return DataFrame(sdf.withColumn(columns, F.lit(index_values))).set_index(columns)
+            sdf = sdf.withColumn(columns, F.lit(index_values))
+            data_columns = [column for column in sdf.columns if column not in columns]
+            index_map = [(column, column) for column in columns]
+            internal = _InternalFrame(sdf=sdf, data_columns=data_columns, index_map=index_map)
+            return DataFrame(internal)
 
     def pivot(self, index=None, columns=None, values=None):
         """
@@ -4364,9 +4372,6 @@ def sort_index(self, axis: int = 0,
         a 1  2  1
         b 1  0  3
         """
-        if len(self._internal.index_map) == 0:
-            raise ValueError("Index should be set.")
-
         if axis != 0:
             raise ValueError("No other axes than 0 are supported at the moment")
         if kind is not None:

databricks/koalas/indexing.py

@@ -26,6 +26,7 @@
 from pyspark.sql.types import BooleanType
 from pyspark.sql.utils import AnalysisException
 
+from databricks.koalas.internal import _InternalFrame
 from databricks.koalas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError
 from databricks.koalas.utils import column_index_level
 
@@ -450,15 +451,17 @@ def raiseNotImplemented(description):
                     raise KeyError("['{}'] not in index".format(key))
 
         try:
-            kdf = DataFrame(sdf.select(self._kdf._internal.index_scols + columns))
+            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)
+            kdf = DataFrame(internal)
         except AnalysisException:
             raise KeyError('[{}] don\'t exist in columns'
                            .format([col._jc.toString() for col in columns]))
 
-        kdf._internal = kdf._internal.copy(
-            data_columns=kdf._internal.data_columns[-len(columns):],
-            index_map=self._kdf._internal.index_map,
-            column_index=column_index)
         if cols_sel is not None and isinstance(cols_sel, spark.Column):
             from databricks.koalas.series import _col
             return _col(kdf)
@@ -686,7 +689,12 @@ def raiseNotImplemented(description):
                              "listlike of integers, boolean array] types, got {}".format(cols_sel))
 
         try:
-            kdf = DataFrame(sdf.select(self._kdf._internal.index_scols + columns))
+            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)
+            kdf = DataFrame(internal)
         except AnalysisException:
             raise KeyError('[{}] don\'t exist in columns'
                            .format([col._jc.toString() for col in columns]))
@@ -699,10 +707,7 @@ def raiseNotImplemented(description):
                 column_index = \
                     pd.MultiIndex.from_tuples(self._kdf._internal.column_index)[cols_sel].tolist()
 
-        kdf._internal = kdf._internal.copy(
-            data_columns=kdf._internal.data_columns[-len(columns):],
-            index_map=self._kdf._internal.index_map,
-            column_index=column_index)
+        kdf = DataFrame(kdf._internal.copy(column_index=column_index))
         if cols_sel is not None and isinstance(cols_sel, (Series, int)):
             from databricks.koalas.series import _col
             return _col(kdf)

databricks/koalas/internal.py

@@ -73,7 +73,7 @@ class _InternalFrame(object):
     However, all columns including index column are also stored in Spark DataFrame internally
     as below.
 
-    >>> kdf.to_spark().show()  # doctest: +NORMALIZE_WHITESPACE
+    >>> kdf._internal.spark_internal_df.show()  # doctest: +NORMALIZE_WHITESPACE
     +-----------------+---+---+---+---+---+
     |__index_level_0__|  A|  B|  C|  D|  E|
     +-----------------+---+---+---+---+---+
@@ -122,7 +122,7 @@ class _InternalFrame(object):
     [None]
     >>> internal.index_map
     [('__index_level_0__', None)]
-    >>> internal.spark_df.show()  # doctest: +NORMALIZE_WHITESPACE
+    >>> internal.spark_internal_df.show()  # doctest: +NORMALIZE_WHITESPACE
     +-----------------+---+---+---+---+---+
     |__index_level_0__|  A|  B|  C|  D|  E|
     +-----------------+---+---+---+---+---+
@@ -149,7 +149,7 @@ class _InternalFrame(object):
     3  7  11  15  19
     4  8  12  16  20
 
-    >>> kdf1.to_spark().show()  # doctest: +NORMALIZE_WHITESPACE
+    >>> kdf1._internal.spark_internal_df.show()  # doctest: +NORMALIZE_WHITESPACE
     +---+---+---+---+---+
     |  A|  B|  C|  D|  E|
     +---+---+---+---+---+
@@ -179,7 +179,7 @@ class _InternalFrame(object):
     ['A']
     >>> internal.index_map
     [('A', 'A')]
-    >>> internal.spark_df.show()  # doctest: +NORMALIZE_WHITESPACE
+    >>> internal.spark_internal_df.show()  # doctest: +NORMALIZE_WHITESPACE
     +---+---+---+---+---+
     |  A|  B|  C|  D|  E|
     +---+---+---+---+---+
@@ -207,7 +207,7 @@ class _InternalFrame(object):
     2 3  7  11  15  19
     3 4  8  12  16  20
 
-    >>> kdf2.to_spark().show()  # doctest: +NORMALIZE_WHITESPACE
+    >>> kdf2._internal.spark_internal_df.show()  # doctest: +NORMALIZE_WHITESPACE
     +-----------------+---+---+---+---+---+
     |__index_level_0__|  A|  B|  C|  D|  E|
     +-----------------+---+---+---+---+---+
@@ -237,7 +237,7 @@ class _InternalFrame(object):
     [None, 'A']
     >>> internal.index_map
     [('__index_level_0__', None), ('A', 'A')]
-    >>> internal.spark_df.show()  # doctest: +NORMALIZE_WHITESPACE
+    >>> internal.spark_internal_df.show()  # doctest: +NORMALIZE_WHITESPACE
     +-----------------+---+---+---+---+---+
     |__index_level_0__|  A|  B|  C|  D|  E|
     +-----------------+---+---+---+---+---+
@@ -322,7 +322,7 @@ class _InternalFrame(object):
     ['A']
     >>> internal.index_map
     [('A', 'A')]
-    >>> internal.spark_df.show()  # doctest: +NORMALIZE_WHITESPACE
+    >>> internal.spark_internal_df.show()  # doctest: +NORMALIZE_WHITESPACE
     +---+---+
     |  A|  B|
     +---+---+
@@ -367,21 +367,20 @@ def __init__(self, sdf: spark.DataFrame,
             # Here is when Koalas DataFrame is created directly from Spark DataFrame.
             assert column_index is None
             assert column_index_names is None
+            assert "__index_level_0__" not in sdf.schema.names
+            # Create default index.
+            index_map = [('__index_level_0__', None)]
+            sdf = _InternalFrame.attach_default_index(sdf)
 
-            if "__index_level_0__" not in sdf.schema.names:
-                # Create default index.
-                index_map = [('__index_level_0__', None)]
-                sdf = _InternalFrame.attach_default_index(sdf)
-
-        assert index_map is None \
-            or all(isinstance(index_field, str)
+        assert index_map is not None
+        assert all(isinstance(index_field, str)
                    and (index_name is None or isinstance(index_name, str))
                    for index_field, index_name in index_map)
         assert scol is None or isinstance(scol, spark.Column)
         assert data_columns is None or all(isinstance(col, str) for col in data_columns)
 
         self._sdf = sdf  # type: spark.DataFrame
-        self._index_map = (index_map if index_map is not None else [])  # type: List[IndexMap]
+        self._index_map = index_map  # type: List[IndexMap]
         self._scol = scol  # type: Optional[spark.Column]
         if scol is not None:
             self._data_columns = sdf.select(scol).columns
@@ -541,6 +540,7 @@ def scols(self) -> List[spark.Column]:
     @property
     def index_map(self) -> List[IndexMap]:
         """ Return the managed index information. """
+        assert len(self._index_map) > 0
         return self._index_map
 
     @lazy_property
@@ -568,6 +568,23 @@ def column_index_names(self) -> Optional[List[str]]:
         """ Return names of the index levels. """
         return self._column_index_names
 
+    @lazy_property
+    def spark_internal_df(self) -> spark.DataFrame:
+        """
+        Return as Spark DataFrame. This contains index columns as well
+        and should be only used for internal purposes.
+        """
+        index_columns = set(self.index_columns)
+        data_columns = []
+        for column, idx in zip(self._data_columns, self.column_index):
+            if column not in index_columns:
+                scol = self.scol_for(column)
+                name = str(idx) if len(idx) > 1 else idx[0]
+                if column != name:
+                    scol = scol.alias(name)
+                data_columns.append(scol)
+        return self._sdf.select(self.index_scols + data_columns)
+
     @lazy_property
     def spark_df(self) -> spark.DataFrame:
         """ Return as Spark DataFrame. """
@@ -580,12 +597,12 @@ def spark_df(self) -> spark.DataFrame:
                 if column != name:
                     scol = scol.alias(name)
                 data_columns.append(scol)
-        return self._sdf.select(self.index_scols + data_columns)
+        return self._sdf.select(data_columns)
 
     @lazy_property
     def pandas_df(self):
         """ Return as pandas DataFrame. """
-        sdf = self.spark_df
+        sdf = self.spark_internal_df
         pdf = sdf.toPandas()
         if len(pdf) == 0 and len(sdf.schema) > 0:
             pdf = pdf.astype({field.name: to_arrow_type(field.dataType).to_pandas_dtype()

databricks/koalas/namespace.py

@@ -273,7 +273,7 @@ def read_delta(path: str, version: Optional[str] = None, timestamp: Optional[str
     Examples
     --------
     >>> ks.range(1).to_delta('%s/read_delta/foo' % path)
-    >>> ks.read_delta('%s/read_delta/foo' % path)  # doctest: +SKIP
+    >>> ks.read_delta('%s/read_delta/foo' % path)
        id
     0   0
     """
@@ -307,7 +307,7 @@ def read_table(name: str) -> DataFrame:
     Examples
     --------
     >>> ks.range(1).to_table('%s.my_table' % db)
-    >>> ks.read_table('%s.my_table' % db)  # doctest: +SKIP
+    >>> ks.read_table('%s.my_table' % db)
        id
     0   0
     """

databricks/koalas/series.py

@@ -1000,7 +1000,7 @@ def to_frame(self, name=None) -> spark.DataFrame:
         2    c
         """
         renamed = self.rename(name)
-        sdf = renamed._internal.spark_df
+        sdf = renamed._internal.spark_internal_df
         internal = _InternalFrame(sdf=sdf,
                                   data_columns=[sdf.schema[-1].name],
                                   index_map=renamed._internal.index_map)

databricks/koalas/sql.py

@@ -92,7 +92,7 @@ def sql(query: str, globals=None, locals=None, **kwargs) -> DataFrame:
 
     >>> mydf = ks.range(10)
     >>> x = range(4)
-    >>> ks.sql("SELECT * from {mydf} WHERE id IN {x}")  # doctest: +SKIP
+    >>> ks.sql("SELECT * from {mydf} WHERE id IN {x}")
        id
     0   0
     1   1
@@ -105,9 +105,9 @@ def sql(query: str, globals=None, locals=None, **kwargs) -> DataFrame:
     ...     mydf2 = ks.DataFrame({"x": range(2)})
     ...     return ks.sql("SELECT * from {mydf2}")
     >>> statement()
-       __index_level_0__  x
-    0                  0  0
-    1                  1  1
+       x
+    0  0
+    1  1
 
     Mixing Koalas and pandas DataFrames in a join operation. Note that the index is dropped.