import numpy as np import pytest import pandas._testing as tm from pandas.core.indexes.api import ( Index, Int64Index, UInt64Index, ) class TestJoinInt64Index: def test_join_non_unique(self): left = Index([4, 4, 3, 3]) joined, lidx, ridx = left.join(left, return_indexers=True) exp_joined = Index([3, 3, 3, 3, 4, 4, 4, 4]) tm.assert_index_equal(joined, exp_joined) exp_lidx = np.array([2, 2, 3, 3, 0, 0, 1, 1], dtype=np.intp) tm.assert_numpy_array_equal(lidx, exp_lidx) exp_ridx = np.array([2, 3, 2, 3, 0, 1, 0, 1], dtype=np.intp) tm.assert_numpy_array_equal(ridx, exp_ridx) def test_join_inner(self): index = Int64Index(range(0, 20, 2)) other = Int64Index([7, 12, 25, 1, 2, 5]) other_mono = Int64Index([1, 2, 5, 7, 12, 25]) # not monotonic res, lidx, ridx = index.join(other, how="inner", return_indexers=True) # no guarantee of sortedness, so sort for comparison purposes ind = res.argsort() res = res.take(ind) lidx = lidx.take(ind) ridx = ridx.take(ind) eres = Int64Index([2, 12]) elidx = np.array([1, 6], dtype=np.intp) eridx = np.array([4, 1], dtype=np.intp) assert isinstance(res, Int64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) # monotonic res, lidx, ridx = index.join(other_mono, how="inner", return_indexers=True) res2 = index.intersection(other_mono) tm.assert_index_equal(res, res2) elidx = np.array([1, 6], dtype=np.intp) eridx = np.array([1, 4], dtype=np.intp) assert isinstance(res, Int64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) def test_join_left(self): index = Int64Index(range(0, 20, 2)) other = Int64Index([7, 12, 25, 1, 2, 5]) other_mono = Int64Index([1, 2, 5, 7, 12, 25]) # not monotonic res, lidx, ridx = index.join(other, how="left", return_indexers=True) eres = index eridx = np.array([-1, 4, -1, -1, -1, -1, 1, -1, -1, -1], dtype=np.intp) assert isinstance(res, Int64Index) tm.assert_index_equal(res, eres) assert lidx is None tm.assert_numpy_array_equal(ridx, eridx) # monotonic res, lidx, ridx = index.join(other_mono, how="left", return_indexers=True) eridx = np.array([-1, 1, -1, -1, -1, -1, 4, -1, -1, -1], dtype=np.intp) assert isinstance(res, Int64Index) tm.assert_index_equal(res, eres) assert lidx is None tm.assert_numpy_array_equal(ridx, eridx) # non-unique idx = Index([1, 1, 2, 5]) idx2 = Index([1, 2, 5, 7, 9]) res, lidx, ridx = idx2.join(idx, how="left", return_indexers=True) eres = Index([1, 1, 2, 5, 7, 9]) # 1 is in idx2, so it should be x2 eridx = np.array([0, 1, 2, 3, -1, -1], dtype=np.intp) elidx = np.array([0, 0, 1, 2, 3, 4], dtype=np.intp) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) def test_join_right(self): index = Int64Index(range(0, 20, 2)) other = Int64Index([7, 12, 25, 1, 2, 5]) other_mono = Int64Index([1, 2, 5, 7, 12, 25]) # not monotonic res, lidx, ridx = index.join(other, how="right", return_indexers=True) eres = other elidx = np.array([-1, 6, -1, -1, 1, -1], dtype=np.intp) assert isinstance(other, Int64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) assert ridx is None # monotonic res, lidx, ridx = index.join(other_mono, how="right", return_indexers=True) eres = other_mono elidx = np.array([-1, 1, -1, -1, 6, -1], dtype=np.intp) assert isinstance(other, Int64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) assert ridx is None # non-unique idx = Index([1, 1, 2, 5]) idx2 = Index([1, 2, 5, 7, 9]) res, lidx, ridx = idx.join(idx2, how="right", return_indexers=True) eres = Index([1, 1, 2, 5, 7, 9]) # 1 is in idx2, so it should be x2 elidx = np.array([0, 1, 2, 3, -1, -1], dtype=np.intp) eridx = np.array([0, 0, 1, 2, 3, 4], dtype=np.intp) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) def test_join_non_int_index(self): index = Int64Index(range(0, 20, 2)) other = Index([3, 6, 7, 8, 10], dtype=object) outer = index.join(other, how="outer") outer2 = other.join(index, how="outer") expected = Index([0, 2, 3, 4, 6, 7, 8, 10, 12, 14, 16, 18]) tm.assert_index_equal(outer, outer2) tm.assert_index_equal(outer, expected) inner = index.join(other, how="inner") inner2 = other.join(index, how="inner") expected = Index([6, 8, 10]) tm.assert_index_equal(inner, inner2) tm.assert_index_equal(inner, expected) left = index.join(other, how="left") tm.assert_index_equal(left, index.astype(object)) left2 = other.join(index, how="left") tm.assert_index_equal(left2, other) right = index.join(other, how="right") tm.assert_index_equal(right, other) right2 = other.join(index, how="right") tm.assert_index_equal(right2, index.astype(object)) def test_join_outer(self): index = Int64Index(range(0, 20, 2)) other = Int64Index([7, 12, 25, 1, 2, 5]) other_mono = Int64Index([1, 2, 5, 7, 12, 25]) # not monotonic # guarantee of sortedness res, lidx, ridx = index.join(other, how="outer", return_indexers=True) noidx_res = index.join(other, how="outer") tm.assert_index_equal(res, noidx_res) eres = Int64Index([0, 1, 2, 4, 5, 6, 7, 8, 10, 12, 14, 16, 18, 25]) elidx = np.array([0, -1, 1, 2, -1, 3, -1, 4, 5, 6, 7, 8, 9, -1], dtype=np.intp) eridx = np.array( [-1, 3, 4, -1, 5, -1, 0, -1, -1, 1, -1, -1, -1, 2], dtype=np.intp ) assert isinstance(res, Int64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) # monotonic res, lidx, ridx = index.join(other_mono, how="outer", return_indexers=True) noidx_res = index.join(other_mono, how="outer") tm.assert_index_equal(res, noidx_res) elidx = np.array([0, -1, 1, 2, -1, 3, -1, 4, 5, 6, 7, 8, 9, -1], dtype=np.intp) eridx = np.array( [-1, 0, 1, -1, 2, -1, 3, -1, -1, 4, -1, -1, -1, 5], dtype=np.intp ) assert isinstance(res, Int64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) class TestJoinUInt64Index: @pytest.fixture def index_large(self): # large values used in TestUInt64Index where no compat needed with Int64/Float64 large = [2**63, 2**63 + 10, 2**63 + 15, 2**63 + 20, 2**63 + 25] return UInt64Index(large) def test_join_inner(self, index_large): other = UInt64Index(2**63 + np.array([7, 12, 25, 1, 2, 10], dtype="uint64")) other_mono = UInt64Index( 2**63 + np.array([1, 2, 7, 10, 12, 25], dtype="uint64") ) # not monotonic res, lidx, ridx = index_large.join(other, how="inner", return_indexers=True) # no guarantee of sortedness, so sort for comparison purposes ind = res.argsort() res = res.take(ind) lidx = lidx.take(ind) ridx = ridx.take(ind) eres = UInt64Index(2**63 + np.array([10, 25], dtype="uint64")) elidx = np.array([1, 4], dtype=np.intp) eridx = np.array([5, 2], dtype=np.intp) assert isinstance(res, UInt64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) # monotonic res, lidx, ridx = index_large.join( other_mono, how="inner", return_indexers=True ) res2 = index_large.intersection(other_mono) tm.assert_index_equal(res, res2) elidx = np.array([1, 4], dtype=np.intp) eridx = np.array([3, 5], dtype=np.intp) assert isinstance(res, UInt64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) def test_join_left(self, index_large): other = UInt64Index(2**63 + np.array([7, 12, 25, 1, 2, 10], dtype="uint64")) other_mono = UInt64Index( 2**63 + np.array([1, 2, 7, 10, 12, 25], dtype="uint64") ) # not monotonic res, lidx, ridx = index_large.join(other, how="left", return_indexers=True) eres = index_large eridx = np.array([-1, 5, -1, -1, 2], dtype=np.intp) assert isinstance(res, UInt64Index) tm.assert_index_equal(res, eres) assert lidx is None tm.assert_numpy_array_equal(ridx, eridx) # monotonic res, lidx, ridx = index_large.join(other_mono, how="left", return_indexers=True) eridx = np.array([-1, 3, -1, -1, 5], dtype=np.intp) assert isinstance(res, UInt64Index) tm.assert_index_equal(res, eres) assert lidx is None tm.assert_numpy_array_equal(ridx, eridx) # non-unique idx = UInt64Index(2**63 + np.array([1, 1, 2, 5], dtype="uint64")) idx2 = UInt64Index(2**63 + np.array([1, 2, 5, 7, 9], dtype="uint64")) res, lidx, ridx = idx2.join(idx, how="left", return_indexers=True) # 1 is in idx2, so it should be x2 eres = UInt64Index(2**63 + np.array([1, 1, 2, 5, 7, 9], dtype="uint64")) eridx = np.array([0, 1, 2, 3, -1, -1], dtype=np.intp) elidx = np.array([0, 0, 1, 2, 3, 4], dtype=np.intp) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) def test_join_right(self, index_large): other = UInt64Index(2**63 + np.array([7, 12, 25, 1, 2, 10], dtype="uint64")) other_mono = UInt64Index( 2**63 + np.array([1, 2, 7, 10, 12, 25], dtype="uint64") ) # not monotonic res, lidx, ridx = index_large.join(other, how="right", return_indexers=True) eres = other elidx = np.array([-1, -1, 4, -1, -1, 1], dtype=np.intp) tm.assert_numpy_array_equal(lidx, elidx) assert isinstance(other, UInt64Index) tm.assert_index_equal(res, eres) assert ridx is None # monotonic res, lidx, ridx = index_large.join( other_mono, how="right", return_indexers=True ) eres = other_mono elidx = np.array([-1, -1, -1, 1, -1, 4], dtype=np.intp) assert isinstance(other, UInt64Index) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_index_equal(res, eres) assert ridx is None # non-unique idx = UInt64Index(2**63 + np.array([1, 1, 2, 5], dtype="uint64")) idx2 = UInt64Index(2**63 + np.array([1, 2, 5, 7, 9], dtype="uint64")) res, lidx, ridx = idx.join(idx2, how="right", return_indexers=True) # 1 is in idx2, so it should be x2 eres = UInt64Index(2**63 + np.array([1, 1, 2, 5, 7, 9], dtype="uint64")) elidx = np.array([0, 1, 2, 3, -1, -1], dtype=np.intp) eridx = np.array([0, 0, 1, 2, 3, 4], dtype=np.intp) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) def test_join_non_int_index(self, index_large): other = Index( 2**63 + np.array([1, 5, 7, 10, 20], dtype="uint64"), dtype=object ) outer = index_large.join(other, how="outer") outer2 = other.join(index_large, how="outer") expected = Index( 2**63 + np.array([0, 1, 5, 7, 10, 15, 20, 25], dtype="uint64") ) tm.assert_index_equal(outer, outer2) tm.assert_index_equal(outer, expected) inner = index_large.join(other, how="inner") inner2 = other.join(index_large, how="inner") expected = Index(2**63 + np.array([10, 20], dtype="uint64")) tm.assert_index_equal(inner, inner2) tm.assert_index_equal(inner, expected) left = index_large.join(other, how="left") tm.assert_index_equal(left, index_large.astype(object)) left2 = other.join(index_large, how="left") tm.assert_index_equal(left2, other) right = index_large.join(other, how="right") tm.assert_index_equal(right, other) right2 = other.join(index_large, how="right") tm.assert_index_equal(right2, index_large.astype(object)) def test_join_outer(self, index_large): other = UInt64Index(2**63 + np.array([7, 12, 25, 1, 2, 10], dtype="uint64")) other_mono = UInt64Index( 2**63 + np.array([1, 2, 7, 10, 12, 25], dtype="uint64") ) # not monotonic # guarantee of sortedness res, lidx, ridx = index_large.join(other, how="outer", return_indexers=True) noidx_res = index_large.join(other, how="outer") tm.assert_index_equal(res, noidx_res) eres = UInt64Index( 2**63 + np.array([0, 1, 2, 7, 10, 12, 15, 20, 25], dtype="uint64") ) elidx = np.array([0, -1, -1, -1, 1, -1, 2, 3, 4], dtype=np.intp) eridx = np.array([-1, 3, 4, 0, 5, 1, -1, -1, 2], dtype=np.intp) assert isinstance(res, UInt64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx) # monotonic res, lidx, ridx = index_large.join( other_mono, how="outer", return_indexers=True ) noidx_res = index_large.join(other_mono, how="outer") tm.assert_index_equal(res, noidx_res) elidx = np.array([0, -1, -1, -1, 1, -1, 2, 3, 4], dtype=np.intp) eridx = np.array([-1, 0, 1, 2, 3, 4, -1, -1, 5], dtype=np.intp) assert isinstance(res, UInt64Index) tm.assert_index_equal(res, eres) tm.assert_numpy_array_equal(lidx, elidx) tm.assert_numpy_array_equal(ridx, eridx)