def test_cogroup_multiple_keys(self):
a = MockRDD.of(('k1', 0), ('k2', 2)).persist()
b = MockRDD.of(('k1', 1), ('k3', 3)).persist()
def convert(x):
k, (i0, i1) = x
return k, (tuple(i0), tuple(i1))
col_eq(a.cogroup(b).collect(), [('k1', ([0], [1])), ('k2', ([2], [])),
('k3', ([], [3]))],
element_converter=convert)
col_eq(b.cogroup(a).collect(), [('k1', ([1], [0])), ('k2', ([], [2])),
('k3', ([3], []))],
element_converter=convert)
def test_cogroup_match(self):
zero = MockRDD.of(('k', 0)).persist()
one = MockRDD.of(('k', 1)).persist()
self.assertEqual(zero.cogroup(one).collect(), [('k', ([0], [1]))])
self.assertEqual(one.cogroup(zero).collect(), [('k', ([1], [0]))])
def test_cogroup_only_one(self):
empty = MockRDD.empty().persist()
one = MockRDD.of(('k', 1)).persist()
self.assertEqual(empty.cogroup(one).collect(), [('k', ([], [1]))])
self.assertEqual(one.cogroup(empty).collect(), [('k', ([1], []))])
