content: /*
   * 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.runners.spark.translation.streaming;
  
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.containsInAnyOrder;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.fail;
  
  import org.apache.beam.runners.spark.SparkPipelineOptions;
  import org.apache.beam.runners.spark.StreamingTest;
  import org.apache.beam.runners.spark.TestSparkPipelineOptions;
  import org.apache.beam.runners.spark.io.CreateStream;
  import org.apache.beam.sdk.coders.KvCoder;
  import org.apache.beam.sdk.coders.VarIntCoder;
  import org.apache.beam.sdk.options.PipelineOptions;
  import org.apache.beam.sdk.testing.PAssert;
  import org.apache.beam.sdk.testing.TestPipeline;
  import org.apache.beam.sdk.transforms.SerializableFunction;
  import org.apache.beam.sdk.transforms.join.CoGbkResult;
  import org.apache.beam.sdk.transforms.join.CoGroupByKey;
  import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple;
  import org.apache.beam.sdk.transforms.windowing.FixedWindows;
  import org.apache.beam.sdk.transforms.windowing.Window;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.TimestampedValue;
  import org.apache.beam.sdk.values.TupleTag;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
  import org.joda.time.Duration;
  import org.joda.time.Instant;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.experimental.categories.Category;
  
  /** A test that verifies that CoGroupByKey works in streaming mode in spark runner. */
  public class SparkCoGroupByKeyStreamingTest {
  
    private static final TupleTag<Integer> INPUT1_TAG = new TupleTag<>("input1");
    private static final TupleTag<Integer> INPUT2_TAG = new TupleTag<>("input2");
  
    @Rule public final TestPipeline pipeline = TestPipeline.fromOptions(streamingOptions());
  
    private Duration batchDuration() {
      return Duration.millis(
          pipeline.getOptions().as(SparkPipelineOptions.class).getBatchIntervalMillis());
    }
  
    @Category(StreamingTest.class)
    @Test
    public void testInStreamingMode() throws Exception {
      Instant instant = new Instant(0);
      CreateStream<KV<Integer, Integer>> source1 =
          CreateStream.of(KvCoder.of(VarIntCoder.of(), VarIntCoder.of()), batchDuration())
              .emptyBatch()
              .advanceWatermarkForNextBatch(instant)
              .nextBatch(
                  TimestampedValue.of(KV.of(1, 1), instant),
                  TimestampedValue.of(KV.of(1, 2), instant),
                  TimestampedValue.of(KV.of(1, 3), instant))
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardSeconds(1L)))
              .nextBatch(
                  TimestampedValue.of(KV.of(2, 4), instant.plus(Duration.standardSeconds(1L))),
                  TimestampedValue.of(KV.of(2, 5), instant.plus(Duration.standardSeconds(1L))),
                  TimestampedValue.of(KV.of(2, 6), instant.plus(Duration.standardSeconds(1L))))
              .advanceNextBatchWatermarkToInfinity();
  
      CreateStream<KV<Integer, Integer>> source2 =
          CreateStream.of(KvCoder.of(VarIntCoder.of(), VarIntCoder.of()), batchDuration())
              .emptyBatch()
              .advanceWatermarkForNextBatch(instant)
              .nextBatch(
                  TimestampedValue.of(KV.of(1, 11), instant),
                  TimestampedValue.of(KV.of(1, 12), instant),
                  TimestampedValue.of(KV.of(1, 13), instant))
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardSeconds(1L)))
              .nextBatch(
                  TimestampedValue.of(KV.of(2, 14), instant.plus(Duration.standardSeconds(1L))),
                  TimestampedValue.of(KV.of(2, 15), instant.plus(Duration.standardSeconds(1L))),
                  TimestampedValue.of(KV.of(2, 16), instant.plus(Duration.standardSeconds(1L))))
              .advanceNextBatchWatermarkToInfinity();
  
      PCollection<KV<Integer, Integer>> input1 =
          pipeline
              .apply("create source1", source1)
              .apply(
                  "window input1",
                  Window.<KV<Integer, Integer>>into(FixedWindows.of(Duration.standardSeconds(3L)))
                      .withAllowedLateness(Duration.ZERO));
  
      PCollection<KV<Integer, Integer>> input2 =
          pipeline
              .apply("create source2", source2)
              .apply(
                  "window input2",
                  Window.<KV<Integer, Integer>>into(FixedWindows.of(Duration.standardSeconds(3L)))
                      .withAllowedLateness(Duration.ZERO));
  
      PCollection<KV<Integer, CoGbkResult>> output =
          KeyedPCollectionTuple.of(INPUT1_TAG, input1)
              .and(INPUT2_TAG, input2)
              .apply(CoGroupByKey.create());
      PAssert.that("Wrong output of the join using CoGroupByKey in streaming mode", output)
          .satisfies(
              (SerializableFunction<Iterable<KV<Integer, CoGbkResult>>, Void>)
                  input -> {
                    assertEquals("Wrong size of the output PCollection", 2, Iterables.size(input));
                    for (KV<Integer, CoGbkResult> element : input) {
                      if (element.getKey() == 1) {
                        Iterable<Integer> input1Elements = element.getValue().getAll(INPUT1_TAG);
                        assertEquals(
                            "Wrong number of values for output elements for tag input1 and key 1",
                            3,
                            Iterables.size(input1Elements));
                        assertThat(
                            "Elements of PCollection input1 for key \"1\" are not present in the output PCollection",
                            input1Elements,
                            containsInAnyOrder(1, 2, 3));
                        Iterable<Integer> input2Elements = element.getValue().getAll(INPUT2_TAG);
                        assertEquals(
                            "Wrong number of values for output elements for tag input2 and key 1",
                            3,
                            Iterables.size(input2Elements));
                        assertThat(
                            "Elements of PCollection input2 for key \"1\" are not present in the output PCollection",
                            input2Elements,
                            containsInAnyOrder(11, 12, 13));
                      } else if (element.getKey() == 2) {
                        Iterable<Integer> input1Elements = element.getValue().getAll(INPUT1_TAG);
                        assertEquals(
                            "Wrong number of values for output elements for tag input1 and key 2",
                            3,
                            Iterables.size(input1Elements));
                        assertThat(
                            "Elements of PCollection input1 for key \"2\" are not present in the output PCollection",
                            input1Elements,
                            containsInAnyOrder(4, 5, 6));
                        Iterable<Integer> input2Elements = element.getValue().getAll(INPUT2_TAG);
                        assertEquals(
                            "Wrong number of values for output elements for tag input2 and key 2",
                            3,
                            Iterables.size(input2Elements));
                        assertThat(
                            "Elements of PCollection input2 for key \"2\" are not present in the output PCollection",
                            input2Elements,
                            containsInAnyOrder(14, 15, 16));
                      } else {
                        fail("Unknown key in the output PCollection");
                      }
                    }
                    return null;
                  });
      pipeline.run();
    }
  
    private static PipelineOptions streamingOptions() {
      PipelineOptions options = TestPipeline.testingPipelineOptions();
      options.as(TestSparkPipelineOptions.class).setStreaming(true);
      return options;
    }
  }
  