Impelment DataFrame.keys

databricks/koalas/frame.py

@@ -7405,6 +7405,31 @@ def gen_new_column_index_entry(column_index_entry):
         else:
             return DataFrame(internal)
 
+    def keys(self):
+        """
+        Return alias for columns.
+
+        Returns
+        -------
+        Index
+            Columns of the DataFrame.
+
+        Examples
+        --------
+        >>> df = ks.DataFrame([[1, 2], [4, 5], [7, 8]],
+        ...                   index=['cobra', 'viper', 'sidewinder'],
+        ...                   columns=['max_speed', 'shield'])
+        >>> df
+                    max_speed  shield
+        cobra               1       2
+        viper               4       5
+        sidewinder          7       8
+
+        >>> df.keys()
+        Index(['max_speed', 'shield'], dtype='object')
+        """
+        return self.columns
+
     def _get_from_multiindex_column(self, key):
         """ Select columns from multi-index columns.
 

databricks/koalas/missing/frame.py

@@ -66,7 +66,6 @@ class _MissingPandasLikeDataFrame(object):
     interpolate = unsupported_function('interpolate')
     iterrows = unsupported_function('iterrows')
     itertuples = unsupported_function('itertuples')
-    keys = unsupported_function('keys')
     last = unsupported_function('last')
     last_valid_index = unsupported_function('last_valid_index')
     lookup = unsupported_function('lookup')

databricks/koalas/tests/test_dataframe.py

@@ -2161,3 +2161,11 @@ def test_to_spark(self):
 
         with self.assertRaisesRegex(ValueError, "length of index columns.*1.*3"):
             kdf.to_spark(index_col=["x", "y", "z"])
+
+    def test_keys(self):
+        kdf = ks.DataFrame([[1, 2], [4, 5], [7, 8]],
+                           index=['cobra', 'viper', 'sidewinder'],
+                           columns=['max_speed', 'shield'])
+        pdf = kdf.to_pandas()
+
+        self.assert_eq(kdf.keys(), pdf.keys())

docs/source/reference/frame.rst

@@ -54,6 +54,7 @@ Indexing, iteration
    DataFrame.iloc
    DataFrame.items
    DataFrame.iteritems
+   DataFrame.keys
    DataFrame.xs
    DataFrame.get