Join the collection using SideInput

Question

Trying to join two Pcollection using SideInput transform. In the ParDo function while mapping the value, from the sideinput collection we may get the multiple mapping records as a collection. In such a case how to handle the collection and how to return those collection of values to the PCollection.

It would be good if some one help to solve this case. Here is the code snippet that I tried.

PCollection<TableRow> pc1 = ...;
PCollection<Row> pc1Rows = pc1.apply(
    ParDo.of(new fnConvertTableRowToRow())).setRowSchema(schemaPc1);
PCollection<KV<Integer, Row>> keyed_pc1Rows = pc1Rows.apply(
    WithKeys.of(new SerializableFunction<Row, Integer>() {
       public Integer apply(Row s) {
         return Integer.parseInt(s.getValue("LOCATION_ID").toString());
       }
    }));

PCollection<TableRow> pc2 = ...;

PCollection<Row> pc2Rows = pc2.apply(
    ParDo.of(new fnConvertTableRowToRow())).setRowSchema(schemaPc2);

PCollection<KV<Integer, Iterable<Row>>> keywordGroups = pc2Rows.apply(
    new fnGroupKeyWords());

PCollectionView<Map<Integer, Iterable<Row>>> sideInputView =
    keywordGroups.apply("Side Input",
       View.<Integer, Iterable<Row>>asMap());

PCollection<Row> finalResultCollection = keyed_pc1Rows.apply("Process",
  ParDo.of(new DoFn<KV<Integer,Row>, Row>() {

    @ProcessElement
    public void processElement(ProcessContext c) {

      Integer key = Integer.parseInt(c.element().getKey().toString());

      Row leftRow = c.element().getValue();

      Map<Integer, Iterable<Row>> key2Rows = c.sideInput(sideInputView);

      Iterable<Row> rightRowsIterable = key2Rows.get(key);

      for (Iterator<Row> i = rightRowsIterable.iterator(); i.hasNext(); ) {
        Row suit = (Row) i.next();
        Row targetRow = Row.withSchema(schemaOutput)
                           .addValues(leftRow.getValues())
                           .addValues(suit.getValues())
                           .build();
        c.output(targetRow);
      }
    }
}).withSideInputs(sideInputView));

public static class fnGroupKeyWords extends
  PTransform<PCollection<Row>, PCollection<KV<Integer, Iterable<Row>>>> {

  @Override
  public PCollection<KV<Integer, Iterable<Row>>> expand(
    PCollection<Row> rows) {

      PCollection<KV<Integer, Row>> kvs = rows.apply(
          ParDo.of(new TransferKeyValueFn()));
      PCollection<KV<Integer, Iterable<Row>>> group = kvs.apply(
          GroupByKey.<Integer, Row> create());
      return group;
  }
}

public static class TransferKeyValueFn extends
  DoFn<Row, KV<Integer, Row>> {

  @ProcessElement
  public void processElement(ProcessContext c) throws ParseException {
    Row tRow = c.element();

    c.output(
       KV.of(
          Integer.parseInt(tRow.getValue("DW_LOCATION_ID").toString()),
          tRow));
  }
}

Answer 1

If you wish to join two PCollections together using a common key. the CoGroupByKey might make more sense. Please consider this approach instead of side inputs

Also this blog post has a great explanation as well.

Answer 2

I think using the SideInput suggestion would perform well if you have a very small collection which could fit into memory. You could use it as a side input with view.asMultimap. Then in a ParDo processing the larger PCollection (After a GBK, to give you an iterable over all elements for the key), lookup the key you are interested in from the side input. Here is an example test pipeline using a multimap pcollection.

However, if your collection is quite large then using Flatten to combine both pcollections together would be a better approach. Then using a GroupByKey afterward, which will give you an iterable for element under the same key. This will still be processed sequentially. Though, I believe you will will have issues with performance, unless you eliminate the hot key. Please see the explanation of using combiners to alleviate this.