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.sdk.transforms;
  
  import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.allOf;
  import static org.hamcrest.Matchers.everyItem;
  import static org.hamcrest.Matchers.in;
  import static org.hamcrest.core.Is.is;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertTrue;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  import java.util.TreeSet;
  import org.apache.beam.sdk.TestUtils;
  import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
  import org.apache.beam.sdk.coders.StringUtf8Coder;
  import org.apache.beam.sdk.testing.CombineFnTester;
  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.transforms.display.DisplayData;
  import org.apache.beam.sdk.transforms.windowing.FixedWindows;
  import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
  import org.apache.beam.sdk.transforms.windowing.Window;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.TimestampedValue;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Joiner;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
  import org.hamcrest.Matchers;
  import org.joda.time.Duration;
  import org.joda.time.Instant;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.experimental.categories.Category;
  import org.junit.runner.RunWith;
  import org.junit.runners.JUnit4;
  import org.junit.runners.Parameterized;
  
  /** Tests for Sample transform. */
  public class SampleTest {
    private static final Integer[] EMPTY = new Integer[] {};
    private static final Integer[] DATA = new Integer[] {1, 2, 3, 4, 5};
    private static final Integer[] REPEATED_DATA = new Integer[] {1, 1, 2, 2, 3, 3, 4, 4, 5, 5};
  
    /** Test variations for Sample transform. */
    @RunWith(Parameterized.class)
    public static class PickAnyTest {
      @Rule public final transient TestPipeline p = TestPipeline.create();
  
      @Parameterized.Parameters(name = "limit_{1}")
      public static Iterable<Object[]> data() throws IOException {
        return ImmutableList.<Object[]>builder()
            .add(
                new Object[] {TestUtils.NO_LINES, 0},
                new Object[] {TestUtils.NO_LINES, 1},
                new Object[] {TestUtils.LINES, 1},
                new Object[] {TestUtils.LINES, TestUtils.LINES.size() / 2},
                new Object[] {TestUtils.LINES, TestUtils.LINES.size() * 2},
                new Object[] {TestUtils.LINES, TestUtils.LINES.size() - 1},
                new Object[] {TestUtils.LINES, TestUtils.LINES.size()},
                new Object[] {TestUtils.LINES, TestUtils.LINES.size() + 1})
            .build();
      }
  
      @SuppressWarnings("DefaultAnnotationParam")
      @Parameterized.Parameter(0)
      public List<String> lines;
  
      @Parameterized.Parameter(1)
      public int limit;
  
      private static class VerifyAnySample implements SerializableFunction<Iterable<String>, Void> {
        private final List<String> lines;
        private final int limit;
  
        private VerifyAnySample(List<String> lines, int limit) {
          this.lines = lines;
          this.limit = limit;
        }
  
        @Override
        public Void apply(Iterable<String> actualIter) {
          final int expectedSize = Math.min(limit, lines.size());
  
          // Make sure actual is the right length, and is a
          // subset of expected.
          List<String> actual = new ArrayList<>();
          for (String s : actualIter) {
            actual.add(s);
          }
          assertEquals(expectedSize, actual.size());
          Set<String> actualAsSet = new TreeSet<>(actual);
          Set<String> linesAsSet = new TreeSet<>(lines);
          assertEquals(actual.size(), actualAsSet.size());
          assertEquals(lines.size(), linesAsSet.size());
          assertTrue(linesAsSet.containsAll(actualAsSet));
          return null;
        }
      }
  
      void runPickAnyTest(final List<String> lines, int limit) {
        checkArgument(new HashSet<>(lines).size() == lines.size(), "Duplicates are unsupported.");
  
        PCollection<String> input = p.apply(Create.of(lines).withCoder(StringUtf8Coder.of()));
  
        PCollection<String> output = input.apply(Sample.any(limit));
  
        PAssert.that(output).satisfies(new VerifyAnySample(lines, limit));
  
        p.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testPickAny() {
        runPickAnyTest(lines, limit);
      }
  
      @Test
      public void testCombineFn() {
        CombineFnTester.testCombineFn(
            Sample.combineFn(limit),
            lines,
            allOf(
                Matchers.<String>iterableWithSize(Math.min(lines.size(), limit)),
                everyItem(is(in(lines)))));
      }
    }
  
    /** Further tests for Sample transform. */
    @RunWith(JUnit4.class)
    public static class MiscTest {
  
      @Rule public final transient TestPipeline pipeline = TestPipeline.create();
  
      /**
       * Verifies that the result of a Sample operation contains the expected number of elements, and
       * that those elements are a subset of the items in expected.
       */
      @SuppressWarnings("rawtypes")
      public static class VerifyCorrectSample<T extends Comparable>
          implements SerializableFunction<Iterable<T>, Void> {
        private Object[] expectedValues;
        private int expectedSize;
  
        /**
         * expectedSize is the number of elements that the Sample should contain. expected is the set
         * of elements that the sample may contain.
         */
        @SafeVarargs
        VerifyCorrectSample(int expectedSize, T... expected) {
          this.expectedValues = expected;
          this.expectedSize = expectedSize;
        }
  
        /**
         * expectedSize is the number of elements that the Sample should contain. expected is the set
         * of elements that the sample may contain.
         */
        VerifyCorrectSample(int expectedSize, Collection<T> expected) {
          this.expectedValues = expected.toArray();
          this.expectedSize = expectedSize;
        }
  
        @Override
        @SuppressWarnings("unchecked")
        public Void apply(Iterable<T> in) {
          List<T> actual = new ArrayList<>();
          for (T elem : in) {
            actual.add(elem);
          }
  
          assertEquals(expectedSize, actual.size());
  
          Collections.sort(actual); // We assume that @expected is already sorted.
          int i = 0; // Index into @expected
          for (T s : actual) {
            boolean matchFound = false;
            for (; i < expectedValues.length; i++) {
              if (s.equals(expectedValues[i])) {
                matchFound = true;
                break;
              }
            }
            assertTrue("Invalid sample: " + Joiner.on(',').join(actual), matchFound);
            i++; // Don't match the same element again.
          }
          return null;
        }
      }
  
      private static TimestampedValue<Integer> tv(int i) {
        return TimestampedValue.of(i, new Instant(i * 1000));
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testSampleAny() {
        PCollection<Integer> input =
            pipeline
                .apply(
                    Create.timestamped(ImmutableList.of(tv(0), tv(1), tv(2), tv(3), tv(4), tv(5)))
                        .withCoder(BigEndianIntegerCoder.of()))
                .apply(Window.into(FixedWindows.of(Duration.standardSeconds(3))));
        PCollection<Integer> output = input.apply(Sample.any(2));
  
        PAssert.that(output)
            .inWindow(new IntervalWindow(new Instant(0), Duration.standardSeconds(3)))
            .satisfies(new VerifyCorrectSample<>(2, Arrays.asList(0, 1, 2)));
        PAssert.that(output)
            .inWindow(new IntervalWindow(new Instant(3000), Duration.standardSeconds(3)))
            .satisfies(new VerifyCorrectSample<>(2, Arrays.asList(3, 4, 5)));
        pipeline.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testSampleAnyEmpty() {
        PCollection<Integer> input = pipeline.apply(Create.empty(BigEndianIntegerCoder.of()));
        PCollection<Integer> output =
            input
                .apply(Window.into(FixedWindows.of(Duration.standardSeconds(3))))
                .apply(Sample.any(2));
  
        PAssert.that(output).satisfies(new VerifyCorrectSample<>(0, EMPTY));
        pipeline.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testSampleAnyZero() {
        PCollection<Integer> input =
            pipeline.apply(
                Create.timestamped(ImmutableList.of(tv(0), tv(1), tv(2), tv(3), tv(4), tv(5)))
                    .withCoder(BigEndianIntegerCoder.of()));
        PCollection<Integer> output =
            input
                .apply(Window.into(FixedWindows.of(Duration.standardSeconds(3))))
                .apply(Sample.any(0));
  
        PAssert.that(output)
            .inWindow(new IntervalWindow(new Instant(0), Duration.standardSeconds(3)))
            .satisfies(new VerifyCorrectSample<>(0, EMPTY));
        PAssert.that(output)
            .inWindow(new IntervalWindow(new Instant(3000), Duration.standardSeconds(3)))
            .satisfies(new VerifyCorrectSample<>(0, EMPTY));
        pipeline.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testSampleAnyInsufficientElements() {
        PCollection<Integer> input = pipeline.apply(Create.empty(BigEndianIntegerCoder.of()));
        PCollection<Integer> output =
            input
                .apply(Window.into(FixedWindows.of(Duration.standardSeconds(3))))
                .apply(Sample.any(10));
  
        PAssert.that(output)
            .inWindow(new IntervalWindow(new Instant(0), Duration.standardSeconds(3)))
            .satisfies(new VerifyCorrectSample<>(0, EMPTY));
        pipeline.run();
      }
  
      @Test(expected = IllegalArgumentException.class)
      public void testSampleAnyNegative() {
        pipeline.enableAbandonedNodeEnforcement(false);
        pipeline.apply(Create.empty(BigEndianIntegerCoder.of())).apply(Sample.any(-10));
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testSample() {
  
        PCollection<Integer> input =
            pipeline.apply(
                Create.of(ImmutableList.copyOf(DATA)).withCoder(BigEndianIntegerCoder.of()));
        PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(3));
  
        PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(3, DATA));
        pipeline.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testSampleEmpty() {
  
        PCollection<Integer> input = pipeline.apply(Create.empty(BigEndianIntegerCoder.of()));
        PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(3));
  
        PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(0, EMPTY));
        pipeline.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testSampleZero() {
  
        PCollection<Integer> input =
            pipeline.apply(
                Create.of(ImmutableList.copyOf(DATA)).withCoder(BigEndianIntegerCoder.of()));
        PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(0));
  
        PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(0, DATA));
        pipeline.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testSampleInsufficientElements() {
  
        PCollection<Integer> input =
            pipeline.apply(
                Create.of(ImmutableList.copyOf(DATA)).withCoder(BigEndianIntegerCoder.of()));
        PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(10));
  
        PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(5, DATA));
        pipeline.run();
      }
  
      @Test(expected = IllegalArgumentException.class)
      public void testSampleNegative() {
        pipeline.enableAbandonedNodeEnforcement(false);
  
        PCollection<Integer> input =
            pipeline.apply(
                Create.of(ImmutableList.copyOf(DATA)).withCoder(BigEndianIntegerCoder.of()));
        input.apply(Sample.fixedSizeGlobally(-1));
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testSampleMultiplicity() {
  
        PCollection<Integer> input =
            pipeline.apply(
                Create.of(ImmutableList.copyOf(REPEATED_DATA)).withCoder(BigEndianIntegerCoder.of()));
        // At least one value must be selected with multiplicity.
        PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(6));
  
        PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(6, REPEATED_DATA));
        pipeline.run();
      }
  
      @Test
      public void testSampleGetName() {
        assertEquals("Sample.Any", Sample.<String>any(1).getName());
      }
  
      @Test
      public void testDisplayData() {
        PTransform<?, ?> sampleAny = Sample.any(1234);
        DisplayData sampleAnyDisplayData = DisplayData.from(sampleAny);
        assertThat(sampleAnyDisplayData, hasDisplayItem("sampleSize", 1234));
  
        PTransform<?, ?> samplePerKey = Sample.fixedSizePerKey(2345);
        DisplayData perKeyDisplayData = DisplayData.from(samplePerKey);
        assertThat(perKeyDisplayData, hasDisplayItem("sampleSize", 2345));
      }
    }
  }
  