You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2024/06/04 00:24:29 UTC
(spark) branch master updated: [SPARK-48482][PYTHON] dropDuplicates and dropDuplicatesWIthinWatermark should accept variable length args
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 560c08332b35 [SPARK-48482][PYTHON] dropDuplicates and dropDuplicatesWIthinWatermark should accept variable length args
560c08332b35 is described below
commit 560c08332b35941260169124b4f522bdc82b84d8
Author: Wei Liu <we...@databricks.com>
AuthorDate: Tue Jun 4 09:24:17 2024 +0900
[SPARK-48482][PYTHON] dropDuplicates and dropDuplicatesWIthinWatermark should accept variable length args
### What changes were proposed in this pull request?
In scala, `dropDuplicates` and `dropDuplicatesWIthinWatermark` accepts varargs, i.e. `df.dropDuplicates("id", "value")`. However this is not supported in Python, users have to wrap them with list. This PR fixes it.
### Why are the changes needed?
Better API, integrated scala and python experience
### Does this PR introduce _any_ user-facing change?
Yes, now users can use var args as parameters to `dropDuplicates` and `dropDuplicatesWithinWatermark`
### How was this patch tested?
Added & modified existing unit tests
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #46817 from WweiL/dropDuplicates-accept-vararg.
Authored-by: Wei Liu <we...@databricks.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
docs/structured-streaming-programming-guide.md | 6 +--
python/pyspark/sql/classic/dataframe.py | 39 +++++++++--------
python/pyspark/sql/connect/dataframe.py | 49 ++++++++++++++--------
python/pyspark/sql/connect/plan.py | 2 +-
python/pyspark/sql/dataframe.py | 22 +++++++---
.../sql/tests/connect/test_connect_basic.py | 12 ++++++
.../pyspark/sql/tests/connect/test_connect_plan.py | 20 +++++++--
python/pyspark/sql/tests/test_dataframe.py | 16 ++++---
8 files changed, 109 insertions(+), 57 deletions(-)
diff --git a/docs/structured-streaming-programming-guide.md b/docs/structured-streaming-programming-guide.md
index 4c3eca6b6d55..fabe7f17b78b 100644
--- a/docs/structured-streaming-programming-guide.md
+++ b/docs/structured-streaming-programming-guide.md
@@ -2082,12 +2082,12 @@ You can deduplicate records in data streams using a unique identifier in the eve
streamingDf = spark.readStream. ...
# Without watermark using guid column
-streamingDf.dropDuplicates(["guid"])
+streamingDf.dropDuplicates("guid")
# With watermark using guid and eventTime columns
streamingDf \
.withWatermark("eventTime", "10 seconds") \
- .dropDuplicates(["guid", "eventTime"])
+ .dropDuplicates("guid", "eventTime")
{% endhighlight %}
</div>
@@ -2163,7 +2163,7 @@ streamingDf = spark.readStream. ...
# deduplicate using guid column with watermark based on eventTime column
streamingDf \
.withWatermark("eventTime", "10 hours") \
- .dropDuplicatesWithinWatermark(["guid"])
+ .dropDuplicatesWithinWatermark("guid")
{% endhighlight %}
</div>
diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py
index 9b6790d29aaa..a03467aff194 100644
--- a/python/pyspark/sql/classic/dataframe.py
+++ b/python/pyspark/sql/classic/dataframe.py
@@ -1220,28 +1220,34 @@ class DataFrame(ParentDataFrame, PandasMapOpsMixin, PandasConversionMixin):
def subtract(self, other: ParentDataFrame) -> ParentDataFrame:
return DataFrame(getattr(self._jdf, "except")(other._jdf), self.sparkSession)
- def dropDuplicates(self, subset: Optional[List[str]] = None) -> ParentDataFrame:
- if subset is not None and (not isinstance(subset, Iterable) or isinstance(subset, str)):
- raise PySparkTypeError(
- error_class="NOT_LIST_OR_TUPLE",
- message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__},
- )
-
- if subset is None:
+ def dropDuplicates(self, *subset: Union[str, List[str]]) -> ParentDataFrame:
+ # Acceptable args should be str, ... or a single List[str]
+ # So if subset length is 1, it can be either single str, or a list of str
+ # if subset length is greater than 1, it must be a sequence of str
+ if len(subset) > 1:
+ assert all(isinstance(c, str) for c in subset)
+
+ if not subset:
jdf = self._jdf.dropDuplicates()
+ elif len(subset) == 1 and isinstance(subset[0], list):
+ jdf = self._jdf.dropDuplicates(self._jseq(subset[0]))
else:
jdf = self._jdf.dropDuplicates(self._jseq(subset))
return DataFrame(jdf, self.sparkSession)
- def dropDuplicatesWithinWatermark(self, subset: Optional[List[str]] = None) -> ParentDataFrame:
- if subset is not None and (not isinstance(subset, Iterable) or isinstance(subset, str)):
- raise PySparkTypeError(
- error_class="NOT_LIST_OR_TUPLE",
- message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__},
- )
+ drop_duplicates = dropDuplicates
- if subset is None:
+ def dropDuplicatesWithinWatermark(self, *subset: Union[str, List[str]]) -> ParentDataFrame:
+ # Acceptable args should be str, ... or a single List[str]
+ # So if subset length is 1, it can be either single str, or a list of str
+ # if subset length is greater than 1, it must be a sequence of str
+ if len(subset) > 1:
+ assert all(isinstance(c, str) for c in subset)
+
+ if not subset:
jdf = self._jdf.dropDuplicatesWithinWatermark()
+ elif len(subset) == 1 and isinstance(subset[0], list):
+ jdf = self._jdf.dropDuplicatesWithinWatermark(self._jseq(subset[0]))
else:
jdf = self._jdf.dropDuplicatesWithinWatermark(self._jseq(subset))
return DataFrame(jdf, self.sparkSession)
@@ -1787,9 +1793,6 @@ class DataFrame(ParentDataFrame, PandasMapOpsMixin, PandasConversionMixin):
def groupby(self, *cols: "ColumnOrNameOrOrdinal") -> "GroupedData": # type: ignore[misc]
return self.groupBy(*cols)
- def drop_duplicates(self, subset: Optional[List[str]] = None) -> ParentDataFrame:
- return self.dropDuplicates(subset)
-
def writeTo(self, table: str) -> DataFrameWriterV2:
return DataFrameWriterV2(self, table)
diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py
index ee60968a84ce..35503d0efbd5 100644
--- a/python/pyspark/sql/connect/dataframe.py
+++ b/python/pyspark/sql/connect/dataframe.py
@@ -408,39 +408,54 @@ class DataFrame(ParentDataFrame):
},
)
- def dropDuplicates(self, subset: Optional[List[str]] = None) -> ParentDataFrame:
- if subset is not None and not isinstance(subset, (list, tuple)):
- raise PySparkTypeError(
- error_class="NOT_LIST_OR_TUPLE",
- message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__},
- )
-
- if subset is None:
+ def dropDuplicates(self, *subset: Union[str, List[str]]) -> ParentDataFrame:
+ # Acceptable args should be str, ... or a single List[str]
+ # So if subset length is 1, it can be either single str, or a list of str
+ # if subset length is greater than 1, it must be a sequence of str
+ if len(subset) > 1:
+ assert all(isinstance(c, str) for c in subset)
+
+ if not subset:
return DataFrame(
plan.Deduplicate(child=self._plan, all_columns_as_keys=True), session=self._session
)
+ elif len(subset) == 1 and isinstance(subset[0], list):
+ return DataFrame(
+ plan.Deduplicate(child=self._plan, column_names=subset[0]),
+ session=self._session,
+ )
else:
return DataFrame(
- plan.Deduplicate(child=self._plan, column_names=subset), session=self._session
+ plan.Deduplicate(child=self._plan, column_names=cast(List[str], subset)),
+ session=self._session,
)
drop_duplicates = dropDuplicates
- def dropDuplicatesWithinWatermark(self, subset: Optional[List[str]] = None) -> ParentDataFrame:
- if subset is not None and not isinstance(subset, (list, tuple)):
- raise PySparkTypeError(
- error_class="NOT_LIST_OR_TUPLE",
- message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__},
- )
+ def dropDuplicatesWithinWatermark(self, *subset: Union[str, List[str]]) -> ParentDataFrame:
+ # Acceptable args should be str, ... or a single List[str]
+ # So if subset length is 1, it can be either single str, or a list of str
+ # if subset length is greater than 1, it must be a sequence of str
+ if len(subset) > 1:
+ assert all(isinstance(c, str) for c in subset)
- if subset is None:
+ if not subset:
return DataFrame(
plan.Deduplicate(child=self._plan, all_columns_as_keys=True, within_watermark=True),
session=self._session,
)
+ elif len(subset) == 1 and isinstance(subset[0], list):
+ return DataFrame(
+ plan.Deduplicate(child=self._plan, column_names=subset[0], within_watermark=True),
+ session=self._session,
+ )
else:
return DataFrame(
- plan.Deduplicate(child=self._plan, column_names=subset, within_watermark=True),
+ plan.Deduplicate(
+ child=self._plan,
+ column_names=cast(List[str], subset),
+ within_watermark=True,
+ ),
session=self._session,
)
diff --git a/python/pyspark/sql/connect/plan.py b/python/pyspark/sql/connect/plan.py
index 868bd4fb57aa..b5a71c4d2536 100644
--- a/python/pyspark/sql/connect/plan.py
+++ b/python/pyspark/sql/connect/plan.py
@@ -682,7 +682,7 @@ class Deduplicate(LogicalPlan):
self,
child: Optional["LogicalPlan"],
all_columns_as_keys: bool = False,
- column_names: Optional[List[str]] = None,
+ column_names: Optional[Sequence[str]] = None,
within_watermark: bool = False,
) -> None:
super().__init__(child)
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index 88cb4546fcab..62c46cfec93c 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -4568,7 +4568,7 @@ class DataFrame:
...
@dispatch_df_method
- def dropDuplicates(self, subset: Optional[List[str]] = None) -> "DataFrame":
+ def dropDuplicates(self, *subset: Union[str, List[str]]) -> "DataFrame":
"""Return a new :class:`DataFrame` with duplicate rows removed,
optionally only considering certain columns.
@@ -4585,6 +4585,9 @@ class DataFrame:
.. versionchanged:: 3.4.0
Supports Spark Connect.
+ .. versionchanged:: 4.0.0
+ Supports variable-length argument
+
Parameters
----------
subset : list of column names, optional
@@ -4616,7 +4619,7 @@ class DataFrame:
Deduplicate values on 'name' and 'height' columns.
- >>> df.dropDuplicates(['name', 'height']).show()
+ >>> df.dropDuplicates('name', 'height').show()
+-----+---+------+
| name|age|height|
+-----+---+------+
@@ -4626,7 +4629,7 @@ class DataFrame:
...
@dispatch_df_method
- def dropDuplicatesWithinWatermark(self, subset: Optional[List[str]] = None) -> "DataFrame":
+ def dropDuplicatesWithinWatermark(self, *subset: Union[str, List[str]]) -> "DataFrame":
"""Return a new :class:`DataFrame` with duplicate rows removed,
optionally only considering certain columns, within watermark.
@@ -4643,6 +4646,9 @@ class DataFrame:
.. versionadded:: 3.5.0
+ .. versionchanged:: 4.0.0
+ Supports variable-length argument
+
Parameters
----------
subset : List of column names, optional
@@ -4672,7 +4678,7 @@ class DataFrame:
Deduplicate values on 'value' columns.
- >>> df.dropDuplicatesWithinWatermark(['value']) # doctest: +SKIP
+ >>> df.dropDuplicatesWithinWatermark('value') # doctest: +SKIP
"""
...
@@ -5929,11 +5935,17 @@ class DataFrame:
...
@dispatch_df_method
- def drop_duplicates(self, subset: Optional[List[str]] = None) -> "DataFrame":
+ def drop_duplicates(self, *subset: Union[str, List[str]]) -> "DataFrame":
"""
:func:`drop_duplicates` is an alias for :func:`dropDuplicates`.
.. versionadded:: 1.4.0
+
+ .. versionchanged:: 3.4.0
+ Supports Spark Connect
+
+ .. versionchanged:: 4.0.0
+ Supports variable-length argument
"""
...
diff --git a/python/pyspark/sql/tests/connect/test_connect_basic.py b/python/pyspark/sql/tests/connect/test_connect_basic.py
index eb5cb18d11e6..74d08424cafc 100755
--- a/python/pyspark/sql/tests/connect/test_connect_basic.py
+++ b/python/pyspark/sql/tests/connect/test_connect_basic.py
@@ -657,6 +657,18 @@ class SparkConnectBasicTests(SparkConnectSQLTestCase):
self.assert_eq(
df.dropDuplicates(["name"]).toPandas(), df2.dropDuplicates(["name"]).toPandas()
)
+ self.assert_eq(
+ df.drop_duplicates(["name"]).toPandas(), df2.drop_duplicates(["name"]).toPandas()
+ )
+ self.assert_eq(
+ df.dropDuplicates(["name", "id"]).toPandas(),
+ df2.dropDuplicates(["name", "id"]).toPandas(),
+ )
+ self.assert_eq(
+ df.drop_duplicates(["name", "id"]).toPandas(),
+ df2.drop_duplicates(["name", "id"]).toPandas(),
+ )
+ self.assert_eq(df.dropDuplicates("name").toPandas(), df2.dropDuplicates("name").toPandas())
def test_drop(self):
# SPARK-41169: test drop
diff --git a/python/pyspark/sql/tests/connect/test_connect_plan.py b/python/pyspark/sql/tests/connect/test_connect_plan.py
index e8d04aeada74..47e3fb5a9602 100644
--- a/python/pyspark/sql/tests/connect/test_connect_plan.py
+++ b/python/pyspark/sql/tests/connect/test_connect_plan.py
@@ -553,13 +553,25 @@ class SparkConnectPlanTests(PlanOnlyTestFixture):
self.assertEqual(deduplicate_on_all_columns_plan.root.deduplicate.all_columns_as_keys, True)
self.assertEqual(len(deduplicate_on_all_columns_plan.root.deduplicate.column_names), 0)
- deduplicate_on_subset_columns_plan = df.dropDuplicates(["name", "height"])._plan.to_proto(
- self.connect
+ deduplicate_on_subset_columns_plan_list_arg = df.dropDuplicates(
+ ["name", "height"]
+ )._plan.to_proto(self.connect)
+ self.assertEqual(
+ deduplicate_on_subset_columns_plan_list_arg.root.deduplicate.all_columns_as_keys, False
+ )
+ self.assertEqual(
+ len(deduplicate_on_subset_columns_plan_list_arg.root.deduplicate.column_names), 2
+ )
+
+ deduplicate_on_subset_columns_plan_var_arg = df.dropDuplicates(
+ "name", "height"
+ )._plan.to_proto(self.connect)
+ self.assertEqual(
+ deduplicate_on_subset_columns_plan_var_arg.root.deduplicate.all_columns_as_keys, False
)
self.assertEqual(
- deduplicate_on_subset_columns_plan.root.deduplicate.all_columns_as_keys, False
+ len(deduplicate_on_subset_columns_plan_var_arg.root.deduplicate.column_names), 2
)
- self.assertEqual(len(deduplicate_on_subset_columns_plan.root.deduplicate.column_names), 2)
def test_relation_alias(self):
df = self.connect.readTable(table_name=self.tbl_name)
diff --git a/python/pyspark/sql/tests/test_dataframe.py b/python/pyspark/sql/tests/test_dataframe.py
index ac46da587495..36a856b62719 100644
--- a/python/pyspark/sql/tests/test_dataframe.py
+++ b/python/pyspark/sql/tests/test_dataframe.py
@@ -252,24 +252,22 @@ class DataFrameTestsMixin:
self.assertEqual(df2.columns, ["a"])
def test_drop_duplicates(self):
- # SPARK-36034 test that drop duplicates throws a type error when in correct type provided
df = self.spark.createDataFrame([("Alice", 50), ("Alice", 60)], ["name", "age"])
# shouldn't drop a non-null row
self.assertEqual(df.dropDuplicates().count(), 2)
self.assertEqual(df.dropDuplicates(["name"]).count(), 1)
-
self.assertEqual(df.dropDuplicates(["name", "age"]).count(), 2)
- with self.assertRaises(PySparkTypeError) as pe:
- df.dropDuplicates("name")
+ self.assertEqual(df.drop_duplicates(["name"]).count(), 1)
+ self.assertEqual(df.drop_duplicates(["name", "age"]).count(), 2)
- self.check_error(
- exception=pe.exception,
- error_class="NOT_LIST_OR_TUPLE",
- message_parameters={"arg_name": "subset", "arg_type": "str"},
- )
+ # SPARK-48482 dropDuplicates should take varargs
+ self.assertEqual(df.dropDuplicates("name").count(), 1)
+ self.assertEqual(df.dropDuplicates("name", "age").count(), 2)
+ self.assertEqual(df.drop_duplicates("name").count(), 1)
+ self.assertEqual(df.drop_duplicates("name", "age").count(), 2)
def test_drop_duplicates_with_ambiguous_reference(self):
df1 = self.spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"])
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org