Apache Beam - join two collections with unequal numer of rows - apache-beam

Question

I have two collection with unequal number of elements. Lets say A and B. I need to split A and B into groups and randomly assign rows inside groups.
Assumptions:
A and B are split into equal number of groups.
Number of elements in each group derived from B is 10 times higher than in group derived from A.
Element from group B can be assigned to only one element from group A (or not matched at all).
Elements between groups should be assigned randomly
My current code looks like this:
def join(grouped_ab):
(key, (rows_a, rows_b)) = grouped_ab
sample = random.sample(rows_b, len(rows_a))
for row_a, row_b in zip(rows_a, sample):
yield {
"a": row_a,
"b": row_b
}
a = p | "read a" >> ReadFromText("a.csv") \
| "split a" >> beam.ParDo(Split())
b = p | "read b" >> ReadFromText("b.csv") \
| "split b" >> beam.ParDo(Split())
(a, b) | beam.CoGroupByKey() \
| beam.ParDo(join) \
| WriteToText("out")
Above code works fine but the problem is that distibution of number of elements in groups are unequal. There are hot groups which can contain up to 30% of all elements.
How I can improve my code?

Answer 1

Looking at the problem statement it looks like B has many more rows compared to A. Also, all the rows of A should be present along with each group of B so that the lookups can happen locally. In the case of CoGroupByKey there is a shuffle that is involved.
For this usecase it would be better to use Side Inputs. The idea of Side input is that when joining two collections if one of the collections is small enough to fit in memory then the smaller dataset is broadcasted to all the workers so that when the lookup happens they are done locally thereby speeding up performance.
Details can be found here - https://beam.apache.org/documentation/programming-guide/#side-inputs
Thanks,
Jayadeep