This is an automated email from the ASF dual-hosted git repository. gurwls223 pushed a commit to branch branch-3.2 in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-3.2 by this push: new 88bba0c [SPARK-36697][PYTHON] Fix dropping all columns of a DataFrame 88bba0c is described below commit 88bba0c94be6946add41c5ea8510b912133205e9 Author: Xinrong Meng <xinrong.m...@databricks.com> AuthorDate: Thu Sep 9 09:59:42 2021 +0900 [SPARK-36697][PYTHON] Fix dropping all columns of a DataFrame ### What changes were proposed in this pull request? Fix dropping all columns of a DataFrame ### Why are the changes needed? When dropping all columns of a pandas-on-Spark DataFrame, a ValueError is raised. Whereas in pandas, an empty DataFrame reserving the index is returned. We should follow pandas. ### Does this PR introduce _any_ user-facing change? Yes. From ```py >>> psdf = ps.DataFrame({"x": [1, 2], "y": [3, 4], "z": [5, 6]}) >>> psdf x y z 0 1 3 5 1 2 4 6 >>> psdf.drop(['x', 'y', 'z']) Traceback (most recent call last): ... ValueError: not enough values to unpack (expected 2, got 0) ``` To ```py >>> psdf = ps.DataFrame({"x": [1, 2], "y": [3, 4], "z": [5, 6]}) >>> psdf x y z 0 1 3 5 1 2 4 6 >>> psdf.drop(['x', 'y', 'z']) Empty DataFrame Columns: [] Index: [0, 1] ``` ### How was this patch tested? Unit tests. Closes #33938 from xinrong-databricks/frame_drop_col. Authored-by: Xinrong Meng <xinrong.m...@databricks.com> Signed-off-by: Hyukjin Kwon <gurwls...@apache.org> (cherry picked from commit 33bb7b39e927ee8c357fd4f8c073be6f3a1d2fb0) Signed-off-by: Hyukjin Kwon <gurwls...@apache.org> --- python/pyspark/pandas/frame.py | 17 ++++++++++------- python/pyspark/pandas/tests/test_dataframe.py | 13 +++++++++++++ 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/python/pyspark/pandas/frame.py b/python/pyspark/pandas/frame.py index 48be138..ec6b261 100644 --- a/python/pyspark/pandas/frame.py +++ b/python/pyspark/pandas/frame.py @@ -6726,14 +6726,17 @@ defaultdict(<class 'list'>, {'col..., 'col...})] ) if len(drop_column_labels) == 0: raise KeyError(columns) - cols, labels = zip( - *( - (column, label) - for column, label in zip( - self._internal.data_spark_column_names, self._internal.column_labels - ) - if label not in drop_column_labels + + keep_columns_and_labels = [ + (column, label) + for column, label in zip( + self._internal.data_spark_column_names, self._internal.column_labels ) + if label not in drop_column_labels + ] + + cols, labels = ( + zip(*keep_columns_and_labels) if len(keep_columns_and_labels) > 0 else ([], []) ) internal = self._internal.with_new_columns([self._psser_for(label) for label in labels]) return DataFrame(internal) diff --git a/python/pyspark/pandas/tests/test_dataframe.py b/python/pyspark/pandas/tests/test_dataframe.py index 5fe1165..11da18c 100644 --- a/python/pyspark/pandas/tests/test_dataframe.py +++ b/python/pyspark/pandas/tests/test_dataframe.py @@ -1263,9 +1263,11 @@ class DataFrameTest(PandasOnSparkTestCase, SQLTestUtils): self.assert_eq(psdf.drop("x"), pdf.drop("x", axis=1)) # Assert using a list for 'labels' works self.assert_eq(psdf.drop(["y", "z"], axis=1), pdf.drop(["y", "z"], axis=1)) + self.assert_eq(psdf.drop(["x", "y", "z"], axis=1), pdf.drop(["x", "y", "z"], axis=1)) # Assert using 'columns' instead of 'labels' produces the same results self.assert_eq(psdf.drop(columns="x"), pdf.drop(columns="x")) self.assert_eq(psdf.drop(columns=["y", "z"]), pdf.drop(columns=["y", "z"])) + self.assert_eq(psdf.drop(columns=["x", "y", "z"]), pdf.drop(columns=["x", "y", "z"])) # Assert 'labels' being used when both 'labels' and 'columns' are specified # TODO: should throw an error? @@ -1279,10 +1281,21 @@ class DataFrameTest(PandasOnSparkTestCase, SQLTestUtils): self.assert_eq(psdf.drop(columns=1), pdf.drop(columns=1)) self.assert_eq(psdf.drop(columns=(1, "x")), pdf.drop(columns=(1, "x"))) self.assert_eq(psdf.drop(columns=[(1, "x"), 2]), pdf.drop(columns=[(1, "x"), 2])) + self.assert_eq( + psdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), + pdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), + ) self.assertRaises(KeyError, lambda: psdf.drop(columns=3)) self.assertRaises(KeyError, lambda: psdf.drop(columns=(1, "z"))) + pdf.index = pd.MultiIndex.from_tuples([("i", 0), ("j", 1)]) + psdf = ps.from_pandas(pdf) + self.assert_eq( + psdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), + pdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), + ) + # non-string names pdf = pd.DataFrame({10: [1, 2], 20: [3, 4], 30: [5, 6]}, index=np.random.rand(2)) psdf = ps.from_pandas(pdf) --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org For additional commands, e-mail: commits-h...@spark.apache.org