import numpy as np import pytest import pandas.util._test_decorators as td import pandas as pd import pandas._testing as tm from pandas.core.reshape.concat import concat from pandas.core.reshape.merge import merge @pytest.fixture def left(): """left dataframe (not multi-indexed) for multi-index join tests""" # a little relevant example with NAs key1 = ["bar", "bar", "bar", "foo", "foo", "baz", "baz", "qux", "qux", "snap"] key2 = ["two", "one", "three", "one", "two", "one", "two", "two", "three", "one"] data = np.random.default_rng(2).standard_normal(len(key1)) return pd.DataFrame({"key1": key1, "key2": key2, "data": data}) @pytest.fixture def right(multiindex_dataframe_random_data): """right dataframe (multi-indexed) for multi-index join tests""" df = multiindex_dataframe_random_data df.index.names = ["key1", "key2"] df.columns = ["j_one", "j_two", "j_three"] return df @pytest.fixture def left_multi(): return pd.DataFrame( { "Origin": ["A", "A", "B", "B", "C"], "Destination": ["A", "B", "A", "C", "A"], "Period": ["AM", "AM", "IP", "AM", "OP"], "TripPurp": ["hbw", "nhb", "hbo", "nhb", "hbw"], "Trips": [1987, 3647, 2470, 4296, 4444], }, columns=["Origin", "Destination", "Period", "TripPurp", "Trips"], ).set_index(["Origin", "Destination", "Period", "TripPurp"]) @pytest.fixture def right_multi(): return pd.DataFrame( { "Origin": ["A", "A", "B", "B", "C", "C", "E"], "Destination": ["A", "B", "A", "B", "A", "B", "F"], "Period": ["AM", "AM", "IP", "AM", "OP", "IP", "AM"], "LinkType": ["a", "b", "c", "b", "a", "b", "a"], "Distance": [100, 80, 90, 80, 75, 35, 55], }, columns=["Origin", "Destination", "Period", "LinkType", "Distance"], ).set_index(["Origin", "Destination", "Period", "LinkType"]) @pytest.fixture def on_cols_multi(): return ["Origin", "Destination", "Period"] class TestMergeMulti: def test_merge_on_multikey(self, left, right, join_type): on_cols = ["key1", "key2"] result = left.join(right, on=on_cols, how=join_type).reset_index(drop=True) expected = merge(left, right.reset_index(), on=on_cols, how=join_type) tm.assert_frame_equal(result, expected) result = left.join(right, on=on_cols, how=join_type, sort=True).reset_index( drop=True ) expected = merge( left, right.reset_index(), on=on_cols, how=join_type, sort=True ) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( "infer_string", [False, pytest.param(True, marks=td.skip_if_no("pyarrow"))] ) @pytest.mark.filterwarnings( "ignore:The 'future.infer_string' option:pandas.errors.Pandas4Warning" ) def test_left_join_multi_index(self, sort, infer_string): with pd.option_context("future.infer_string", infer_string): icols = ["1st", "2nd", "3rd"] def bind_cols(df): iord = lambda a: 0 if a != a else ord(a) f = lambda ts: ts.map(iord) - ord("a") return f(df["1st"]) + f(df["3rd"]) * 1e2 + df["2nd"].fillna(0) * 10 def run_asserts(left, right, sort): res = left.join(right, on=icols, how="left", sort=sort) assert len(left) < len(res) + 1 assert not res["4th"].isna().any() assert not res["5th"].isna().any() tm.assert_series_equal(res["4th"], -res["5th"], check_names=False) result = bind_cols(res.iloc[:, :-2]) tm.assert_series_equal(res["4th"], result, check_names=False) assert result.name is None if sort: tm.assert_frame_equal(res, res.sort_values(icols, kind="mergesort")) out = merge(left, right.reset_index(), on=icols, sort=sort, how="left") res.index = pd.RangeIndex(len(res)) tm.assert_frame_equal(out, res) lc = list(map(chr, np.arange(ord("a"), ord("z") + 1))) left = pd.DataFrame( np.random.default_rng(2).choice(lc, (50, 2)), columns=["1st", "3rd"] ) # Explicit cast to float to avoid implicit cast when setting nan left.insert( 1, "2nd", np.random.default_rng(2).integers(0, 10, len(left)).astype("float"), ) right = left.sample(frac=1, random_state=np.random.default_rng(2)) left["4th"] = bind_cols(left) right["5th"] = -bind_cols(right) right = right.set_index(icols) run_asserts(left, right, sort) # inject some nulls left.loc[1::4, "1st"] = np.nan left.loc[2::5, "2nd"] = np.nan left.loc[3::6, "3rd"] = np.nan left["4th"] = bind_cols(left) i = np.random.default_rng(2).permutation(len(left)) right = left.iloc[i, :-1] right["5th"] = -bind_cols(right) right = right.set_index(icols) run_asserts(left, right, sort) def test_merge_right_vs_left(self, left, right, sort): # compare left vs right merge with multikey on_cols = ["key1", "key2"] merged_left_right = left.merge( right, left_on=on_cols, right_index=True, how="left", sort=sort ) merge_right_left = right.merge( left, right_on=on_cols, left_index=True, how="right", sort=sort ) # Reorder columns merge_right_left = merge_right_left[merged_left_right.columns] tm.assert_frame_equal(merged_left_right, merge_right_left) def test_merge_multiple_cols_with_mixed_cols_index(self): # GH29522 s = pd.Series( range(6), pd.MultiIndex.from_product([["A", "B"], [1, 2, 3]], names=["lev1", "lev2"]), name="Amount", ) df = pd.DataFrame( {"lev1": list("AAABBB"), "lev2": [1, 2, 3, 1, 2, 3], "col": 0} ) result = merge(df, s.reset_index(), on=["lev1", "lev2"]) expected = pd.DataFrame( { "lev1": list("AAABBB"), "lev2": [1, 2, 3, 1, 2, 3], "col": [0] * 6, "Amount": range(6), } ) tm.assert_frame_equal(result, expected) def test_compress_group_combinations(self): # ~ 40000000 possible unique groups key1 = [str(i) for i in range(10000)] key1 = np.tile(key1, 2) key2 = key1[::-1] df = pd.DataFrame( { "key1": key1, "key2": key2, "value1": np.random.default_rng(2).standard_normal(20000), } ) df2 = pd.DataFrame( { "key1": key1[::2], "key2": key2[::2], "value2": np.random.default_rng(2).standard_normal(10000), } ) # just to hit the label compression code path merge(df, df2, how="outer") def test_left_join_index_preserve_order(self): on_cols = ["k1", "k2"] left = pd.DataFrame( { "k1": [0, 1, 2] * 8, "k2": ["foo", "bar"] * 12, "v": np.array(np.arange(24), dtype=np.int64), } ) index = pd.MultiIndex.from_tuples([(2, "bar"), (1, "foo")]) right = pd.DataFrame({"v2": [5, 7]}, index=index) result = left.join(right, on=on_cols) expected = left.copy() expected["v2"] = np.nan expected.loc[(expected.k1 == 2) & (expected.k2 == "bar"), "v2"] = 5 expected.loc[(expected.k1 == 1) & (expected.k2 == "foo"), "v2"] = 7 tm.assert_frame_equal(result, expected) result = left.join(right, on=on_cols, sort=True) expected = expected.sort_values(on_cols, kind="mergesort") tm.assert_frame_equal(result, expected) # test join with multi dtypes blocks left = pd.DataFrame( { "k1": [0, 1, 2] * 8, "k2": ["foo", "bar"] * 12, "k3": np.array([0, 1, 2] * 8, dtype=np.float32), "v": np.array(np.arange(24), dtype=np.int32), } ) index = pd.MultiIndex.from_tuples([(2, "bar"), (1, "foo")]) right = pd.DataFrame({"v2": [5, 7]}, index=index) result = left.join(right, on=on_cols) expected = left.copy() expected["v2"] = np.nan expected.loc[(expected.k1 == 2) & (expected.k2 == "bar"), "v2"] = 5 expected.loc[(expected.k1 == 1) & (expected.k2 == "foo"), "v2"] = 7 tm.assert_frame_equal(result, expected) result = result.sort_values(on_cols, kind="mergesort") expected = left.join(right, on=on_cols, sort=True) tm.assert_frame_equal(result, expected) def test_left_join_index_multi_match_multiindex(self): left = pd.DataFrame( [ ["X", "Y", "C", "a"], ["W", "Y", "C", "e"], ["V", "Q", "A", "h"], ["V", "R", "D", "i"], ["X", "Y", "D", "b"], ["X", "Y", "A", "c"], ["W", "Q", "B", "f"], ["W", "R", "C", "g"], ["V", "Y", "C", "j"], ["X", "Y", "B", "d"], ], columns=["cola", "colb", "colc", "tag"], index=[3, 2, 0, 1, 7, 6, 4, 5, 9, 8], ) right = pd.DataFrame( [ ["W", "R", "C", 0], ["W", "Q", "B", 3], ["W", "Q", "B", 8], ["X", "Y", "A", 1], ["X", "Y", "A", 4], ["X", "Y", "B", 5], ["X", "Y", "C", 6], ["X", "Y", "C", 9], ["X", "Q", "C", -6], ["X", "R", "C", -9], ["V", "Y", "C", 7], ["V", "R", "D", 2], ["V", "R", "D", -1], ["V", "Q", "A", -3], ], columns=["col1", "col2", "col3", "val"], ).set_index(["col1", "col2", "col3"]) result = left.join(right, on=["cola", "colb", "colc"], how="left") expected = pd.DataFrame( [ ["X", "Y", "C", "a", 6], ["X", "Y", "C", "a", 9], ["W", "Y", "C", "e", np.nan], ["V", "Q", "A", "h", -3], ["V", "R", "D", "i", 2], ["V", "R", "D", "i", -1], ["X", "Y", "D", "b", np.nan], ["X", "Y", "A", "c", 1], ["X", "Y", "A", "c", 4], ["W", "Q", "B", "f", 3], ["W", "Q", "B", "f", 8], ["W", "R", "C", "g", 0], ["V", "Y", "C", "j", 7], ["X", "Y", "B", "d", 5], ], columns=["cola", "colb", "colc", "tag", "val"], index=[3, 3, 2, 0, 1, 1, 7, 6, 6, 4, 4, 5, 9, 8], ) tm.assert_frame_equal(result, expected) result = left.join(right, on=["cola", "colb", "colc"], how="left", sort=True) expected = expected.sort_values(["cola", "colb", "colc"], kind="mergesort") tm.assert_frame_equal(result, expected) def test_left_join_index_multi_match(self): left = pd.DataFrame( [["c", 0], ["b", 1], ["a", 2], ["b", 3]], columns=["tag", "val"], index=[2, 0, 1, 3], ) right = pd.DataFrame( [ ["a", "v"], ["c", "w"], ["c", "x"], ["d", "y"], ["a", "z"], ["c", "r"], ["e", "q"], ["c", "s"], ], columns=["tag", "char"], ).set_index("tag") result = left.join(right, on="tag", how="left") expected = pd.DataFrame( [ ["c", 0, "w"], ["c", 0, "x"], ["c", 0, "r"], ["c", 0, "s"], ["b", 1, np.nan], ["a", 2, "v"], ["a", 2, "z"], ["b", 3, np.nan], ], columns=["tag", "val", "char"], index=[2, 2, 2, 2, 0, 1, 1, 3], ) tm.assert_frame_equal(result, expected) result = left.join(right, on="tag", how="left", sort=True) expected2 = expected.sort_values("tag", kind="mergesort") tm.assert_frame_equal(result, expected2) # GH7331 - maintain left frame order in left merge result = merge(left, right.reset_index(), how="left", on="tag") expected.index = pd.RangeIndex(len(expected)) tm.assert_frame_equal(result, expected) def test_left_merge_na_buglet(self): left = pd.DataFrame( { "id": list("abcde"), "v1": np.random.default_rng(2).standard_normal(5), "v2": np.random.default_rng(2).standard_normal(5), "dummy": list("abcde"), "v3": np.random.default_rng(2).standard_normal(5), }, columns=["id", "v1", "v2", "dummy", "v3"], ) right = pd.DataFrame( { "id": ["a", "b", np.nan, np.nan, np.nan], "sv3": [1.234, 5.678, np.nan, np.nan, np.nan], } ) result = merge(left, right, on="id", how="left") rdf = right.drop(["id"], axis=1) expected = left.join(rdf) tm.assert_frame_equal(result, expected) def test_merge_na_keys(self): data = [ [1950, "A", 1.5], [1950, "B", 1.5], [1955, "B", 1.5], [1960, "B", np.nan], [1970, "B", 4.0], [1950, "C", 4.0], [1960, "C", np.nan], [1965, "C", 3.0], [1970, "C", 4.0], ] frame = pd.DataFrame(data, columns=["year", "panel", "data"]) other_data = [ [1960, "A", np.nan], [1970, "A", np.nan], [1955, "A", np.nan], [1965, "A", np.nan], [1965, "B", np.nan], [1955, "C", np.nan], ] other = pd.DataFrame(other_data, columns=["year", "panel", "data"]) result = frame.merge(other, how="outer") expected = frame.fillna(-999).merge(other.fillna(-999), how="outer") expected = expected.replace(-999, np.nan) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("klass", [None, np.asarray, pd.Series, pd.Index]) def test_merge_datetime_index(self, klass): # see gh-19038 df = pd.DataFrame( [1, 2, 3], ["2016-01-01", "2017-01-01", "2018-01-01"], columns=["a"] ) df.index = pd.to_datetime(df.index) on_vector = df.index.year if klass is not None: on_vector = klass(on_vector) exp_years = np.array([2016, 2017, 2018], dtype=np.int32) expected = pd.DataFrame({"a": [1, 2, 3], "key_1": exp_years}) result = df.merge(df, on=["a", on_vector], how="inner") tm.assert_frame_equal(result, expected) expected = pd.DataFrame( {"key_0": exp_years, "a_x": [1, 2, 3], "a_y": [1, 2, 3]} ) result = df.merge(df, on=[df.index.year], how="inner") tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("merge_type", ["left", "right"]) def test_merge_datetime_multi_index_empty_df(self, merge_type): # see gh-36895 left = pd.DataFrame( data={ "data": [1.5, 1.5], }, index=pd.MultiIndex.from_tuples( [[pd.Timestamp("1950-01-01"), "A"], [pd.Timestamp("1950-01-02"), "B"]], names=["date", "panel"], ), ) right = pd.DataFrame( index=pd.MultiIndex.from_tuples([], names=["date", "panel"]), columns=["state"], ) expected_index = pd.MultiIndex.from_tuples( [[pd.Timestamp("1950-01-01"), "A"], [pd.Timestamp("1950-01-02"), "B"]], names=["date", "panel"], ) if merge_type == "left": expected = pd.DataFrame( data={ "data": [1.5, 1.5], "state": np.array([np.nan, np.nan], dtype=object), }, index=expected_index, ) results_merge = left.merge(right, how="left", on=["date", "panel"]) results_join = left.join(right, how="left") else: expected = pd.DataFrame( data={ "state": np.array([np.nan, np.nan], dtype=object), "data": [1.5, 1.5], }, index=expected_index, ) results_merge = right.merge(left, how="right", on=["date", "panel"]) results_join = right.join(left, how="right") tm.assert_frame_equal(results_merge, expected) tm.assert_frame_equal(results_join, expected) @pytest.fixture def household(self): household = pd.DataFrame( { "household_id": [1, 2, 3], "male": [0, 1, 0], "wealth": [196087.3, 316478.7, 294750], }, columns=["household_id", "male", "wealth"], ).set_index("household_id") return household @pytest.fixture def portfolio(self): portfolio = pd.DataFrame( { "household_id": [1, 2, 2, 3, 3, 3, 4], "asset_id": [ "nl0000301109", "nl0000289783", "gb00b03mlx29", "gb00b03mlx29", "lu0197800237", "nl0000289965", np.nan, ], "name": [ "ABN Amro", "Robeco", "Royal Dutch Shell", "Royal Dutch Shell", "AAB Eastern Europe Equity Fund", "Postbank BioTech Fonds", np.nan, ], "share": [1.0, 0.4, 0.6, 0.15, 0.6, 0.25, 1.0], }, columns=["household_id", "asset_id", "name", "share"], ).set_index(["household_id", "asset_id"]) return portfolio @pytest.fixture def expected(self): expected = ( pd.DataFrame( { "male": [0, 1, 1, 0, 0, 0], "wealth": [ 196087.3, 316478.7, 316478.7, 294750.0, 294750.0, 294750.0, ], "name": [ "ABN Amro", "Robeco", "Royal Dutch Shell", "Royal Dutch Shell", "AAB Eastern Europe Equity Fund", "Postbank BioTech Fonds", ], "share": [1.00, 0.40, 0.60, 0.15, 0.60, 0.25], "household_id": [1, 2, 2, 3, 3, 3], "asset_id": [ "nl0000301109", "nl0000289783", "gb00b03mlx29", "gb00b03mlx29", "lu0197800237", "nl0000289965", ], } ) .set_index(["household_id", "asset_id"]) .reindex(columns=["male", "wealth", "name", "share"]) ) return expected def test_join_multi_levels(self, portfolio, household, expected): portfolio = portfolio.copy() household = household.copy() # GH 3662 # merge multi-levels result = household.join(portfolio, how="inner") tm.assert_frame_equal(result, expected) def test_join_multi_levels_merge_equivalence(self, portfolio, household, expected): portfolio = portfolio.copy() household = household.copy() # equivalency result = merge( household.reset_index(), portfolio.reset_index(), on=["household_id"], how="inner", ).set_index(["household_id", "asset_id"]) tm.assert_frame_equal(result, expected) def test_join_multi_levels_outer(self, portfolio, household, expected): portfolio = portfolio.copy() household = household.copy() result = household.join(portfolio, how="outer") expected = concat( [ expected, ( pd.DataFrame( {"share": [1.00]}, index=pd.MultiIndex.from_tuples( [(4, np.nan)], names=["household_id", "asset_id"] ), ) ), ], axis=0, sort=True, ).reindex(columns=expected.columns) tm.assert_frame_equal(result, expected, check_index_type=False) def test_join_multi_levels_invalid(self, portfolio, household): portfolio = portfolio.copy() household = household.copy() # invalid cases household.index.name = "foo" with pytest.raises( ValueError, match="cannot join with no overlapping index names" ): household.join(portfolio, how="inner") portfolio2 = portfolio.copy() portfolio2.index.set_names(["household_id", "foo"]) with pytest.raises(ValueError, match="columns overlap but no suffix specified"): portfolio2.join(portfolio, how="inner") def test_join_multi_levels2(self): # some more advanced merges # GH6360 household = pd.DataFrame( { "household_id": [1, 2, 2, 3, 3, 3, 4], "asset_id": [ "nl0000301109", "nl0000301109", "gb00b03mlx29", "gb00b03mlx29", "lu0197800237", "nl0000289965", np.nan, ], "share": [1.0, 0.4, 0.6, 0.15, 0.6, 0.25, 1.0], }, columns=["household_id", "asset_id", "share"], ).set_index(["household_id", "asset_id"]) log_return = pd.DataFrame( { "asset_id": [ "gb00b03mlx29", "gb00b03mlx29", "gb00b03mlx29", "lu0197800237", "lu0197800237", ], "t": [233, 234, 235, 180, 181], "log_return": [ 0.09604978, -0.06524096, 0.03532373, 0.03025441, 0.036997, ], } ).set_index(["asset_id", "t"]) expected = ( pd.DataFrame( { "household_id": [2, 2, 2, 3, 3, 3, 3, 3], "asset_id": [ "gb00b03mlx29", "gb00b03mlx29", "gb00b03mlx29", "gb00b03mlx29", "gb00b03mlx29", "gb00b03mlx29", "lu0197800237", "lu0197800237", ], "t": [233, 234, 235, 233, 234, 235, 180, 181], "share": [0.6, 0.6, 0.6, 0.15, 0.15, 0.15, 0.6, 0.6], "log_return": [ 0.09604978, -0.06524096, 0.03532373, 0.09604978, -0.06524096, 0.03532373, 0.03025441, 0.036997, ], } ) .set_index(["household_id", "asset_id", "t"]) .reindex(columns=["share", "log_return"]) ) # this is the equivalency result = merge( household.reset_index(), log_return.reset_index(), on=["asset_id"], how="inner", ).set_index(["household_id", "asset_id", "t"]) tm.assert_frame_equal(result, expected) expected = ( pd.DataFrame( { "household_id": [2, 2, 2, 3, 3, 3, 3, 3, 3, 1, 2, 4], "asset_id": [ "gb00b03mlx29", "gb00b03mlx29", "gb00b03mlx29", "gb00b03mlx29", "gb00b03mlx29", "gb00b03mlx29", "lu0197800237", "lu0197800237", "nl0000289965", "nl0000301109", "nl0000301109", None, ], "t": [ 233, 234, 235, 233, 234, 235, 180, 181, None, None, None, None, ], "share": [ 0.6, 0.6, 0.6, 0.15, 0.15, 0.15, 0.6, 0.6, 0.25, 1.0, 0.4, 1.0, ], "log_return": [ 0.09604978, -0.06524096, 0.03532373, 0.09604978, -0.06524096, 0.03532373, 0.03025441, 0.036997, None, None, None, None, ], } ) .set_index(["household_id", "asset_id", "t"]) .reindex(columns=["share", "log_return"]) ) result = merge( household.reset_index(), log_return.reset_index(), on=["asset_id"], how="outer", ).set_index(["household_id", "asset_id", "t"]) tm.assert_frame_equal(result, expected) class TestMergeMultiIndexNaN: def test_merge_multiindex_nan_right_index(self): # GH#64492 - merge with right MultiIndex containing NaN used # lev.take(codes) which mapped -1 codes to the last level value # instead of NaN left = pd.DataFrame( {"key1": [1.0, np.nan, 3.0], "key2": ["a", "b", "c"], "val": [10, 20, 30]} ) right = pd.DataFrame( {"data": [100, 200, 300]}, index=pd.MultiIndex.from_arrays( [[1.0, np.nan, 3.0], ["a", "b", "c"]], names=["key1", "key2"] ), ) result = merge(left, right, left_on=["key1", "key2"], right_index=True) expected = pd.DataFrame( { "key1": [1.0, np.nan, 3.0], "key2": ["a", "b", "c"], "val": [10, 20, 30], "data": [100, 200, 300], }, ) tm.assert_frame_equal(result, expected) def test_merge_multiindex_nan_left_index(self): # GH#64492 - same bug as test_merge_multiindex_nan_right_index # but for the left_index=True, right_on=... code path left = pd.DataFrame( {"data": [100, 200, 300]}, index=pd.MultiIndex.from_arrays( [[1.0, np.nan, 3.0], ["a", "b", "c"]], names=["key1", "key2"] ), ) right = pd.DataFrame( {"key1": [1.0, np.nan, 3.0], "key2": ["a", "b", "c"], "val": [10, 20, 30]} ) result = merge(left, right, left_index=True, right_on=["key1", "key2"]) expected = pd.DataFrame( { "data": [100, 200, 300], "key1": [1.0, np.nan, 3.0], "key2": ["a", "b", "c"], "val": [10, 20, 30], }, ) tm.assert_frame_equal(result, expected) class TestJoinMultiMulti: def test_join_multi_multi(self, left_multi, right_multi, join_type, on_cols_multi): left_names = left_multi.index.names right_names = right_multi.index.names if join_type == "right": level_order = right_names + left_names.difference(right_names) else: level_order = left_names + right_names.difference(left_names) # Multi-index join tests expected = ( merge( left_multi.reset_index(), right_multi.reset_index(), how=join_type, on=on_cols_multi, ) .set_index(level_order) .sort_index() ) result = left_multi.join(right_multi, how=join_type).sort_index() tm.assert_frame_equal(result, expected) def test_join_multi_empty_frames( self, left_multi, right_multi, join_type, on_cols_multi ): left_multi = left_multi.drop(columns=left_multi.columns) right_multi = right_multi.drop(columns=right_multi.columns) left_names = left_multi.index.names right_names = right_multi.index.names if join_type == "right": level_order = right_names + left_names.difference(right_names) else: level_order = left_names + right_names.difference(left_names) expected = ( merge( left_multi.reset_index(), right_multi.reset_index(), how=join_type, on=on_cols_multi, ) .set_index(level_order) .sort_index() ) result = left_multi.join(right_multi, how=join_type).sort_index() tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("box", [None, np.asarray, pd.Series, pd.Index]) def test_merge_datetime_index(self, box): # see gh-19038 df = pd.DataFrame( [1, 2, 3], ["2016-01-01", "2017-01-01", "2018-01-01"], columns=["a"] ) df.index = pd.to_datetime(df.index) on_vector = df.index.year if box is not None: on_vector = box(on_vector) exp_years = np.array([2016, 2017, 2018], dtype=np.int32) expected = pd.DataFrame({"a": [1, 2, 3], "key_1": exp_years}) result = df.merge(df, on=["a", on_vector], how="inner") tm.assert_frame_equal(result, expected) expected = pd.DataFrame( {"key_0": exp_years, "a_x": [1, 2, 3], "a_y": [1, 2, 3]} ) result = df.merge(df, on=[df.index.year], how="inner") tm.assert_frame_equal(result, expected) def test_single_common_level(self): index_left = pd.MultiIndex.from_tuples( [("K0", "X0"), ("K0", "X1"), ("K1", "X2")], names=["key", "X"] ) left = pd.DataFrame( {"A": ["A0", "A1", "A2"], "B": ["B0", "B1", "B2"]}, index=index_left ) index_right = pd.MultiIndex.from_tuples( [("K0", "Y0"), ("K1", "Y1"), ("K2", "Y2"), ("K2", "Y3")], names=["key", "Y"] ) right = pd.DataFrame( {"C": ["C0", "C1", "C2", "C3"], "D": ["D0", "D1", "D2", "D3"]}, index=index_right, ) result = left.join(right) expected = merge( left.reset_index(), right.reset_index(), on=["key"], how="inner" ).set_index(["key", "X", "Y"]) tm.assert_frame_equal(result, expected) def test_join_multi_wrong_order(self): # GH 25760 # GH 28956 midx1 = pd.MultiIndex.from_product([[1, 2], [3, 4]], names=["a", "b"]) midx3 = pd.MultiIndex.from_tuples([(4, 1), (3, 2), (3, 1)], names=["b", "a"]) left = pd.DataFrame(index=midx1, data={"x": [10, 20, 30, 40]}) right = pd.DataFrame(index=midx3, data={"y": ["foo", "bar", "fing"]}) result = left.join(right) expected = pd.DataFrame( index=midx1, data={"x": [10, 20, 30, 40], "y": ["fing", "foo", "bar", np.nan]}, ) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("dtype", [np.int64, np.int32, np.uint16]) def test_merge_multi_int_keys_sort(dtype, join_type): # GH#66124 exercise the fused multi-int-key path; result row order with # sort=True must be lexicographic in the key tuples left = pd.DataFrame( { "k1": np.array([3, 1, 3, 2], dtype=dtype), "k2": np.array([10, 20, 5, 20], dtype=dtype), "lval": [0, 1, 2, 3], } ) right = pd.DataFrame( { "k1": np.array([3, 2, 4], dtype=dtype), "k2": np.array([5, 20, 1], dtype=dtype), "rval": [0, 1, 2], } ) result = merge(left, right, on=["k1", "k2"], how=join_type, sort=True) assert result.set_index(["k1", "k2"]).index.is_monotonic_increasing if join_type == "inner": expected = pd.DataFrame( { "k1": np.array([2, 3], dtype=dtype), "k2": np.array([20, 5], dtype=dtype), "lval": [3, 2], "rval": [1, 0], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_int64_extreme_values(): # GH#66124 values near the int64 boundaries within a small span imin = np.iinfo(np.int64).min left = pd.DataFrame( { "k1": np.array([imin, imin + 5, imin + 3], dtype=np.int64), "k2": np.array([-1, -2, -1], dtype=np.int64), "lval": [0, 1, 2], } ) right = pd.DataFrame( { "k1": np.array([imin + 5, imin], dtype=np.int64), "k2": np.array([-2, -1], dtype=np.int64), "rval": [0, 1], } ) result = merge(left, right, on=["k1", "k2"], how="inner") expected = pd.DataFrame( { "k1": np.array([imin, imin + 5], dtype=np.int64), "k2": np.array([-1, -2], dtype=np.int64), "lval": [0, 1], "rval": [1, 0], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_uint64_extreme_values(): # GH#66124 uint64 keys above the int64 boundary umax = np.iinfo(np.uint64).max left = pd.DataFrame( { "k1": np.array([umax, umax - 5, umax - 3], dtype=np.uint64), "k2": np.array([1, 2, 3], dtype=np.uint64), "lval": [0, 1, 2], } ) right = pd.DataFrame( { "k1": np.array([umax - 5, umax], dtype=np.uint64), "k2": np.array([2, 1], dtype=np.uint64), "rval": [0, 1], } ) result = merge(left, right, on=["k1", "k2"], how="inner") expected = pd.DataFrame( { "k1": np.array([umax, umax - 5], dtype=np.uint64), "k2": np.array([1, 2], dtype=np.uint64), "lval": [0, 1], "rval": [1, 0], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_int_keys_wide_span(): # GH#66124 combined key space exceeding int64 falls back to per-column # factorization left = pd.DataFrame( { "k1": np.array([0, 2**62, 5], dtype=np.int64), "k2": np.array([0, 2**61, 7], dtype=np.int64), "lval": [0, 1, 2], } ) right = pd.DataFrame( { "k1": np.array([5, 0, 2**62], dtype=np.int64), "k2": np.array([7, 0, 2**61], dtype=np.int64), "rval": [0, 1, 2], } ) result = merge(left, right, on=["k1", "k2"], how="inner") expected = pd.DataFrame( { "k1": np.array([0, 2**62, 5], dtype=np.int64), "k2": np.array([0, 2**61, 7], dtype=np.int64), "lval": [0, 1, 2], "rval": [1, 2, 0], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_many_keys_int64_overflow(): # GH#58517 with enough high-cardinality key columns the combined key space # overflows int64, so _get_join_keys densifies the keys handled so far and # recurses over the remaining levels n_keys = 14 n_rows = 120 # distinct moduli give each column a cardinality near 100, making the # product of the cardinalities exceed 2**63, while keeping every row's # combination of keys unique keys = {f"k{i}": np.arange(n_rows) % (97 + i) for i in range(n_keys)} left = pd.DataFrame({**keys, "lval": np.arange(n_rows)}) right = left.iloc[::2].drop(columns=["lval"]).copy() right["rval"] = np.arange(len(right)) * 10 result = merge(left, right, on=list(keys), how="inner") expected = left.iloc[::2].reset_index(drop=True) expected["rval"] = np.arange(len(expected)) * 10 tm.assert_frame_equal(result, expected) def test_merge_multi_datetime_int_keys(): # GH#66124 datetime64 plus integer key columns dti = pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-01"]) left = pd.DataFrame( {"k1": dti, "k2": np.array([1, 2, 1], dtype=np.int64), "lval": [0, 1, 2]} ) right = pd.DataFrame( {"k1": dti[:2], "k2": np.array([1, 2], dtype=np.int64), "rval": [0, 1]} ) result = merge(left, right, on=["k1", "k2"], how="inner") expected = pd.DataFrame( { "k1": dti, "k2": np.array([1, 2, 1], dtype=np.int64), "lval": [0, 1, 2], "rval": [0, 1, 0], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_datetime_keys_with_nat(): # GH#66124 NaT keys match each other, same as with a single datetime key dti = pd.to_datetime(["2020-01-01", "NaT", "2020-01-02"]) left = pd.DataFrame( {"k1": dti, "k2": np.array([1, 2, 3], dtype=np.int64), "lval": [0, 1, 2]} ) right = pd.DataFrame( {"k1": dti[:2], "k2": np.array([1, 2], dtype=np.int64), "rval": [0, 1]} ) result = merge(left, right, on=["k1", "k2"], how="inner") expected = pd.DataFrame( { "k1": dti[:2], "k2": np.array([1, 2], dtype=np.int64), "lval": [0, 1], "rval": [0, 1], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_datetime_nat_fused_sort(): # GH#66124 NaT only reaches the fused path when the dates sit near the # datetime64[ns] lower bound (so the NaT<->date span fits in int64); # NaT must still match NaT and sort first, like its i8 value dti = pd.to_datetime(["1680-01-01", "NaT", "1680-01-03", "NaT"]).as_unit("ns") left = pd.DataFrame( {"k1": dti, "k2": np.array([1, 2, 3, 2], dtype=np.int64), "lval": [0, 1, 2, 3]} ) right = pd.DataFrame( {"k1": dti[[1, 0]], "k2": np.array([2, 1], dtype=np.int64), "rval": [0, 1]} ) result = merge(left, right, on=["k1", "k2"], how="inner", sort=True) expected = pd.DataFrame( { "k1": dti[[1, 1, 0]], "k2": np.array([2, 2, 1], dtype=np.int64), "lval": [1, 3, 0], "rval": [0, 0, 1], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_datetimetz_int_keys(): # GH#66124 tz-aware datetime plus integer key columns dti = pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-01"]).tz_localize( "US/Pacific" ) left = pd.DataFrame( {"k1": dti, "k2": np.array([1, 2, 1], dtype=np.int64), "lval": [0, 1, 2]} ) right = pd.DataFrame( {"k1": dti[:2], "k2": np.array([1, 2], dtype=np.int64), "rval": [0, 1]} ) result = merge(left, right, on=["k1", "k2"], how="inner") expected = pd.DataFrame( { "k1": dti, "k2": np.array([1, 2, 1], dtype=np.int64), "lval": [0, 1, 2], "rval": [0, 1, 0], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_timedelta_int_keys(): # GH#66124 timedelta64 plus integer key columns tdi = pd.to_timedelta([1, 2, 1], unit="s") left = pd.DataFrame( {"k1": tdi, "k2": np.array([1, 2, 1], dtype=np.int64), "lval": [0, 1, 2]} ) right = pd.DataFrame( {"k1": tdi[:2], "k2": np.array([1, 2], dtype=np.int64), "rval": [0, 1]} ) result = merge(left, right, on=["k1", "k2"], how="inner") expected = pd.DataFrame( { "k1": tdi, "k2": np.array([1, 2, 1], dtype=np.int64), "lval": [0, 1, 2], "rval": [0, 1, 0], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_mixed_int_dtypes(): # GH#66124 int32/int64 key column pairs promote losslessly left = pd.DataFrame( { "k1": np.array([1, 2, 3, 2], dtype=np.int32), "k2": np.array([1, 1, 2, 2], dtype=np.int64), "lval": [0, 1, 2, 3], } ) right = pd.DataFrame( { "k1": np.array([2, 3], dtype=np.int64), "k2": np.array([1, 2], dtype=np.int32), "rval": [0, 1], } ) result = merge(left, right, on=["k1", "k2"], how="inner") expected = pd.DataFrame( { # the result keys retain the left frame's dtypes "k1": np.array([2, 3], dtype=np.int32), "k2": np.array([1, 2], dtype=np.int64), "lval": [1, 2], "rval": [0, 1], } ) tm.assert_frame_equal(result, expected) def test_merge_multi_bool_int_keys(): # GH#66124 bool key column alongside an integer key column left = pd.DataFrame( { "k1": np.array([True, False, True, False]), "k2": np.array([1, 1, 2, 2], dtype=np.int64), "lval": [0, 1, 2, 3], } ) right = pd.DataFrame( { "k1": np.array([True, False]), "k2": np.array([2, 1], dtype=np.int64), "rval": [0, 1], } ) result = merge(left, right, on=["k1", "k2"], how="inner") expected = pd.DataFrame( { "k1": np.array([False, True]), "k2": np.array([1, 2], dtype=np.int64), "lval": [1, 2], "rval": [1, 0], } ) tm.assert_frame_equal(result, expected) def test_merge_multiindex_with_tuple_valued_level(): # GH#39100 merging on a MultiIndex whose innermost level holds tuples must # align rows correctly instead of mismatching them a_idx = pd.MultiIndex.from_tuples( [ (0, "TV", ("train", "val")), (1, "TV", ("train", "val")), (2, "Test", ("test",)), ], names=["epoch", "stage", "sl"], ) a_df = pd.DataFrame({"a": [0.1, 0.2, 0.3]}, index=a_idx) b_idx = pd.MultiIndex.from_tuples( [(0, "Reval", ("reval",)), (2, "Test", ("test",))], names=["epoch", "stage", "sl"], ) b_df = pd.DataFrame({"b": [0.03, 0.9]}, index=b_idx) result = a_df.merge(b_df, how="outer", left_index=True, right_index=True) expected_idx = pd.MultiIndex.from_tuples( [ (0, "Reval", ("reval",)), (0, "TV", ("train", "val")), (1, "TV", ("train", "val")), (2, "Test", ("test",)), ], names=["epoch", "stage", "sl"], ) expected = pd.DataFrame( {"a": [np.nan, 0.1, 0.2, 0.3], "b": [0.03, np.nan, np.nan, 0.9]}, index=expected_idx, ) tm.assert_frame_equal(result, expected)