|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | + |
| 3 | +import pytest |
| 4 | +import numpy as np |
| 5 | + |
| 6 | +from pandas import DataFrame, Index |
| 7 | +from pandas.tests.frame.common import TestData |
| 8 | +import pandas.util.testing as tm |
| 9 | + |
| 10 | + |
| 11 | +@pytest.fixture |
| 12 | +def frame(): |
| 13 | + return TestData().frame |
| 14 | + |
| 15 | + |
| 16 | +@pytest.fixture |
| 17 | +def left(): |
| 18 | + return DataFrame({'a': [20, 10, 0]}, index=[2, 1, 0]) |
| 19 | + |
| 20 | + |
| 21 | +@pytest.fixture |
| 22 | +def right(): |
| 23 | + return DataFrame({'b': [300, 100, 200]}, index=[3, 1, 2]) |
| 24 | + |
| 25 | + |
| 26 | +@pytest.mark.parametrize( |
| 27 | + "how, sort, expected", |
| 28 | + [('inner', False, DataFrame({'a': [20, 10], |
| 29 | + 'b': [200, 100]}, |
| 30 | + index=[2, 1])), |
| 31 | + ('inner', True, DataFrame({'a': [10, 20], |
| 32 | + 'b': [100, 200]}, |
| 33 | + index=[1, 2])), |
| 34 | + ('left', False, DataFrame({'a': [20, 10, 0], |
| 35 | + 'b': [200, 100, np.nan]}, |
| 36 | + index=[2, 1, 0])), |
| 37 | + ('left', True, DataFrame({'a': [0, 10, 20], |
| 38 | + 'b': [np.nan, 100, 200]}, |
| 39 | + index=[0, 1, 2])), |
| 40 | + ('right', False, DataFrame({'a': [np.nan, 10, 20], |
| 41 | + 'b': [300, 100, 200]}, |
| 42 | + index=[3, 1, 2])), |
| 43 | + ('right', True, DataFrame({'a': [10, 20, np.nan], |
| 44 | + 'b': [100, 200, 300]}, |
| 45 | + index=[1, 2, 3])), |
| 46 | + ('outer', False, DataFrame({'a': [0, 10, 20, np.nan], |
| 47 | + 'b': [np.nan, 100, 200, 300]}, |
| 48 | + index=[0, 1, 2, 3])), |
| 49 | + ('outer', True, DataFrame({'a': [0, 10, 20, np.nan], |
| 50 | + 'b': [np.nan, 100, 200, 300]}, |
| 51 | + index=[0, 1, 2, 3]))]) |
| 52 | +def test_join(left, right, how, sort, expected): |
| 53 | + |
| 54 | + result = left.join(right, how=how, sort=sort) |
| 55 | + tm.assert_frame_equal(result, expected) |
| 56 | + |
| 57 | + |
| 58 | +def test_join_index(frame): |
| 59 | + # left / right |
| 60 | + |
| 61 | + f = frame.loc[frame.index[:10], ['A', 'B']] |
| 62 | + f2 = frame.loc[frame.index[5:], ['C', 'D']].iloc[::-1] |
| 63 | + |
| 64 | + joined = f.join(f2) |
| 65 | + tm.assert_index_equal(f.index, joined.index) |
| 66 | + expected_columns = Index(['A', 'B', 'C', 'D']) |
| 67 | + tm.assert_index_equal(joined.columns, expected_columns) |
| 68 | + |
| 69 | + joined = f.join(f2, how='left') |
| 70 | + tm.assert_index_equal(joined.index, f.index) |
| 71 | + tm.assert_index_equal(joined.columns, expected_columns) |
| 72 | + |
| 73 | + joined = f.join(f2, how='right') |
| 74 | + tm.assert_index_equal(joined.index, f2.index) |
| 75 | + tm.assert_index_equal(joined.columns, expected_columns) |
| 76 | + |
| 77 | + # inner |
| 78 | + |
| 79 | + joined = f.join(f2, how='inner') |
| 80 | + tm.assert_index_equal(joined.index, f.index[5:10]) |
| 81 | + tm.assert_index_equal(joined.columns, expected_columns) |
| 82 | + |
| 83 | + # outer |
| 84 | + |
| 85 | + joined = f.join(f2, how='outer') |
| 86 | + tm.assert_index_equal(joined.index, frame.index.sort_values()) |
| 87 | + tm.assert_index_equal(joined.columns, expected_columns) |
| 88 | + |
| 89 | + tm.assertRaisesRegexp(ValueError, 'join method', f.join, f2, how='foo') |
| 90 | + |
| 91 | + # corner case - overlapping columns |
| 92 | + for how in ('outer', 'left', 'inner'): |
| 93 | + with tm.assertRaisesRegexp(ValueError, 'columns overlap but ' |
| 94 | + 'no suffix'): |
| 95 | + frame.join(frame, how=how) |
| 96 | + |
| 97 | + |
| 98 | +def test_join_index_more(frame): |
| 99 | + af = frame.loc[:, ['A', 'B']] |
| 100 | + bf = frame.loc[::2, ['C', 'D']] |
| 101 | + |
| 102 | + expected = af.copy() |
| 103 | + expected['C'] = frame['C'][::2] |
| 104 | + expected['D'] = frame['D'][::2] |
| 105 | + |
| 106 | + result = af.join(bf) |
| 107 | + tm.assert_frame_equal(result, expected) |
| 108 | + |
| 109 | + result = af.join(bf, how='right') |
| 110 | + tm.assert_frame_equal(result, expected[::2]) |
| 111 | + |
| 112 | + result = bf.join(af, how='right') |
| 113 | + tm.assert_frame_equal(result, expected.loc[:, result.columns]) |
| 114 | + |
| 115 | + |
| 116 | +def test_join_index_series(frame): |
| 117 | + df = frame.copy() |
| 118 | + s = df.pop(frame.columns[-1]) |
| 119 | + joined = df.join(s) |
| 120 | + |
| 121 | + # TODO should this check_names ? |
| 122 | + tm.assert_frame_equal(joined, frame, check_names=False) |
| 123 | + |
| 124 | + s.name = None |
| 125 | + tm.assertRaisesRegexp(ValueError, 'must have a name', df.join, s) |
| 126 | + |
| 127 | + |
| 128 | +def test_join_overlap(frame): |
| 129 | + df1 = frame.loc[:, ['A', 'B', 'C']] |
| 130 | + df2 = frame.loc[:, ['B', 'C', 'D']] |
| 131 | + |
| 132 | + joined = df1.join(df2, lsuffix='_df1', rsuffix='_df2') |
| 133 | + df1_suf = df1.loc[:, ['B', 'C']].add_suffix('_df1') |
| 134 | + df2_suf = df2.loc[:, ['B', 'C']].add_suffix('_df2') |
| 135 | + |
| 136 | + no_overlap = frame.loc[:, ['A', 'D']] |
| 137 | + expected = df1_suf.join(df2_suf).join(no_overlap) |
| 138 | + |
| 139 | + # column order not necessarily sorted |
| 140 | + tm.assert_frame_equal(joined, expected.loc[:, joined.columns]) |
0 commit comments