def test_join_multiindex(self): df1 = pdw.DataFrame( {'col1': np.arange(8)}, pdw.MultiIndex.from_product( [np.array([1, 2]), np.array([3, 4]), np.array([5, 6])], ['i1', 'i2', 'i3'])) df2 = pdw.DataFrame( {'col2': np.arange(12)}, pdw.MultiIndex.from_product( [np.array([1, 2, 3]), np.array([3, 5]), np.array([5, 6])], ['i1', 'i2', 'i3'])) result = df1.merge(df2) expected_result = pdw.DataFrame( { 'col1': np.array([0, 1, 4, 5]), 'col2': np.array([0, 1, 4, 5]) }, pdw.MultiIndex( [np.array([1, 2]), np.array([3, 4]), np.array([5, 6])], [ np.array([0, 0, 1, 1]), np.array([0, 0, 0, 0]), np.array([0, 1, 0, 1]) ], ['i1', 'i2', 'i3'])) test_equal_multiindex(expected_result.index, result.index) test_equal_series(expected_result['col1'], result['col1']) test_equal_series(expected_result['col2'], result['col2'])
def test_drop_list(self): data = {} index = pdw.MultiIndex.from_product( [np.array([1, 2]), np.array([3, 4])], ['a', 'b']) expected_result = pdw.DataFrame(data, index) result = self.df.drop(['col1', 'col2']) self.assertListEqual(expected_result.data.keys(), result.data.keys()) test_equal_multiindex(expected_result.index, result.index)
def test_drop_str(self): data = {'col2': np.array([5., 6., 7., 8.])} index = pdw.MultiIndex.from_product( [np.array([1, 2]), np.array([3, 4])], ['a', 'b']) expected_result = pdw.DataFrame(data, index) result = self.df.drop('col1') self.assertListEqual(expected_result.data.keys(), result.data.keys()) np.testing.assert_array_equal( evaluate_if_necessary(expected_result['col2']), evaluate_if_necessary(result['col2'])) test_equal_multiindex(expected_result.index, result.index)
def test_getitem_series(self): data = {'col1': np.array([1, 2]), 'col2': np.array([5., 6.])} index = pdw.MultiIndex( [np.array([1, 2]), np.array([3, 4])], [np.array([0, 0]), np.array([0, 1])], ['a', 'b']) expected_result = pdw.DataFrame(data, index) result = self.df[self.df['col1'] < 3] np.testing.assert_array_equal( evaluate_if_necessary(expected_result['col1']), evaluate_if_necessary(result['col1'])) np.testing.assert_array_equal( evaluate_if_necessary(expected_result['col2']), evaluate_if_necessary(result['col2'])) test_equal_multiindex(expected_result.index, result.index)
def test_getitem_list(self): data = { 'col1': np.array([1, 2, 3, 4]), 'col2': np.array([5., 6., 7., 8.]) } index = pdw.MultiIndex.from_product( [np.array([1, 2]), np.array([3, 4])], ['a', 'b']) expected_result = pdw.DataFrame(data, index) result = self.df[['col1', 'col2']] np.testing.assert_array_equal( evaluate_if_necessary(expected_result['col1']), evaluate_if_necessary(result['col1'])) np.testing.assert_array_equal( evaluate_if_necessary(expected_result['col2']), evaluate_if_necessary(result['col2'])) test_equal_multiindex(expected_result.index, result.index)
def test_element_wise_operation(self): expected_data = { 'col1': np.array([2, 4, 6, 8]), 'col2': np.array([10, 12, 14, 16]) } expected_index = pdw.MultiIndex.from_product( [np.array([1, 2]), np.array([3, 4])], ['a', 'b']) expected_result = pdw.DataFrame(expected_data, expected_index) data = {'col1': np.array([1, 2, 3, 4]), 'col2': np.array([5, 6, 7, 8])} index = pdw.MultiIndex.from_product( [np.array([1, 2]), np.array([3, 4])], ['a', 'b']) result = pdw.DataFrame(data, index) * 2 np.testing.assert_array_equal( evaluate_if_necessary(expected_result['col1']), evaluate_if_necessary(result['col1'])) np.testing.assert_array_equal( evaluate_if_necessary(expected_result['col2']), evaluate_if_necessary(result['col2'])) test_equal_multiindex(expected_result.index, result.index)