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)
