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/07/03 00:27:00 UTC
(spark) branch master updated: [SPARK-48714][PYTHON] Implement `DataFrame.mergeInto` in PySpark
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 6d226c1e2aba [SPARK-48714][PYTHON] Implement `DataFrame.mergeInto` in PySpark
6d226c1e2aba is described below
commit 6d226c1e2abad1df230511c6d0add295844a336b
Author: Paddy Xu <xu...@gmail.com>
AuthorDate: Wed Jul 3 09:26:48 2024 +0900
[SPARK-48714][PYTHON] Implement `DataFrame.mergeInto` in PySpark
### What changes were proposed in this pull request?
Spark 4.0 added a new `df.mergeInto` API, but it is missing from PySpark. This PR fixes that.
The support for this API in Spark Connect Python API will be added later by https://github.com/apache/spark/pull/46960.
### Why are the changes needed?
Because PySpark does not support `df.mergeInto`.
### Does this PR introduce _any_ user-facing change?
Yes, the user would be able to use the `df.mergeInto` API.
### How was this patch tested?
New unit tests.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #47086 from xupefei/pyspark-mergeinto.
Authored-by: Paddy Xu <xu...@gmail.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
python/pyspark/sql/classic/dataframe.py | 18 +-
python/pyspark/sql/connect/dataframe.py | 7 +
python/pyspark/sql/dataframe.py | 36 +++
python/pyspark/sql/merge.py | 251 +++++++++++++++++++++
python/pyspark/sql/readwriter.py | 2 +-
.../sql/tests/connect/test_parity_dataframe.py | 4 +
python/pyspark/sql/tests/test_dataframe.py | 73 ++++++
python/pyspark/sql/utils.py | 12 +-
8 files changed, 390 insertions(+), 13 deletions(-)
diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py
index 1bedd624603e..fd632c038603 100644
--- a/python/pyspark/sql/classic/dataframe.py
+++ b/python/pyspark/sql/classic/dataframe.py
@@ -57,6 +57,7 @@ from pyspark.traceback_utils import SCCallSiteSync
from pyspark.sql.column import Column
from pyspark.sql.classic.column import _to_seq, _to_list, _to_java_column
from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2
+from pyspark.sql.merge import MergeIntoWriter
from pyspark.sql.streaming import DataStreamWriter
from pyspark.sql.types import (
StructType,
@@ -68,7 +69,7 @@ from pyspark.sql.dataframe import (
DataFrameNaFunctions as ParentDataFrameNaFunctions,
DataFrameStatFunctions as ParentDataFrameStatFunctions,
)
-from pyspark.sql.utils import get_active_spark_context, toJArray
+from pyspark.sql.utils import get_active_spark_context, to_java_array, to_scala_map
from pyspark.sql.pandas.conversion import PandasConversionMixin
from pyspark.sql.pandas.map_ops import PandasMapOpsMixin
@@ -431,7 +432,7 @@ class DataFrame(ParentDataFrame, PandasMapOpsMixin, PandasConversionMixin):
int: gateway.jvm.long,
}
jclass = mapping[type(parameter[0])]
- return toJArray(gateway, jclass, parameter)
+ return to_java_array(gateway, jclass, parameter)
else:
return parameter
@@ -890,7 +891,7 @@ class DataFrame(ParentDataFrame, PandasMapOpsMixin, PandasConversionMixin):
def _jmap(self, jm: Dict) -> "JavaObject":
"""Return a JVM Scala Map from a dict"""
- return _to_scala_map(self.sparkSession._sc, jm)
+ return to_scala_map(self.sparkSession._sc._jvm, jm)
def _jcols(self, *cols: "ColumnOrName") -> "JavaObject":
"""Return a JVM Seq of Columns from a list of Column or column names
@@ -1797,6 +1798,9 @@ class DataFrame(ParentDataFrame, PandasMapOpsMixin, PandasConversionMixin):
def writeTo(self, table: str) -> DataFrameWriterV2:
return DataFrameWriterV2(self, table)
+ def mergeInto(self, table: str, condition: Column) -> MergeIntoWriter:
+ return MergeIntoWriter(self, table, condition)
+
def pandas_api(
self, index_col: Optional[Union[str, List[str]]] = None
) -> "PandasOnSparkDataFrame":
@@ -1844,14 +1848,6 @@ class DataFrame(ParentDataFrame, PandasMapOpsMixin, PandasConversionMixin):
)
-def _to_scala_map(sc: "SparkContext", jm: Dict) -> "JavaObject":
- """
- Convert a dict into a JVM Map.
- """
- assert sc._jvm is not None
- return sc._jvm.PythonUtils.toScalaMap(jm)
-
-
class DataFrameNaFunctions(ParentDataFrameNaFunctions):
def __init__(self, df: ParentDataFrame):
self.df = df
diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py
index 46698c2530ea..92cadb598c59 100644
--- a/python/pyspark/sql/connect/dataframe.py
+++ b/python/pyspark/sql/connect/dataframe.py
@@ -23,6 +23,7 @@ from pyspark.errors.exceptions.base import (
)
from pyspark.resource import ResourceProfile
from pyspark.sql.connect.utils import check_dependencies
+from pyspark.sql.merge import MergeIntoWriter
check_dependencies(__name__)
@@ -2186,6 +2187,12 @@ class DataFrame(ParentDataFrame):
return DataFrameWriterV2(self._plan, self._session, table, cb)
+ def mergeInto(self, table: str, condition: Column) -> MergeIntoWriter:
+ raise PySparkNotImplementedError(
+ error_class="NOT_IMPLEMENTED",
+ message_parameters={"feature": "mergeInto"},
+ )
+
def offset(self, n: int) -> ParentDataFrame:
return DataFrame(plan.Offset(child=self._plan, offset=n), session=self._session)
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index 625678588bf9..8d16604879bf 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -38,6 +38,7 @@ from pyspark.storagelevel import StorageLevel
from pyspark.resource import ResourceProfile
from pyspark.sql.column import Column
from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2
+from pyspark.sql.merge import MergeIntoWriter
from pyspark.sql.streaming import DataStreamWriter
from pyspark.sql.types import StructType, Row
from pyspark.sql.utils import dispatch_df_method
@@ -5985,6 +5986,41 @@ class DataFrame:
"""
...
+ @dispatch_df_method
+ def mergeInto(self, table: str, condition: Column) -> MergeIntoWriter:
+ """
+ Merges a set of updates, insertions, and deletions based on a source table into
+ a target table.
+
+ .. versionadded:: 4.0.0
+
+ Parameters
+ ----------
+ table : str
+ Target table name to merge into.
+ condition : :class:`Column`
+ The condition that determines whether a row in the target table matches one in the
+ source DataFrame.
+
+ Returns
+ -------
+ :class:`MergeIntoWriter`
+ MergeIntoWriter to use further to specify how to merge the source DataFrame
+ into the target table.
+
+ Examples
+ --------
+ >>> from pyspark.sql.functions import expr
+ >>> source = spark.createDataFrame(
+ ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["id", "name"])
+ >>> (source.mergeInto("target", "id") # doctest: +SKIP
+ ... .whenMatched().update({ "name": source.name })
+ ... .whenNotMatched().insertAll()
+ ... .whenNotMatchedBySource().delete()
+ ... .merge())
+ """
+ ...
+
@dispatch_df_method
def pandas_api(
self, index_col: Optional[Union[str, List[str]]] = None
diff --git a/python/pyspark/sql/merge.py b/python/pyspark/sql/merge.py
new file mode 100644
index 000000000000..c5ddecbee1d1
--- /dev/null
+++ b/python/pyspark/sql/merge.py
@@ -0,0 +1,251 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import sys
+from typing import Dict, Optional, TYPE_CHECKING
+
+from pyspark.sql.column import Column
+from pyspark.sql.utils import to_scala_map
+
+if TYPE_CHECKING:
+ from pyspark.sql.dataframe import DataFrame
+
+__all__ = ["MergeIntoWriter"]
+
+
+class MergeIntoWriter:
+ """
+ `MergeIntoWriter` provides methods to define and execute merge actions based
+ on specified conditions.
+
+ .. versionadded: 4.0.0
+ """
+
+ def __init__(self, df: "DataFrame", table: str, condition: Column):
+ self._df = df
+ self._spark = df.sparkSession
+
+ from pyspark.sql.classic.column import _to_java_column
+
+ self._jwriter = df._jdf.mergeInto(table, _to_java_column(condition))
+
+ def whenMatched(self, condition: Optional[Column] = None) -> "MergeIntoWriter.WhenMatched":
+ """
+ Initialize a `WhenMatched` action with a condition.
+
+ This `WhenMatched` action will be executed when a source row matches a target table row
+ based on the merge condition and the specified `condition` is satisfied.
+
+ This `WhenMatched` can be followed by one of the following merge actions:
+ - `updateAll`: Update all the matched target table rows with source dataset rows.
+ - `update(Dict)`: Update all the matched target table rows while changing only
+ a subset of columns based on the provided assignment.
+ - `delete`: Delete all target rows that have a match in the source table.
+ """
+ return self.WhenMatched(self, condition)
+
+ def whenNotMatched(
+ self, condition: Optional[Column] = None
+ ) -> "MergeIntoWriter.WhenNotMatched":
+ """
+ Initialize a `WhenNotMatched` action with a condition.
+ This `WhenNotMatched` action will be executed when a source row does not match any target
+ row based on the merge condition and the specified `condition` is satisfied.
+
+ This `WhenNotMatched` can be followed by one of the following merge actions:
+ - `insertAll`: Insert all rows from the source that are not already in the target table.
+ - `insert(Dict)`: Insert all rows from the source that are not already in the target
+ table, with the specified columns based on the provided assignment.
+ """
+ return self.WhenNotMatched(self, condition)
+
+ def whenNotMatchedBySource(
+ self, condition: Optional[Column] = None
+ ) -> "MergeIntoWriter.WhenNotMatchedBySource":
+ """
+ Initialize a `WhenNotMatchedBySource` action with a condition.
+ This `WhenNotMatchedBySource` action will be executed when a target row does not match any
+ rows in the source table based on the merge condition and the specified `condition`
+ is satisfied.
+ This `WhenNotMatchedBySource` can be followed by one of the following merge actions:
+ - `updateAll`: Update all the not matched target table rows with source dataset rows.
+ - `update(Dict)`: Update all the not matched target table rows while changing only
+ the specified columns based on the provided assignment.
+ - `delete`: Delete all target rows that have no matches in the source table.
+ """
+ return self.WhenNotMatchedBySource(self, condition)
+
+ def withSchemaEvolution(self) -> "MergeIntoWriter":
+ """
+ Enable automatic schema evolution for this merge operation.
+ """
+ self._jwriter = self._jwriter.withSchemaEvolution()
+ return self
+
+ def merge(self) -> None:
+ """
+ Execute the merge operation.
+ """
+ self._jwriter.merge()
+
+ class WhenMatched:
+ """
+ A class for defining actions to be taken when matching rows in a DataFrame during
+ a merge operation."""
+
+ def __init__(self, writer: "MergeIntoWriter", condition: Optional[Column]):
+ self.writer = writer
+ if condition is None:
+ self.when_matched = writer._jwriter.whenMatched()
+ else:
+ from pyspark.sql.classic.column import _to_java_column
+
+ self.when_matched = writer._jwriter.whenMatched(_to_java_column(condition))
+
+ def updateAll(self) -> "MergeIntoWriter":
+ """
+ Specifies an action to update all matched rows in the DataFrame.
+ """
+ self.writer._jwriter = self.when_matched.updateAll()
+ return self.writer
+
+ def update(self, assignments: Dict[str, Column]) -> "MergeIntoWriter":
+ """
+ Specifies an action to update matched rows in the DataFrame with the provided column
+ assignments.
+ """
+ jvm = self.writer._spark._jvm
+ from pyspark.sql.classic.column import _to_java_column
+
+ jmap = to_scala_map(jvm, {k: _to_java_column(v) for k, v in assignments.items()})
+ self.writer._jwriter = self.when_matched.update(jmap)
+ return self.writer
+
+ def delete(self) -> "MergeIntoWriter":
+ """
+ Specifies an action to delete matched rows from the DataFrame.
+ """
+ self.writer._jwriter = self.when_matched.delete()
+ return self.writer
+
+ class WhenNotMatched:
+ """
+ A class for defining actions to be taken when no matching rows are found in a DataFrame
+ during a merge operation."""
+
+ def __init__(self, writer: "MergeIntoWriter", condition: Optional[Column]):
+ self.writer = writer
+ if condition is None:
+ self.when_not_matched = writer._jwriter.whenNotMatched()
+ else:
+ from pyspark.sql.classic.column import _to_java_column
+
+ self.when_not_matched = writer._jwriter.whenNotMatched(_to_java_column(condition))
+
+ def insertAll(self) -> "MergeIntoWriter":
+ """
+ Specifies an action to insert all non-matched rows into the DataFrame.
+ """
+ self.writer._jwriter = self.when_not_matched.insertAll()
+ return self.writer
+
+ def insert(self, assignments: Dict[str, Column]) -> "MergeIntoWriter":
+ """
+ Specifies an action to insert non-matched rows into the DataFrame with the provided
+ column assignments.
+ """
+ jvm = self.writer._spark._jvm
+ from pyspark.sql.classic.column import _to_java_column
+
+ jmap = to_scala_map(jvm, {k: _to_java_column(v) for k, v in assignments.items()})
+ self.writer._jwriter = self.when_not_matched.insert(jmap)
+ return self.writer
+
+ class WhenNotMatchedBySource:
+ """
+ A class for defining actions to be performed when there is no match by source
+ during a merge operation in a MergeIntoWriter.
+ """
+
+ def __init__(self, writer: "MergeIntoWriter", condition: Optional[Column]):
+ self.writer = writer
+ if condition is None:
+ self.when_not_matched_by_source = writer._jwriter.whenNotMatchedBySource()
+ else:
+ from pyspark.sql.classic.column import _to_java_column
+
+ self.when_not_matched_by_source = writer._jwriter.whenNotMatchedBySource(
+ _to_java_column(condition)
+ )
+
+ def updateAll(self) -> "MergeIntoWriter":
+ """
+ Specifies an action to update all non-matched rows in the target DataFrame when
+ not matched by the source.
+ """
+ self.writer._jwriter = self.when_not_matched_by_source.updateAll()
+ return self.writer
+
+ def update(self, assignments: Dict[str, Column]) -> "MergeIntoWriter":
+ """
+ Specifies an action to update non-matched rows in the target DataFrame with the provided
+ column assignments when not matched by the source.
+ """
+ jvm = self.writer._spark._jvm
+ from pyspark.sql.classic.column import _to_java_column
+
+ jmap = to_scala_map(jvm, {k: _to_java_column(v) for k, v in assignments.items()})
+ self.writer._jwriter = self.when_not_matched_by_source.update(jmap)
+ return self.writer
+
+ def delete(self) -> "MergeIntoWriter":
+ """
+ Specifies an action to delete matched rows from the DataFrame.
+ """
+ self.writer._jwriter = self.when_not_matched_by_source.delete()
+ return self.writer
+
+
+def _test() -> None:
+ import doctest
+ import os
+ import py4j
+ from pyspark.core.context import SparkContext
+ from pyspark.sql import SparkSession
+ import pyspark.sql.merge
+
+ os.chdir(os.environ["SPARK_HOME"])
+
+ globs = pyspark.sql.merge.__dict__.copy()
+ sc = SparkContext("local[4]", "PythonTest")
+ try:
+ spark = SparkSession._getActiveSessionOrCreate()
+ except py4j.protocol.Py4JError:
+ spark = SparkSession(sc)
+
+ globs["spark"] = spark
+ (failure_count, test_count) = doctest.testmod(
+ pyspark.sql.merge,
+ globs=globs,
+ optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF,
+ )
+ spark.stop()
+ if failure_count:
+ sys.exit(-1)
+
+
+if __name__ == "__main__":
+ _test()
diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py
index 0ee235c8bf02..9d05d85da013 100644
--- a/python/pyspark/sql/readwriter.py
+++ b/python/pyspark/sql/readwriter.py
@@ -1173,7 +1173,7 @@ class DataFrameReader(OptionUtils):
if predicates is not None:
gateway = self._spark._sc._gateway
assert gateway is not None
- jpredicates = utils.toJArray(gateway, gateway.jvm.java.lang.String, predicates)
+ jpredicates = utils.to_java_array(gateway, gateway.jvm.java.lang.String, predicates)
return self._df(self._jreader.jdbc(url, table, jpredicates, jprop))
return self._df(self._jreader.jdbc(url, table, jprop))
diff --git a/python/pyspark/sql/tests/connect/test_parity_dataframe.py b/python/pyspark/sql/tests/connect/test_parity_dataframe.py
index 343f485553a9..23a55281354b 100644
--- a/python/pyspark/sql/tests/connect/test_parity_dataframe.py
+++ b/python/pyspark/sql/tests/connect/test_parity_dataframe.py
@@ -30,6 +30,10 @@ class DataFrameParityTests(DataFrameTestsMixin, ReusedConnectTestCase):
def test_toDF_with_schema_string(self):
super().test_toDF_with_schema_string()
+ @unittest.skip("Spark Connect does not support mergeInto at the moment.")
+ def test_df_merge_into(self):
+ super().test_df_merge_into()
+
if __name__ == "__main__":
import unittest
diff --git a/python/pyspark/sql/tests/test_dataframe.py b/python/pyspark/sql/tests/test_dataframe.py
index c7cf43a33454..28ff19dc2e18 100644
--- a/python/pyspark/sql/tests/test_dataframe.py
+++ b/python/pyspark/sql/tests/test_dataframe.py
@@ -742,6 +742,79 @@ class DataFrameTestsMixin:
message_parameters={"arg_name": "truncate", "arg_type": "str"},
)
+ def test_df_merge_into(self):
+ try:
+ # InMemoryRowLevelOperationTableCatalog is a test catalog that is included in the
+ # catalyst-test package. If Spark complains that it can't find this class, make sure
+ # the catalyst-test JAR does exist in "SPARK_HOME/assembly/target/scala-x.xx/jars"
+ # directory. If not, build it with `build/sbt test:assembly` and copy it over.
+ self.spark.conf.set(
+ "spark.sql.catalog.testcat",
+ "org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog",
+ )
+ with self.table("testcat.ns1.target"):
+
+ def reset_target_table():
+ self.spark.createDataFrame(
+ [(1, "Alice"), (2, "Bob")], ["id", "name"]
+ ).write.mode("overwrite").saveAsTable("testcat.ns1.target")
+
+ source = self.spark.createDataFrame(
+ [(1, "Charlie"), (3, "David")], ["id", "name"]
+ ) # type: DataFrame
+
+ from pyspark.sql.functions import col
+
+ # Match -> update, NotMatch -> insert, NotMatchedBySource -> delete
+ reset_target_table()
+ # fmt: off
+ source.mergeInto("testcat.ns1.target", source.id == col("target.id")) \
+ .whenMatched(source.id == 1).update({"name": source.name}) \
+ .whenNotMatched().insert({"id": source.id, "name": source.name}) \
+ .whenNotMatchedBySource().delete() \
+ .merge()
+ # fmt: on
+ self.assertEqual(
+ self.spark.table("testcat.ns1.target").orderBy("id").collect(),
+ [Row(id=1, name="Charlie"), Row(id=3, name="David")],
+ )
+
+ # Match -> updateAll, NotMatch -> insertAll, NotMatchedBySource -> update
+ reset_target_table()
+ # fmt: off
+ source.mergeInto("testcat.ns1.target", source.id == col("target.id")) \
+ .whenMatched(source.id == 1).updateAll() \
+ .whenNotMatched(source.id == 3).insertAll() \
+ .whenNotMatchedBySource(col("target.id") == lit(2)) \
+ .update({"name": lit("not_matched")}) \
+ .merge()
+ # fmt: on
+ self.assertEqual(
+ self.spark.table("testcat.ns1.target").orderBy("id").collect(),
+ [
+ Row(id=1, name="Charlie"),
+ Row(id=2, name="not_matched"),
+ Row(id=3, name="David"),
+ ],
+ )
+
+ # Match -> delete, NotMatchedBySource -> delete
+ reset_target_table()
+ # fmt: off
+ self.spark.createDataFrame([(1, "AliceJr")], ["id", "name"]) \
+ .write.mode("append").saveAsTable("testcat.ns1.target")
+ source.mergeInto("testcat.ns1.target", source.id == col("target.id")) \
+ .whenMatched(col("target.name") != lit("AliceJr")).delete() \
+ .whenNotMatchedBySource().delete() \
+ .merge()
+ # fmt: on
+ self.assertEqual(
+ self.spark.table("testcat.ns1.target").orderBy("id").collect(),
+ [Row(id=1, name="AliceJr")],
+ )
+ finally:
+ self.spark.conf.unset("spark.sql.catalog.testcat")
+
@unittest.skipIf(
not have_pandas or not have_pyarrow,
cast(str, pandas_requirement_message or pyarrow_requirement_message),
diff --git a/python/pyspark/sql/utils.py b/python/pyspark/sql/utils.py
index 76227851f9fe..e6094502c084 100644
--- a/python/pyspark/sql/utils.py
+++ b/python/pyspark/sql/utils.py
@@ -20,6 +20,7 @@ import os
from typing import (
Any,
Callable,
+ Dict,
Optional,
List,
Sequence,
@@ -52,6 +53,7 @@ if TYPE_CHECKING:
JavaClass,
JavaGateway,
JavaObject,
+ JVMView,
)
from pyspark import SparkContext
from pyspark.sql.session import SparkSession
@@ -70,7 +72,7 @@ except ImportError:
FuncT = TypeVar("FuncT", bound=Callable[..., Any])
-def toJArray(gateway: "JavaGateway", jtype: "JavaClass", arr: Sequence[Any]) -> "JavaArray":
+def to_java_array(gateway: "JavaGateway", jtype: "JavaClass", arr: Sequence[Any]) -> "JavaArray":
"""
Convert python list to java type array
@@ -89,6 +91,14 @@ def toJArray(gateway: "JavaGateway", jtype: "JavaClass", arr: Sequence[Any]) ->
return jarray
+def to_scala_map(jvm: "JVMView", dic: Dict) -> "JavaObject":
+ """
+ Convert a dict into a Scala Map.
+ """
+ assert jvm is not None
+ return jvm.PythonUtils.toScalaMap(dic)
+
+
def require_test_compiled() -> None:
"""Raise Exception if test classes are not compiled"""
import os
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org