/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.sdk.schemas.transforms;

import static junit.framework.TestCase.assertEquals;
import static org.hamcrest.CoreMatchers.allOf;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.collection.IsIterableContainingInAnyOrder.containsInAnyOrder;

import java.util.List;
import java.util.Objects;
import org.apache.beam.sdk.schemas.JavaFieldSchema;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.schemas.Schema.TypeName;
import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
import org.apache.beam.sdk.schemas.transforms.CoGroup.By;
import org.apache.beam.sdk.schemas.utils.SchemaTestUtils.RowFieldMatcherIterableFieldAnyOrder;
import org.apache.beam.sdk.testing.NeedsRunner;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.UsesSchema;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
import org.hamcrest.Matcher;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for {@link CoGroup}. */
@RunWith(JUnit4.class)
@Category(UsesSchema.class)
@SuppressWarnings({
  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
})
public class CoGroupTest {
  @Rule public final transient TestPipeline pipeline = TestPipeline.create();
  @Rule public transient ExpectedException thrown = ExpectedException.none();

  private static final Schema CG_SCHEMA_1 =
      Schema.builder()
          .addStringField("user")
          .addInt32Field("count")
          .addStringField("country")
          .build();

  @DefaultSchema(JavaFieldSchema.class)
  public static class CgPojo {
    public String user;
    public int count;
    public String country;

    public CgPojo() {}

    public CgPojo(String user, int count, String country) {
      this.user = user;
      this.count = count;
      this.country = country;
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }
      CgPojo cgPojo = (CgPojo) o;
      return count == cgPojo.count
          && Objects.equals(user, cgPojo.user)
          && Objects.equals(country, cgPojo.country);
    }

    @Override
    public int hashCode() {
      return Objects.hash(user, count, country);
    }
  }

  private static final Schema SIMPLE_CG_KEY_SCHEMA =
      Schema.builder().addStringField("user").addStringField("country").build();

  @Test
  @Category(NeedsRunner.class)
  public void testCoGroupByFieldNames() {
    // Input
    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);

    // Output
    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();
  }

  private static final Schema CG_SCHEMA_2 =
      Schema.builder()
          .addStringField("user2")
          .addInt32Field("count2")
          .addStringField("country2")
          .build();

  private static final Schema CG_SCHEMA_3 =
      Schema.builder()
          .addStringField("user3")
          .addInt32Field("count3")
          .addStringField("country3")
          .build();

  @Test
  @Category(NeedsRunner.class)
  public void testCoGroupByDifferentFields() {
    // Inputs.
    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);

    // Expected outputs
    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);
    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);
    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);
    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);
    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);

    // Full outer join, so any field might be null.
    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);
    // Manually add the outer-join rows to the list of expected results.
    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);

    // Partial outer join. Missing entries in the "pc2" PCollection will be filled in with nulls,
    // but not others.
    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);

    // Manually add the outer-join rows to the list of expected results. Missing results from the
    // middle (pc2) PCollection are filled in with nulls. Missing events from other PCollections
    // are not. Events with key ("user2", "ar) show up in pc1 and pc3 but not in pc2, so we expect
    // the outer join to still produce those rows, with nulls for pc2. Events with key
    // ("user3", "ar) however show up in p2 and pc3, but not in pc1; since pc1 is marked for
    // full participation (no outer join), these events should not be included in the join.
    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();
  }

  private static Void containsJoinedFields(List<Row> expected, Iterable<Row> actual) {
    List<Matcher<? super Row>> matchers = Lists.newArrayList();
    for (Row row : expected) {
      List<Matcher> fieldMatchers = Lists.newArrayList();
      Schema schema = row.getSchema();
      fieldMatchers.add(
          new RowFieldMatcherIterableFieldAnyOrder(row.getSchema(), 0, row.getRow(0)));
      for (int i = 1; i < schema.getFieldCount(); ++i) {
        assertEquals(TypeName.ITERABLE, schema.getField(i).getType().getTypeName());
        fieldMatchers.add(
            new RowFieldMatcherIterableFieldAnyOrder(row.getSchema(), i, row.getIterable(i)));
      }
      matchers.add(allOf(fieldMatchers.toArray(new Matcher[0])));
    }
    assertThat(actual, containsInAnyOrder(matchers.toArray(new Matcher[0])));
    return null;
  }
}
