ENH: add validate as a param to join (#46622) (#46740)

Author: gaotian98

doc/source/whatsnew/v1.5.0.rst

@@ -95,6 +95,7 @@ Other enhancements
 - :meth:`pd.concat` now raises when ``levels`` is given but ``keys`` is None (:issue:`46653`)
 - :meth:`pd.concat` now raises when ``levels`` contains duplicate values (:issue:`46653`)
 - Added ``numeric_only`` argument to :meth:`DataFrame.corr`, :meth:`DataFrame.corrwith`, and :meth:`DataFrame.cov` (:issue:`46560`)
+- Added ``validate`` argument to :meth:`DataFrame.join` (:issue:`46622`)
 - A :class:`errors.PerformanceWarning` is now thrown when using ``string[pyarrow]`` dtype with methods that don't dispatch to ``pyarrow.compute`` methods (:issue:`42613`)
 - Added ``numeric_only`` argument to :meth:`Resampler.sum`, :meth:`Resampler.prod`, :meth:`Resampler.min`, :meth:`Resampler.max`, :meth:`Resampler.first`, and :meth:`Resampler.last` (:issue:`46442`)
 

pandas/core/frame.py

@@ -9385,6 +9385,7 @@ def join(
         lsuffix: str = "",
         rsuffix: str = "",
         sort: bool = False,
+        validate: str | None = None,
     ) -> DataFrame:
         """
         Join columns of another DataFrame.
@@ -9428,6 +9429,14 @@ def join(
         sort : bool, default False
             Order result DataFrame lexicographically by the join key. If False,
             the order of the join key depends on the join type (how keyword).
+        validate : str, optional
+            If specified, checks if join is of specified type.
+            * "one_to_one" or "1:1": check if join keys are unique in both left
+            and right datasets.
+            * "one_to_many" or "1:m": check if join keys are unique in left dataset.
+            * "many_to_one" or "m:1": check if join keys are unique in right dataset.
+            * "many_to_many" or "m:m": allowed, but does not result in checks.
+            .. versionadded:: 1.5.0
 
         Returns
         -------
@@ -9522,7 +9531,7 @@ def join(
         4  K0  A4
         5  K1  A5
 
-        >>> df.join(other.set_index('key'), on='key')
+        >>> df.join(other.set_index('key'), on='key', validate='m:1')
           key   A    B
         0  K0  A0   B0
         1  K1  A1   B1
@@ -9532,7 +9541,13 @@ def join(
         5  K1  A5   B1
         """
         return self._join_compat(
-            other, on=on, how=how, lsuffix=lsuffix, rsuffix=rsuffix, sort=sort
+            other,
+            on=on,
+            how=how,
+            lsuffix=lsuffix,
+            rsuffix=rsuffix,
+            sort=sort,
+            validate=validate,
         )
 
     def _join_compat(
@@ -9543,6 +9558,7 @@ def _join_compat(
         lsuffix: str = "",
         rsuffix: str = "",
         sort: bool = False,
+        validate: str | None = None,
     ):
         from pandas.core.reshape.concat import concat
         from pandas.core.reshape.merge import merge
@@ -9561,6 +9577,7 @@ def _join_compat(
                     on=on,
                     suffixes=(lsuffix, rsuffix),
                     sort=sort,
+                    validate=validate,
                 )
             return merge(
                 self,
@@ -9571,6 +9588,7 @@ def _join_compat(
                 right_index=True,
                 suffixes=(lsuffix, rsuffix),
                 sort=sort,
+                validate=validate,
             )
         else:
             if on is not None:
@@ -9603,7 +9621,12 @@ def _join_compat(
 
             for frame in frames[1:]:
                 joined = merge(
-                    joined, frame, how=how, left_index=True, right_index=True
+                    joined,
+                    frame,
+                    how=how,
+                    left_index=True,
+                    right_index=True,
+                    validate=validate,
                 )
 
             return joined

pandas/tests/frame/methods/test_join.py

@@ -3,6 +3,8 @@
 import numpy as np
 import pytest
 
+from pandas.errors import MergeError
+
 import pandas as pd
 from pandas import (
     DataFrame,
@@ -12,6 +14,7 @@
     period_range,
 )
 import pandas._testing as tm
+from pandas.core.reshape.concat import concat
 
 
 @pytest.fixture
@@ -33,6 +36,39 @@ def right():
     return DataFrame({"b": [300, 100, 200]}, index=[3, 1, 2])
 
 
+@pytest.fixture
+def left_no_dup():
+    return DataFrame(
+        {"a": ["a", "b", "c", "d"], "b": ["cat", "dog", "weasel", "horse"]},
+        index=range(4),
+    )
+
+
+@pytest.fixture
+def right_no_dup():
+    return DataFrame(
+        {
+            "a": ["a", "b", "c", "d", "e"],
+            "c": ["meow", "bark", "um... weasel noise?", "nay", "chirp"],
+        },
+        index=range(5),
+    ).set_index("a")
+
+
+@pytest.fixture
+def left_w_dups(left_no_dup):
+    return concat(
+        [left_no_dup, DataFrame({"a": ["a"], "b": ["cow"]}, index=[3])], sort=True
+    )
+
+
+@pytest.fixture
+def right_w_dups(right_no_dup):
+    return concat(
+        [right_no_dup, DataFrame({"a": ["e"], "c": ["moo"]}, index=[3])]
+    ).set_index("a")
+
+
 @pytest.mark.parametrize(
     "how, sort, expected",
     [
@@ -78,7 +114,7 @@ def right():
 )
 def test_join(left, right, how, sort, expected):
 
-    result = left.join(right, how=how, sort=sort)
+    result = left.join(right, how=how, sort=sort, validate="1:1")
     tm.assert_frame_equal(result, expected)
 
 
@@ -104,6 +140,112 @@ def test_suffix_on_list_join():
     tm.assert_frame_equal(arr_joined, norm_joined)
 
 
+def test_join_invalid_validate(left_no_dup, right_no_dup):
+    # GH 46622
+    # Check invalid arguments
+    msg = "Not a valid argument for validate"
+    with pytest.raises(ValueError, match=msg):
+        left_no_dup.merge(right_no_dup, on="a", validate="invalid")
+
+
+def test_join_on_single_col_dup_on_right(left_no_dup, right_w_dups):
+    # GH 46622
+    # Dups on right allowed by one_to_many constraint
+    left_no_dup.join(
+        right_w_dups,
+        on="a",
+        validate="one_to_many",
+    )
+
+    # Dups on right not allowed by one_to_one constraint
+    msg = "Merge keys are not unique in right dataset; not a one-to-one merge"
+    with pytest.raises(MergeError, match=msg):
+        left_no_dup.join(
+            right_w_dups,
+            on="a",
+            validate="one_to_one",
+        )
+
+
+def test_join_on_single_col_dup_on_left(left_w_dups, right_no_dup):
+    # GH 46622
+    # Dups on left allowed by many_to_one constraint
+    left_w_dups.join(
+        right_no_dup,
+        on="a",
+        validate="many_to_one",
+    )
+
+    # Dups on left not allowed by one_to_one constraint
+    msg = "Merge keys are not unique in left dataset; not a one-to-one merge"
+    with pytest.raises(MergeError, match=msg):
+        left_w_dups.join(
+            right_no_dup,
+            on="a",
+            validate="one_to_one",
+        )
+
+
+def test_join_on_single_col_dup_on_both(left_w_dups, right_w_dups):
+    # GH 46622
+    # Dups on both allowed by many_to_many constraint
+    left_w_dups.join(right_w_dups, on="a", validate="many_to_many")
+
+    # Dups on both not allowed by many_to_one constraint
+    msg = "Merge keys are not unique in right dataset; not a many-to-one merge"
+    with pytest.raises(MergeError, match=msg):
+        left_w_dups.join(
+            right_w_dups,
+            on="a",
+            validate="many_to_one",
+        )
+
+    # Dups on both not allowed by one_to_many constraint
+    msg = "Merge keys are not unique in left dataset; not a one-to-many merge"
+    with pytest.raises(MergeError, match=msg):
+        left_w_dups.join(
+            right_w_dups,
+            on="a",
+            validate="one_to_many",
+        )
+
+
+def test_join_on_multi_col_check_dup():
+    # GH 46622
+    # Two column join, dups in both, but jointly no dups
+    left = DataFrame(
+        {
+            "a": ["a", "a", "b", "b"],
+            "b": [0, 1, 0, 1],
+            "c": ["cat", "dog", "weasel", "horse"],
+        },
+        index=range(4),
+    ).set_index(["a", "b"])
+
+    right = DataFrame(
+        {
+            "a": ["a", "a", "b"],
+            "b": [0, 1, 0],
+            "d": ["meow", "bark", "um... weasel noise?"],
+        },
+        index=range(3),
+    ).set_index(["a", "b"])
+
+    expected_multi = DataFrame(
+        {
+            "a": ["a", "a", "b"],
+            "b": [0, 1, 0],
+            "c": ["cat", "dog", "weasel"],
+            "d": ["meow", "bark", "um... weasel noise?"],
+        },
+        index=range(3),
+    ).set_index(["a", "b"])
+
+    # Jointly no dups allowed by one_to_one constraint
+    result = left.join(right, how="inner", validate="1:1")
+    tm.assert_frame_equal(result, expected_multi)
+
+
 def test_join_index(float_frame):
     # left / right