Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix DataFrame.drop() to remove fields from Spark DataFrame also. #794

Merged
merged 11 commits into from
Sep 20, 2019
Merged

Fix DataFrame.drop() to remove fields from Spark DataFrame also. #794

Changes from 9 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
7c9538e
Drop internal spark dataframe columns also
itholic Sep 18, 2019
811aed0
Fix to not modifying internal
itholic Sep 18, 2019
43d3fbe
Support for MultiIndex column
itholic Sep 18, 2019
bca770d
Add doctest for multiindex
itholic Sep 18, 2019
fed4d46
Add doctest NORMALIZE_WHITESPACE
itholic Sep 18, 2019
3f304ff
Remove trailing whitespace
itholic Sep 18, 2019
99f16ec
Add NORMALIZE
itholic Sep 18, 2019
a50b104
Remove multi-index test since lint failed
itholic Sep 18, 2019
4792fed
Fix & Readd doctest for python 3.5 compaitability
itholic Sep 19, 2019
382987e
Fix doctest
itholic Sep 19, 2019
e23f9ee
Remove pandas part from doctest
itholic Sep 19, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions databricks/koalas/frame.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4344,6 +4344,20 @@ def drop(self, labels=None, axis=1,
0 1 7
1 2 8

>>> df = ks.DataFrame({('a', 'x'): [1, 2], ('a', 'y'): [3, 4],
... ('b', 'z'): [5, 6], ('b', 'w'): [7, 8]},
... columns=[('a', 'x'), ('a', 'y'), ('b', 'z'), ('b', 'w')])
>>> df # doctest: +NORMALIZE_WHITESPACE
a b
x y z w
0 1 3 5 7
1 2 4 6 8
>>> df.drop('a') # doctest: +NORMALIZE_WHITESPACE
b
z w
0 5 7
1 6 8

Notes
-----
Currently only axis = 1 is supported in this function,
Expand All @@ -4367,11 +4381,15 @@ def drop(self, labels=None, axis=1,
if idx[:len(col)] == col)
if len(drop_column_index) == 0:
raise KeyError(columns)
cols, idx = zip(*((column, idx)
cols, idxes = zip(*((column, idx)
for column, idx
in zip(self._internal.data_columns, self._internal.column_index)
if idx not in drop_column_index))
internal = self._internal.copy(data_columns=list(cols), column_index=list(idx))
internal = self._internal.copy(
sdf=self._sdf.select(
self._internal.index_scols + [self._internal.scol_for(idx) for idx in idxes]),
data_columns=list(cols),
column_index=list(idxes))
return DataFrame(internal)
else:
raise ValueError("Need to specify at least one of 'labels' or 'columns'")
Expand Down