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