Someone has asked for a Java-based solution for this question. Here is the Java code for this. It's more verbose, but it does essentially the same thing.
// First we want to load all departments, and put them into a PCollection
// of key-value pairs, where the Key is their identifier. We assume that it is String-type.
PCollection<KV<String, Department>> departments =
p.apply(new LoadDepts())
.apply("getKey", MapElements.via((Department dept) -> KV.of(dept.getId(), dept)));
// We then convert this PCollection into a map-type PCollectionView.
// We can access this map directly within a ParDo.
PCollectionView<Map<String, Department>> departmentSideInput =
departments.apply("ToMapSideInput", View.<String, Department>asMap());
// We load the PCollection of employees
PCollection<Employee> employees = p.apply(new LoadEmployees());
// Let us suppose that we will *extend* an employee information with their
// Department information. I have assumed the existence of an ExtendedEmployee
// class to represent an employee extended with department information.
class JoinDeptEmployeeDoFn extends DoFn<Employee, ExtendedEmployee> {
@ProcessElement
public void processElement(ProcessContext c) {
// We obtain the Map-type side input with department information.
Map<String, Department> departmentMap = c.sideInput(departmentSideInput);
Employee empl = c.element();
Department dept = departmentMap.get(empl.getDepartmentId(), null);
if (department == null) return;
ExtendedEmployee result = empl.extendWith(dept);
c.output(result);
}
}
// We apply the ParDo to extend the employee with department information
// and specify that it takes in a departmentSideInput.
PCollection<ExtendedEmployee> extendedEmployees =
employees.apply(
ParDo.of(new JoinDeptEmployeeDoFn()).withSideInput(departmentSideInput));
With CoGroupByKey, you can use dept_id as a key to group both collections. The way this looks in Beam Java SDK is a CoGbkResult
.
// We load the departments, and make them a key-value collection, to Join them
// later with employees.
PCollection<KV<String, Department>> departments =
p.apply(new LoadDepts())
.apply("getKey", MapElements.via((Department dept) -> KV.of(dept.getId(), dept)));
// Because we will perform a join, employees also need to be put into
// key-value pairs, where their key is their *department id*.
PCollection<KV<String, Employee>> employees =
p.apply(new LoadEmployees())
.apply("getKey", MapElements.via((Employee empl) -> KV.of(empl.getDepartmentId(), empl)));
// We define a DoFn that is able to join a single department with multiple
// employees.
class JoinEmployeesWithDepartments extends DoFn<KV<String, CoGbkResult>, ExtendedEmployee> {
@ProcessElement
public void processElement(ProcessContext c) {
KV<...> elm = c.element();
// We assume one department with the same ID, and assume that
// employees always have a department available.
Department dept = elm.getValue().getOnly(departmentsTag);
Iterable<Employee> employees = elm.getValue().getAll(employeesTag);
for (Employee empl : employees) {
ExtendedEmployee result = empl.extendWith(dept);
c.output(result);
}
}
}
// The syntax for a CoGroupByKey operation is a bit verbose.
// In this step we define a TupleTag, which serves as identifier for a
// PCollection.
final TupleTag<String> employeesTag = new TupleTag<>();
final TupleTag<String> departmentsTag = new TupleTag<>();
// We use the PCollection tuple-tags to join the two PCollections.
PCollection<KV<String, CoGbkResult>> results =
KeyedPCollectionTuple.of(departmentsTag, departments)
.and(employeesTag, employees)
.apply(CoGroupByKey.create());
// Finally, we convert the joined PCollections into a kind that
// we can use: ExtendedEmployee.
PCollection<ExtendedEmployee> extendedEmployees =
results.apply("ExtendInformation", ParDo.of(new JoinEmployeesWithDepartments()));