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.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.both;
  import static org.hamcrest.Matchers.greaterThan;
  import static org.hamcrest.Matchers.lessThan;
  import static org.hamcrest.core.Is.is;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertTrue;
  import static org.junit.Assert.fail;
  
  import java.io.IOException;
  import java.io.Serializable;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.List;
  import java.util.stream.Collectors;
  import java.util.stream.LongStream;
  import org.apache.beam.sdk.TestUtils;
  import org.apache.beam.sdk.coders.DoubleCoder;
  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.ApproximateUnique.ApproximateUniqueCombineFn;
  import org.apache.beam.sdk.transforms.display.DisplayData;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.PCollectionView;
  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.runner.RunWith;
  import org.junit.runners.JUnit4;
  import org.junit.runners.Parameterized;
  
  /** Tests for the ApproximateUnique transform. */
  public class ApproximateUniqueTest implements Serializable {
    // implements Serializable just to make it easy to use anonymous inner DoFn subclasses
  
    @Rule public final transient TestPipeline p = TestPipeline.create();
  
    private static class VerifyEstimateFn implements SerializableFunction<Long, Void> {
      private final long uniqueCount;
      private final int sampleSize;
  
      private VerifyEstimateFn(final long uniqueCount, final int sampleSize) {
        this.uniqueCount = uniqueCount;
        this.sampleSize = sampleSize;
      }
  
      @Override
      public Void apply(final Long estimate) {
        verifyEstimate(uniqueCount, sampleSize, estimate);
        return null;
      }
    }
  
    /**
     * Checks that the estimation error, i.e., the difference between {@code uniqueCount} and {@code
     * estimate} is less than {@code 2 / sqrt(sampleSize}).
     */
    private static void verifyEstimate(
        final long uniqueCount, final int sampleSize, final long estimate) {
      if (uniqueCount < sampleSize) {
        assertEquals(
            "Number of hashes is less than the sample size. " + "Estimate should be exact",
            uniqueCount,
            estimate);
      }
  
      final double error = 100.0 * Math.abs(estimate - uniqueCount) / uniqueCount;
      final double maxError = 100.0 * 2 / Math.sqrt(sampleSize);
  
      assertTrue(
          "Estimate="
              + estimate
              + " Actual="
              + uniqueCount
              + " Error="
              + error
              + "%, MaxError="
              + maxError
              + "%.",
          error < maxError);
  
      assertTrue(
          "Estimate="
              + estimate
              + " Actual="
              + uniqueCount
              + " Error="
              + error
              + "%, MaxError="
              + maxError
              + "%.",
          error < maxError);
    }
  
    private static Matcher<Long> estimateIsWithinRangeFor(
        final long uniqueCount, final int sampleSize) {
      if (uniqueCount <= sampleSize) {
        return is(uniqueCount);
      } else {
        long maxError = (long) Math.ceil(2.0 * uniqueCount / Math.sqrt(sampleSize));
        return both(lessThan(uniqueCount + maxError)).and(greaterThan(uniqueCount - maxError));
      }
    }
  
    private static class VerifyEstimatePerKeyFn
        implements SerializableFunction<Iterable<KV<Long, Long>>, Void> {
  
      private final int sampleSize;
  
      private VerifyEstimatePerKeyFn(final int sampleSize) {
        this.sampleSize = sampleSize;
      }
  
      @Override
      public Void apply(final Iterable<KV<Long, Long>> estimatePerKey) {
        for (final KV<Long, Long> result : estimatePerKey) {
          verifyEstimate(result.getKey(), sampleSize, result.getValue());
        }
        return null;
      }
    }
  
    /** Tests for ApproximateUnique with duplicates. */
    @RunWith(Parameterized.class)
    public static class ApproximateUniqueWithDuplicatesTest extends ApproximateUniqueTest {
  
      @Parameterized.Parameter public int elementCount;
  
      @Parameterized.Parameter(1)
      public int uniqueCount;
  
      @Parameterized.Parameter(2)
      public int sampleSize;
  
      @Parameterized.Parameters(name = "total_{0}_unique_{1}_sample_{2}")
      public static Iterable<Object[]> data() throws IOException {
        return ImmutableList.<Object[]>builder()
            .add(
                new Object[] {100, 100, 100},
                new Object[] {1000, 1000, 100},
                new Object[] {1500, 1000, 100},
                new Object[] {10000, 1000, 100})
            .build();
      }
  
      private void runApproximateUniqueWithDuplicates(
          final int elementCount, final int uniqueCount, final int sampleSize) {
  
        assert elementCount >= uniqueCount;
        final List<Double> elements = Lists.newArrayList();
        for (int i = 0; i < elementCount; i++) {
          elements.add(1.0 / (i % uniqueCount + 1));
        }
        Collections.shuffle(elements);
  
        final PCollection<Double> input = p.apply(Create.of(elements));
        final PCollection<Long> estimate = input.apply(ApproximateUnique.globally(sampleSize));
  
        PAssert.thatSingleton(estimate).satisfies(new VerifyEstimateFn(uniqueCount, sampleSize));
  
        p.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testApproximateUniqueWithDuplicates() {
        runApproximateUniqueWithDuplicates(elementCount, uniqueCount, sampleSize);
      }
    }
  
    /** Tests for ApproximateUnique with different sample sizes. */
    @RunWith(Parameterized.class)
    public static class ApproximateUniqueVariationsTest extends ApproximateUniqueTest {
  
      private static final int TEST_PAGES = 100;
      private static final List<String> TEST_LINES =
          new ArrayList<>(TEST_PAGES * TestUtils.LINES.size());
  
      static {
        for (int i = 0; i < TEST_PAGES; i++) {
          TEST_LINES.addAll(TestUtils.LINES);
        }
      }
  
      @Parameterized.Parameter public int sampleSize;
  
      @Parameterized.Parameters(name = "sampleSize_{0}")
      public static Iterable<Object[]> data() throws IOException {
        return ImmutableList.<Object[]>builder()
            .add(
                new Object[] {16},
                new Object[] {64},
                new Object[] {128},
                new Object[] {256},
                new Object[] {512},
                new Object[] {1000},
                new Object[] {2014},
                new Object[] {15})
            .build();
      }
  
      /**
       * Applies {@code ApproximateUnique(sampleSize)} verifying that the estimation error falls
       * within the maximum allowed error of {@code 2/sqrt(sampleSize)}.
       */
      private void runApproximateUniquePipeline(final int sampleSize) {
        final PCollection<String> input = p.apply(Create.of(TEST_LINES));
        final PCollection<Long> approximate = input.apply(ApproximateUnique.globally(sampleSize));
        final PCollectionView<Long> exact =
            input.apply(Distinct.create()).apply(Count.globally()).apply(View.asSingleton());
  
        final PCollection<KV<Long, Long>> approximateAndExact =
            approximate.apply(
                ParDo.of(
                        new DoFn<Long, KV<Long, Long>>() {
  
                          @ProcessElement
                          public void processElement(final ProcessContext c) {
                            c.output(KV.of(c.element(), c.sideInput(exact)));
                          }
                        })
                    .withSideInputs(exact));
  
        PAssert.that(approximateAndExact).satisfies(new VerifyEstimatePerKeyFn(sampleSize));
  
        p.run();
      }
  
      /**
       * Applies {@link ApproximateUnique} for different sample sizes and verifies that the estimation
       * error falls within the maximum allowed error of {@code 2 / sqrt(sampleSize)}.
       */
      @Test
      @Category(NeedsRunner.class)
      public void testApproximateUniqueWithDifferentSampleSizes() {
        if (sampleSize > 16) {
          runApproximateUniquePipeline(sampleSize);
        } else {
          try {
            p.enableAbandonedNodeEnforcement(false);
            runApproximateUniquePipeline(15);
            fail("Accepted sampleSize < 16");
          } catch (final IllegalArgumentException e) {
            assertTrue(
                "Expected an exception due to sampleSize < 16",
                e.getMessage().startsWith("ApproximateUnique needs a sampleSize >= 16"));
          }
        }
      }
    }
  
    /** Test ApproximateUniqueCombineFn. TestPipeline does seem to test merging partial results. */
    @RunWith(JUnit4.class)
    public static class ApproximateUniqueCombineFnTest {
  
      private void runCombineFnTest(long elementCount, long uniqueCount, int sampleSize) {
        List<Double> input =
            LongStream.range(0, elementCount)
                .mapToObj(i -> 1.0 / (i % uniqueCount + 1))
                .collect(Collectors.toList());
  
        CombineFnTester.testCombineFn(
            new ApproximateUniqueCombineFn<>(sampleSize, DoubleCoder.of()),
            input,
            estimateIsWithinRangeFor(uniqueCount, sampleSize));
      }
  
      @Test
      public void testFnWithSmallerFractionOfUniques() {
        runCombineFnTest(1000, 100, 16);
      }
  
      @Test
      public void testWithLargerFractionOfUniques() {
        runCombineFnTest(1000, 800, 100);
      }
  
      @Test
      public void testWithLargeSampleSize() {
        runCombineFnTest(200, 100, 150);
      }
    }
  
    /** Further tests for ApproximateUnique. */
    @RunWith(JUnit4.class)
    public static class ApproximateUniqueMiscTest extends ApproximateUniqueTest {
  
      @Test
      public void testEstimationErrorToSampleSize() {
        assertEquals(40000, ApproximateUnique.sampleSizeFromEstimationError(0.01));
        assertEquals(10000, ApproximateUnique.sampleSizeFromEstimationError(0.02));
        assertEquals(2500, ApproximateUnique.sampleSizeFromEstimationError(0.04));
        assertEquals(1600, ApproximateUnique.sampleSizeFromEstimationError(0.05));
        assertEquals(400, ApproximateUnique.sampleSizeFromEstimationError(0.1));
        assertEquals(100, ApproximateUnique.sampleSizeFromEstimationError(0.2));
        assertEquals(25, ApproximateUnique.sampleSizeFromEstimationError(0.4));
        assertEquals(16, ApproximateUnique.sampleSizeFromEstimationError(0.5));
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testApproximateUniqueWithSmallInput() {
        final PCollection<Integer> input = p.apply(Create.of(Arrays.asList(1, 2, 3, 3)));
  
        final PCollection<Long> estimate = input.apply(ApproximateUnique.globally(1000));
  
        PAssert.thatSingleton(estimate).isEqualTo(3L);
  
        p.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testApproximateUniqueWithSkewedDistributionsAndLargeSampleSize() {
        runApproximateUniqueWithSkewedDistributions(10000, 2000, 1000);
      }
  
      private void runApproximateUniqueWithSkewedDistributions(
          final int elementCount, final int uniqueCount, final int sampleSize) {
        final List<Integer> elements = Lists.newArrayList();
        // Zipf distribution with approximately elementCount items.
        final double s = 1 - 1.0 * uniqueCount / elementCount;
        final double maxCount = Math.pow(uniqueCount, s);
        for (int k = 0; k < uniqueCount; k++) {
          final int count = Math.max(1, (int) Math.round(maxCount * Math.pow(k, -s)));
          // Element k occurs count times.
          for (int c = 0; c < count; c++) {
            elements.add(k);
          }
        }
  
        final PCollection<Integer> input = p.apply(Create.of(elements));
        final PCollection<Long> estimate = input.apply(ApproximateUnique.globally(sampleSize));
  
        PAssert.thatSingleton(estimate).satisfies(new VerifyEstimateFn(uniqueCount, sampleSize));
  
        p.run();
      }
  
      @Test
      @Category(NeedsRunner.class)
      public void testApproximateUniquePerKey() {
        final List<KV<Long, Long>> elements = Lists.newArrayList();
        final List<Long> keys = ImmutableList.of(20L, 50L, 100L);
        final int elementCount = 1000;
        final int sampleSize = 100;
        // Use the key as the number of unique values.
        for (final long uniqueCount : keys) {
          for (long value = 0; value < elementCount; value++) {
            elements.add(KV.of(uniqueCount, value % uniqueCount));
          }
        }
  
        final PCollection<KV<Long, Long>> input = p.apply(Create.of(elements));
        final PCollection<KV<Long, Long>> counts = input.apply(ApproximateUnique.perKey(sampleSize));
  
        PAssert.that(counts).satisfies(new VerifyEstimatePerKeyFn(sampleSize));
  
        p.run();
      }
  
      @Test
      public void testApproximateUniqueGetName() {
        assertEquals("ApproximateUnique.PerKey", ApproximateUnique.<Long, Long>perKey(16).getName());
        assertEquals("ApproximateUnique.Globally", ApproximateUnique.<Integer>globally(16).getName());
      }
  
      @Test
      public void testDisplayData() {
        final ApproximateUnique.Globally<Integer> specifiedSampleSize =
            ApproximateUnique.globally(1234);
        final ApproximateUnique.PerKey<String, Integer> specifiedMaxError =
            ApproximateUnique.perKey(0.1234);
  
        assertThat(DisplayData.from(specifiedSampleSize), hasDisplayItem("sampleSize", 1234));
  
        final DisplayData maxErrorDisplayData = DisplayData.from(specifiedMaxError);
        assertThat(maxErrorDisplayData, hasDisplayItem("maximumEstimationError", 0.1234));
        assertThat(
            "calculated sampleSize should be included",
            maxErrorDisplayData,
            hasDisplayItem("sampleSize"));
      }
    }
  }
  