Skip to content

Commit 1fe8eb7

Browse files
itholicHyukjinKwon
authored andcommitted
Impelment DataFrame.keys (#937)
An alias for columns. ```python >>> 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') >>> df.to_pandas().keys() Index(['max_speed', 'shield'], dtype='object') ```
1 parent 465861c commit 1fe8eb7

File tree

4 files changed

+34
-1
lines changed

4 files changed

+34
-1
lines changed

databricks/koalas/frame.py

+25
Original file line numberDiff line numberDiff line change
@@ -7405,6 +7405,31 @@ def gen_new_column_index_entry(column_index_entry):
74057405
else:
74067406
return DataFrame(internal)
74077407

7408+
def keys(self):
7409+
"""
7410+
Return alias for columns.
7411+
7412+
Returns
7413+
-------
7414+
Index
7415+
Columns of the DataFrame.
7416+
7417+
Examples
7418+
--------
7419+
>>> df = ks.DataFrame([[1, 2], [4, 5], [7, 8]],
7420+
... index=['cobra', 'viper', 'sidewinder'],
7421+
... columns=['max_speed', 'shield'])
7422+
>>> df
7423+
max_speed shield
7424+
cobra 1 2
7425+
viper 4 5
7426+
sidewinder 7 8
7427+
7428+
>>> df.keys()
7429+
Index(['max_speed', 'shield'], dtype='object')
7430+
"""
7431+
return self.columns
7432+
74087433
def _get_from_multiindex_column(self, key):
74097434
""" Select columns from multi-index columns.
74107435

databricks/koalas/missing/frame.py

-1
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,6 @@ class _MissingPandasLikeDataFrame(object):
6666
interpolate = unsupported_function('interpolate')
6767
iterrows = unsupported_function('iterrows')
6868
itertuples = unsupported_function('itertuples')
69-
keys = unsupported_function('keys')
7069
last = unsupported_function('last')
7170
last_valid_index = unsupported_function('last_valid_index')
7271
lookup = unsupported_function('lookup')

databricks/koalas/tests/test_dataframe.py

+8
Original file line numberDiff line numberDiff line change
@@ -2161,3 +2161,11 @@ def test_to_spark(self):
21612161

21622162
with self.assertRaisesRegex(ValueError, "length of index columns.*1.*3"):
21632163
kdf.to_spark(index_col=["x", "y", "z"])
2164+
2165+
def test_keys(self):
2166+
kdf = ks.DataFrame([[1, 2], [4, 5], [7, 8]],
2167+
index=['cobra', 'viper', 'sidewinder'],
2168+
columns=['max_speed', 'shield'])
2169+
pdf = kdf.to_pandas()
2170+
2171+
self.assert_eq(kdf.keys(), pdf.keys())

docs/source/reference/frame.rst

+1
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ Indexing, iteration
5454
DataFrame.iloc
5555
DataFrame.items
5656
DataFrame.iteritems
57+
DataFrame.keys
5758
DataFrame.xs
5859
DataFrame.get
5960

0 commit comments

Comments
 (0)