CoGroup { public static Impl join(By clause) { return new Impl(new JoinArguments(clause)); }  static Impl join(By clause); static Impl join(String tag, By clause);  }
@Test @Category(NeedsRunner.class) public void testCoGroupByFieldNames() { PCollection<Row> pc1 = pipeline .apply( "Create1", Create.of( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 5, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 6, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build())) .setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline .apply( "Create2", Create.of( Row.withSchema(CG_SCHEMA_1).addValues("user1", 9, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 10, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 11, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 12, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 13, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 14, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 15, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 16, "ar").build())) .setRowSchema(CG_SCHEMA_1); PCollection<Row> pc3 = pipeline .apply( "Create3", Create.of( Row.withSchema(CG_SCHEMA_1).addValues("user1", 17, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 18, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 19, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 20, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 21, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 22, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 23, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 24, "ar").build())) .setRowSchema(CG_SCHEMA_1); Schema expectedSchema = Schema.builder() .addRowField("key", SIMPLE_CG_KEY_SCHEMA) .addIterableField("pc1", FieldType.row(CG_SCHEMA_1)) .addIterableField("pc2", FieldType.row(CG_SCHEMA_1)) .addIterableField("pc3", FieldType.row(CG_SCHEMA_1)) .build(); Row key1Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user1", "us").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 9, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 10, "us").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 17, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 18, "us").build())) .build(); Row key2Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user1", "il").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 11, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 12, "il").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 19, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 20, "il").build())) .build(); Row key3Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user2", "fr").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 5, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 6, "fr").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 13, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 14, "fr").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 21, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 22, "fr").build())) .build(); Row key4Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user2", "ar").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 15, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 16, "ar").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 23, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 24, "ar").build())) .build(); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply("CoGroup", CoGroup.join(By.fieldNames("user", "country"))); List<Row> expected = ImmutableList.of(key1Joined, key2Joined, key3Joined, key4Joined); PAssert.that(joined).satisfies(actual -> containsJoinedFields(expected, actual)); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testCoGroupByDifferentFields() { PCollection<Row> pc1 = pipeline .apply( "Create1", Create.of( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 5, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 6, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build())) .setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline .apply( "Create2", Create.of( Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 10, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 11, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 12, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 13, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 14, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 15, "ar").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 16, "ar").build())) .setRowSchema(CG_SCHEMA_2); PCollection<Row> pc3 = pipeline .apply( "Create3", Create.of( Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 18, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 19, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 20, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 21, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 22, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 23, "ar").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 24, "ar").build())) .setRowSchema(CG_SCHEMA_3); Schema expectedSchema = Schema.builder() .addRowField("key", SIMPLE_CG_KEY_SCHEMA) .addIterableField("pc1", FieldType.row(CG_SCHEMA_1)) .addIterableField("pc2", FieldType.row(CG_SCHEMA_2)) .addIterableField("pc3", FieldType.row(CG_SCHEMA_3)) .build(); Row key1Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user1", "us").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 10, "us").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 18, "us").build())) .build(); Row key2Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user1", "il").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user1", 11, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 12, "il").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user1", 19, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 20, "il").build())) .build(); Row key3Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user2", "fr").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 5, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 6, "fr").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user2", 13, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 14, "fr").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user2", 21, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 22, "fr").build())) .build(); Row key4Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user2", "ar").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user2", 15, "ar").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 16, "ar").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user2", 23, "ar").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 24, "ar").build())) .build(); PCollection<Row> joined1 = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup1", CoGroup.join("pc1", By.fieldNames("user", "country")) .join("pc2", By.fieldNames("user2", "country2")) .join("pc3", By.fieldNames("user3", "country3"))); PCollection<Row> joined2 = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup2", CoGroup.join("pc1", By.fieldNames("user", "country")) .join("pc2", By.fieldNames("user2", "country2").withSideInput()) .join("pc3", By.fieldNames("user3", "country3"))); PCollection<Row> joined3 = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup3", CoGroup.join("pc1", By.fieldNames("user", "country")) .join("pc2", By.fieldNames("user2", "country2").withSideInput()) .join("pc3", By.fieldNames("user3", "country3").withSideInput())); List<Row> expected = ImmutableList.of(key1Joined, key2Joined, key3Joined, key4Joined); PAssert.that(joined1).satisfies(actual -> containsJoinedFields(expected, actual)); PAssert.that(joined2).satisfies(actual -> containsJoinedFields(expected, actual)); PAssert.that(joined3).satisfies(actual -> containsJoinedFields(expected, actual)); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testNoMainInput() { PCollection<Row> pc1 = pipeline .apply( "Create1", Create.of(Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build())) .setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline .apply( "Create2", Create.of(Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build())) .setRowSchema(CG_SCHEMA_2); PCollection<Row> pc3 = pipeline .apply( "Create3", Create.of(Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build())) .setRowSchema(CG_SCHEMA_3); thrown.expect(IllegalArgumentException.class); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup1", CoGroup.join("pc1", By.fieldNames("user", "country").withSideInput()) .join("pc2", By.fieldNames("user2", "country2").withSideInput()) .join("pc3", By.fieldNames("user3", "country3").withSideInput())); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testIllegalOuterJoinWithSideInput() { PCollection<Row> pc1 = pipeline .apply( "Create1", Create.of(Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build())) .setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline .apply( "Create2", Create.of(Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build())) .setRowSchema(CG_SCHEMA_2); PCollection<Row> pc3 = pipeline .apply( "Create3", Create.of(Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build())) .setRowSchema(CG_SCHEMA_3); thrown.expect(IllegalArgumentException.class); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup1", CoGroup.join("pc1", By.fieldNames("user", "country").withOptionalParticipation()) .join("pc2", By.fieldNames("user2", "country2").withOptionalParticipation()) .join("pc3", By.fieldNames("user3", "country3").withSideInput()) .crossProductJoin()); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testUnderspecifiedCoGroup() { PCollection<Row> pc1 = pipeline .apply( "Create1", Create.of(Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build())) .setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline .apply( "Create2", Create.of(Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build())) .setRowSchema(CG_SCHEMA_2); PCollection<Row> pc3 = pipeline.apply( "Create3", Create.of(Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build())); thrown.expect(IllegalArgumentException.class); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup", CoGroup.join("pc1", By.fieldNames("user", "country")) .join("pc2", By.fieldNames("user2", "country2"))); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testMismatchingKeys() { PCollection<Row> pc1 = pipeline .apply( "Create1", Create.of(Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build())) .setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline .apply( "Create2", Create.of(Row.withSchema(CG_SCHEMA_1).addValues("user1", 9, "us").build())) .setRowSchema(CG_SCHEMA_1); thrown.expect(IllegalArgumentException.class); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2) .apply( "CoGroup", CoGroup.join("pc1", By.fieldNames("user")).join("pc2", By.fieldNames("count"))); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testInnerJoin() { List<Row> pc1Rows = Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 5, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 6, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build()); List<Row> pc2Rows = Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 10, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 11, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 12, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 13, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 14, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 15, "ar").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 16, "ar").build()); List<Row> pc3Rows = Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 18, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 19, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 20, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 21, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 22, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 23, "ar").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 24, "ar").build()); PCollection<Row> pc1 = pipeline.apply("Create1", Create.of(pc1Rows)).setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline.apply("Create2", Create.of(pc2Rows)).setRowSchema(CG_SCHEMA_2); PCollection<Row> pc3 = pipeline.apply("Create3", Create.of(pc3Rows)).setRowSchema(CG_SCHEMA_3); Schema expectedSchema = Schema.builder() .addRowField("pc1", CG_SCHEMA_1) .addRowField("pc2", CG_SCHEMA_2) .addRowField("pc3", CG_SCHEMA_3) .build(); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup", CoGroup.join("pc1", By.fieldNames("user", "country")) .join("pc2", By.fieldNames("user2", "country2")) .join("pc3", By.fieldNames("user3", "country3")) .crossProductJoin()); assertEquals(expectedSchema, joined.getSchema()); List<Row> expectedJoinedRows = JoinTestUtils.innerJoin( pc1Rows, pc2Rows, pc3Rows, new String[] {"user", "country"}, new String[] {"user2", "country2"}, new String[] {"user3", "country3"}, expectedSchema); PAssert.that(joined).containsInAnyOrder(expectedJoinedRows); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testFullOuterJoin() { List<Row> pc1Rows = Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 5, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 6, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user3", 7, "ar").build()); List<Row> pc2Rows = Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 10, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 11, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 12, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 13, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 14, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 15, "ar").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 16, "ar").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 16, "es").build()); List<Row> pc3Rows = Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 18, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 19, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 20, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 21, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 22, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 23, "ar").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 24, "ar").build(), Row.withSchema(CG_SCHEMA_3).addValues("user27", 24, "se").build()); PCollection<Row> pc1 = pipeline.apply("Create1", Create.of(pc1Rows)).setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline.apply("Create2", Create.of(pc2Rows)).setRowSchema(CG_SCHEMA_2); PCollection<Row> pc3 = pipeline.apply("Create3", Create.of(pc3Rows)).setRowSchema(CG_SCHEMA_3); Schema expectedSchema = Schema.builder() .addNullableField("pc1", FieldType.row(CG_SCHEMA_1)) .addNullableField("pc2", FieldType.row(CG_SCHEMA_2)) .addNullableField("pc3", FieldType.row(CG_SCHEMA_3)) .build(); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup", CoGroup.join("pc1", By.fieldNames("user", "country").withOptionalParticipation()) .join("pc2", By.fieldNames("user2", "country2").withOptionalParticipation()) .join("pc3", By.fieldNames("user3", "country3").withOptionalParticipation()) .crossProductJoin()); assertEquals(expectedSchema, joined.getSchema()); List<Row> expectedJoinedRows = JoinTestUtils.innerJoin( pc1Rows, pc2Rows, pc3Rows, new String[] {"user", "country"}, new String[] {"user2", "country2"}, new String[] {"user3", "country3"}, expectedSchema); expectedJoinedRows.add( Row.withSchema(expectedSchema) .addValues(Row.withSchema(CG_SCHEMA_1).addValues("user3", 7, "ar").build(), null, null) .build()); expectedJoinedRows.add( Row.withSchema(expectedSchema) .addValues(null, Row.withSchema(CG_SCHEMA_2).addValues("user2", 16, "es").build(), null) .build()); expectedJoinedRows.add( Row.withSchema(expectedSchema) .addValues( null, null, Row.withSchema(CG_SCHEMA_3).addValues("user27", 24, "se").build()) .build()); PAssert.that(joined).containsInAnyOrder(expectedJoinedRows); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testPartialOuterJoin() { List<Row> pc1Rows = Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 5, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 6, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build()); List<Row> pc2Rows = Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 10, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 11, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 12, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 13, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 14, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user3", 7, "ar").build()); List<Row> pc3Rows = Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 18, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 19, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 20, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 21, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 22, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 23, "ar").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 24, "ar").build(), Row.withSchema(CG_SCHEMA_3).addValues("user3", 25, "ar").build()); PCollection<Row> pc1 = pipeline.apply("Create1", Create.of(pc1Rows)).setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline.apply("Create2", Create.of(pc2Rows)).setRowSchema(CG_SCHEMA_2); PCollection<Row> pc3 = pipeline.apply("Create3", Create.of(pc3Rows)).setRowSchema(CG_SCHEMA_3); Schema expectedSchema = Schema.builder() .addField("pc1", FieldType.row(CG_SCHEMA_1)) .addNullableField("pc2", FieldType.row(CG_SCHEMA_2)) .addField("pc3", FieldType.row(CG_SCHEMA_3)) .build(); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup", CoGroup.join("pc1", By.fieldNames("user", "country")) .join("pc2", By.fieldNames("user2", "country2").withOptionalParticipation()) .join("pc3", By.fieldNames("user3", "country3")) .crossProductJoin()); assertEquals(expectedSchema, joined.getSchema()); List<Row> expectedJoinedRows = JoinTestUtils.innerJoin( pc1Rows, pc2Rows, pc3Rows, new String[] {"user", "country"}, new String[] {"user2", "country2"}, new String[] {"user3", "country3"}, expectedSchema); expectedJoinedRows.add( Row.withSchema(expectedSchema) .addValues( Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), null, Row.withSchema(CG_SCHEMA_3).addValues("user2", 23, "ar").build()) .build()); expectedJoinedRows.add( Row.withSchema(expectedSchema) .addValues( Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), null, Row.withSchema(CG_SCHEMA_3).addValues("user2", 24, "ar").build()) .build()); expectedJoinedRows.add( Row.withSchema(expectedSchema) .addValues( Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build(), null, Row.withSchema(CG_SCHEMA_3).addValues("user2", 23, "ar").build()) .build()); expectedJoinedRows.add( Row.withSchema(expectedSchema) .addValues( Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build(), null, Row.withSchema(CG_SCHEMA_3).addValues("user2", 24, "ar").build()) .build()); PAssert.that(joined).containsInAnyOrder(expectedJoinedRows); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testUnmatchedTags() { PCollection<Row> pc1 = pipeline.apply("Create1", Create.empty(CG_SCHEMA_1)); PCollection<Row> pc2 = pipeline.apply("Create2", Create.empty(CG_SCHEMA_2)); thrown.expect(IllegalArgumentException.class); PCollectionTuple.of("pc1", pc1, "pc2", pc2) .apply( CoGroup.join("pc1", By.fieldNames("user")) .join("pc3", By.fieldNames("user3")) .crossProductJoin()); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testPojo() { List<CgPojo> pc1Rows = Lists.newArrayList( new CgPojo("user1", 1, "us"), new CgPojo("user1", 2, "us"), new CgPojo("user1", 3, "il"), new CgPojo("user1", 4, "il")); List<CgPojo> pc2Rows = Lists.newArrayList( new CgPojo("user1", 3, "us"), new CgPojo("user1", 4, "us"), new CgPojo("user1", 5, "il"), new CgPojo("user1", 6, "il")); PCollection<CgPojo> pc1 = pipeline.apply("Create1", Create.of(pc1Rows)); PCollection<CgPojo> pc2 = pipeline.apply("Create2", Create.of(pc2Rows)); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1) .and("pc2", pc2) .apply( CoGroup.join("pc1", By.fieldNames("user", "country")) .join("pc2", By.fieldNames("user", "country"))); Schema expectedSchema = Schema.builder() .addRowField("key", SIMPLE_CG_KEY_SCHEMA) .addIterableField("pc1", FieldType.row(CG_SCHEMA_1)) .addIterableField("pc2", FieldType.row(CG_SCHEMA_1)) .build(); List<Row> expected = Lists.newArrayList( Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user1", "us").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "us").build())) .build(), Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user1", "il").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 5, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 6, "il").build())) .build()); assertEquals(expectedSchema, joined.getSchema()); PAssert.that(joined).satisfies(actual -> containsJoinedFields(expected, actual)); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testCoGroupByDifferentFields() { PCollection<Row> pc1 = pipeline .apply( "Create1", Create.of( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 5, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 6, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build())) .setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline .apply( "Create2", Create.of( Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 10, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 11, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 12, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 13, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 14, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 15, "ar").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 16, "ar").build())) .setRowSchema(CG_SCHEMA_2); PCollection<Row> pc3 = pipeline .apply( "Create3", Create.of( Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 18, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 19, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 20, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 21, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 22, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 23, "ar").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 24, "ar").build())) .setRowSchema(CG_SCHEMA_3); Schema expectedSchema = Schema.builder() .addRowField("key", SIMPLE_CG_KEY_SCHEMA) .addIterableField("pc1", FieldType.row(CG_SCHEMA_1)) .addIterableField("pc2", FieldType.row(CG_SCHEMA_2)) .addIterableField("pc3", FieldType.row(CG_SCHEMA_3)) .build(); Row key1Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user1", "us").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 2, "us").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user1", 9, "us").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 10, "us").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user1", 17, "us").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 18, "us").build())) .build(); Row key2Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user1", "il").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user1", 3, "il").build(), Row.withSchema(CG_SCHEMA_1).addValues("user1", 4, "il").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user1", 11, "il").build(), Row.withSchema(CG_SCHEMA_2).addValues("user1", 12, "il").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user1", 19, "il").build(), Row.withSchema(CG_SCHEMA_3).addValues("user1", 20, "il").build())) .build(); Row key3Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user2", "fr").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 5, "fr").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 6, "fr").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user2", 13, "fr").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 14, "fr").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user2", 21, "fr").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 22, "fr").build())) .build(); Row key4Joined = Row.withSchema(expectedSchema) .addValue(Row.withSchema(SIMPLE_CG_KEY_SCHEMA).addValues("user2", "ar").build()) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_1).addValues("user2", 7, "ar").build(), Row.withSchema(CG_SCHEMA_1).addValues("user2", 8, "ar").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_2).addValues("user2", 15, "ar").build(), Row.withSchema(CG_SCHEMA_2).addValues("user2", 16, "ar").build())) .addIterable( Lists.newArrayList( Row.withSchema(CG_SCHEMA_3).addValues("user2", 23, "ar").build(), Row.withSchema(CG_SCHEMA_3).addValues("user2", 24, "ar").build())) .build(); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2, "pc3", pc3) .apply( "CoGroup", CoGroup.join("pc1", By.fieldNames("user", "country")) .join("pc2", By.fieldNames("user2", "country2")) .join("pc3", By.fieldNames("user3", "country3"))); List<Row> expected = ImmutableList.of(key1Joined, key2Joined, key3Joined, key4Joined); PAssert.that(joined).satisfies(actual -> containsJoinedFields(expected, actual)); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testMismatchingKeys() { PCollection<Row> pc1 = pipeline .apply( "Create1", Create.of(Row.withSchema(CG_SCHEMA_1).addValues("user1", 1, "us").build())) .setRowSchema(CG_SCHEMA_1); PCollection<Row> pc2 = pipeline .apply( "Create2", Create.of(Row.withSchema(CG_SCHEMA_1).addValues("user1", 9, "us").build())) .setRowSchema(CG_SCHEMA_1); thrown.expect(IllegalStateException.class); PCollection<Row> joined = PCollectionTuple.of("pc1", pc1, "pc2", pc2) .apply( "CoGroup", CoGroup.join("pc1", By.fieldNames("user")).join("pc2", By.fieldNames("count"))); pipeline.run(); }