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.extensions.zetasketch;
  
  import com.google.zetasketch.HyperLogLogPlusPlus;
  import java.nio.ByteBuffer;
  import org.apache.beam.sdk.transforms.Combine;
  import org.apache.beam.sdk.transforms.DoFn;
  import org.apache.beam.sdk.transforms.PTransform;
  import org.apache.beam.sdk.transforms.ParDo;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * {@code PTransform}s to compute HyperLogLogPlusPlus (HLL++) sketches on data streams based on the
   * <a href="https://github.com/google/zetasketch">ZetaSketch</a> implementation.
   *
   * <p>HLL++ is an algorithm implemented by Google that estimates the count of distinct elements in a
   * data stream. HLL++ requires significantly less memory than the linear memory needed for exact
   * computation, at the cost of a small error. Cardinalities of arbitrary breakdowns can be computed
   * using the HLL++ sketch. See this <a
   * href="http://static.googleusercontent.com/media/research.google.com/en/us/pubs/archive/40671.pdf">published
   * paper</a> for details about the algorithm.
   *
   * <p>HLL++ functions are also supported in <a
   * href="https://cloud.google.com/bigquery/docs/reference/standard-sql/hll_functions">Google Cloud
   * BigQuery</a>. The {@code HllCount PTransform}s provided here produce and consume sketches
   * compatible with BigQuery.
   *
   * <p>For detailed design of this class, see https://s.apache.org/hll-in-beam.
   *
   * <h3>Examples</h3>
   *
   * <h4>Example 1: Create long-type sketch for a {@code PCollection<Long>} and specify precision</h4>
   *
   * <pre>{@code
   * PCollection<Long> input = ...;
   * int p = ...;
   * PCollection<byte[]> sketch = input.apply(HllCount.Init.forLongs().withPrecision(p).globally());
   * }</pre>
   *
   * <h4>Example 2: Create bytes-type sketch for a {@code PCollection<KV<String, byte[]>>}</h4>
   *
   * <pre>{@code
   * PCollection<KV<String, byte[]>> input = ...;
   * PCollection<KV<String, byte[]>> sketch = input.apply(HllCount.Init.forBytes().perKey());
   * }</pre>
   *
   * <h4>Example 3: Merge existing sketches in a {@code PCollection<byte[]>} into a new one</h4>
   *
   * <pre>{@code
   * PCollection<byte[]> sketches = ...;
   * PCollection<byte[]> mergedSketch = sketches.apply(HllCount.MergePartial.globally());
   * }</pre>
   *
   * <h4>Example 4: Estimates the count of distinct elements in a {@code PCollection<String>}</h4>
   *
   * <pre>{@code
   * PCollection<String> input = ...;
   * PCollection<Long> countDistinct =
   *     input.apply(HllCount.Init.forStrings().globally()).apply(HllCount.Extract.globally());
   * }</pre>
   *
   * Note: Currently HllCount does not work on FnAPI workers. See <a
   * href="https://github.com/apache/beam/issues/19698">Issue #19698</a>.
   */
  public final class HllCount {
  
    private static final Logger LOG = LoggerFactory.getLogger(HllCount.class);
  
    /**
     * The minimum {@code precision} value you can set in {@link Init.Builder#withPrecision(int)} is
     * {@value}.
     */
    public static final int MINIMUM_PRECISION = HyperLogLogPlusPlus.MINIMUM_PRECISION;
  
    /**
     * The maximum {@code precision} value you can set in {@link Init.Builder#withPrecision(int)} is
     * {@value}.
     */
    public static final int MAXIMUM_PRECISION = HyperLogLogPlusPlus.MAXIMUM_PRECISION;
  
    /**
     * The default {@code precision} value used in {@link Init.Builder#withPrecision(int)} is
     * {@value}.
     */
    public static final int DEFAULT_PRECISION = HyperLogLogPlusPlus.DEFAULT_NORMAL_PRECISION;
  
    // Cannot be instantiated. This class is intended to be a namespace only.
    private HllCount() {}
  
    /**
     * Converts the passed-in sketch from {@code ByteBuffer} to {@code byte[]}, mapping {@code null
     * ByteBuffer}s (representing empty sketches) to empty {@code byte[]}s.
     *
     * <p>Utility method to convert sketches materialized with ZetaSQL/BigQuery to valid inputs for
     * Beam {@code HllCount} transforms.
     */
    public static byte[] getSketchFromByteBuffer(@Nullable ByteBuffer bf) {
      if (bf == null) {
        return new byte[0];
      } else {
        byte[] result = new byte[bf.remaining()];
        bf.get(result);
        return result;
      }
    }
  
    /**
     * Provides {@code PTransform}s to aggregate inputs into HLL++ sketches. The four supported input
     * types are {@code Integer}, {@code Long}, {@code String}, and {@code byte[]}.
     *
     * <p>Sketches are represented using the {@code byte[]} type. Sketches of the same type can be
     * merged into a new sketch using {@link HllCount.MergePartial}. Estimated count of distinct
     * elements can be extracted from sketches using {@link HllCount.Extract}.
     *
     * <p>An "empty sketch" represented by an byte array of length 0 is returned if the input {@code
     * PCollection} is empty.
     *
     * <p>Corresponds to the {@code HLL_COUNT.INIT(input [, precision])} function in <a
     * href="https://cloud.google.com/bigquery/docs/reference/standard-sql/hll_functions">BigQuery</a>.
     */
    public static final class Init {
  
      // Cannot be instantiated. This class is intended to be a namespace only.
      private Init() {}
  
      /**
       * Returns a {@link Builder} for a {@code HllCount.Init} combining {@code PTransform} that
       * computes integer-type HLL++ sketches. Call {@link Builder#globally()} or {@link
       * Builder#perKey()} on the returning {@link Builder} to finalize the {@code PTransform}.
       *
       * <p>Calling {@link Builder#globally()} returns a {@code PTransform} that takes an input {@code
       * PCollection<Integer>} and returns a {@code PCollection<byte[]>} which consists of the
       * integer-type HLL++ sketch computed from the elements in the input {@code PCollection}.
       *
       * <p>Calling {@link Builder#perKey()} returns a {@code PTransform} that takes an input {@code
       * PCollection<KV<K, Integer>>} and returns a {@code PCollection<KV<K, byte[]>>} which consists
       * of the per-key integer-type HLL++ sketch computed from the values matching each key in the
       * input {@code PCollection}.
       *
       * <p>Integer-type sketches cannot be merged with sketches of other types.
       */
      public static Builder<Integer> forIntegers() {
        return new Builder<>(HllCountInitFn.forInteger());
      }
  
      /**
       * Returns a {@link Builder} for a {@code HllCount.Init} combining {@code PTransform} that
       * computes long-type HLL++ sketches. Call {@link Builder#globally()} or {@link
       * Builder#perKey()} on the returning {@link Builder} to finalize the {@code PTransform}.
       *
       * <p>Calling {@link Builder#globally()} returns a {@code PTransform} that takes an input {@code
       * PCollection<Long>} and returns a {@code PCollection<byte[]>} which consists of the long-type
       * HLL++ sketch computed from the elements in the input {@code PCollection}.
       *
       * <p>Calling {@link Builder#perKey()} returns a {@code PTransform} that takes an input {@code
       * PCollection<KV<K, Long>>} and returns a {@code PCollection<KV<K, byte[]>>} which consists of
       * the per-key long-type HLL++ sketch computed from the values matching each key in the input
       * {@code PCollection}.
       *
       * <p>Long-type sketches cannot be merged with sketches of other types.
       */
      public static Builder<Long> forLongs() {
        return new Builder<>(HllCountInitFn.forLong());
      }
  
      /**
       * Returns a {@link Builder} for a {@code HllCount.Init} combining {@code PTransform} that
       * computes string-type HLL++ sketches. Call {@link Builder#globally()} or {@link
       * Builder#perKey()} on the returning {@link Builder} to finalize the {@code PTransform}.
       *
       * <p>Calling {@link Builder#globally()} returns a {@code PTransform} that takes an input {@code
       * PCollection<String>} and returns a {@code PCollection<byte[]>} which consists of the
       * string-type HLL++ sketch computed from the elements in the input {@code PCollection}.
       *
       * <p>Calling {@link Builder#perKey()} returns a {@code PTransform} that takes an input {@code
       * PCollection<KV<K, String>>} and returns a {@code PCollection<KV<K, byte[]>>} which consists
       * of the per-key string-type HLL++ sketch computed from the values matching each key in the
       * input {@code PCollection}.
       *
       * <p>String-type sketches cannot be merged with sketches of other types.
       */
      public static Builder<String> forStrings() {
        return new Builder<>(HllCountInitFn.forString());
      }
  
      /**
       * Returns a {@link Builder} for a {@code HllCount.Init} combining {@code PTransform} that
       * computes bytes-type HLL++ sketches. Call {@link Builder#globally()} or {@link
       * Builder#perKey()} on the returning {@link Builder} to finalize the {@code PTransform}.
       *
       * <p>Calling {@link Builder#globally()} returns a {@code PTransform} that takes an input {@code
       * PCollection<byte[]>} and returns a {@code PCollection<byte[]>} which consists of the
       * bytes-type HLL++ sketch computed from the elements in the input {@code PCollection}.
       *
       * <p>Calling {@link Builder#perKey()} returns a {@code PTransform} that takes an input {@code
       * PCollection<KV<K, byte[]>>} and returns a {@code PCollection<KV<K, byte[]>>} which consists
       * of the per-key bytes-type HLL++ sketch computed from the values matching each key in the
       * input {@code PCollection}.
       *
       * <p>Bytes-type sketches cannot be merged with sketches of other types.
       */
      public static Builder<byte[]> forBytes() {
        return new Builder<>(HllCountInitFn.forBytes());
      }
  
      /**
       * Builder for the {@code HllCount.Init} combining {@code PTransform}.
       *
       * <p>Call {@link #withPrecision(int)} to customize the {@code precision} parameter of the
       * sketch.
       *
       * <p>Call {@link #globally()} or {@link #perKey()} to finalize the {@code PTransform}.
       *
       * @param <InputT> element type or value type in {@code KV}s of the input {@code PCollection} to
       *     the {@code PTransform} being built
       */
      public static final class Builder<InputT extends @Nullable Object> {
  
        private final HllCountInitFn<InputT, ?> initFn;
  
        private Builder(HllCountInitFn<InputT, ?> initFn) {
          this.initFn = initFn;
        }
  
        /**
         * Explicitly set the {@code precision} parameter used to compute HLL++ sketch.
         *
         * <p>Valid range is between {@link #MINIMUM_PRECISION} and {@link #MAXIMUM_PRECISION}. If
         * this method is not called, {@link #DEFAULT_PRECISION} will be used. Sketches computed using
         * different {@code precision}s cannot be merged together.
         *
         * @param precision the {@code precision} parameter used to compute HLL++ sketch
         */
        public Builder<InputT> withPrecision(int precision) {
          initFn.setPrecision(precision);
          return this;
        }
  
        /**
         * Returns a {@link Combine.Globally} {@code PTransform} that takes an input {@code
         * PCollection<InputT>} and returns a {@code PCollection<byte[]>} which consists of the HLL++
         * sketch computed from the elements in the input {@code PCollection}.
         *
         * <p>Returns a singleton {@code PCollection} with an "empty sketch" (byte array of length 0)
         * if the input {@code PCollection} is empty.
         */
        public Combine.Globally<InputT, byte[]> globally() {
          return Combine.globally(initFn);
        }
  
        /**
         * Returns a {@link Combine.PerKey} {@code PTransform} that takes an input {@code
         * PCollection<KV<K, InputT>>} and returns a {@code PCollection<KV<K, byte[]>>} which consists
         * of the per-key HLL++ sketch computed from the values matching each key in the input {@code
         * PCollection}.
         */
        public <K> Combine.PerKey<K, InputT, byte[]> perKey() {
          return Combine.perKey(initFn);
        }
  
        HllCountInitFn<InputT, ?> asUdaf() {
          return initFn;
        }
      }
    }
  
    /**
     * Provides {@code PTransform}s to merge HLL++ sketches into a new sketch.
     *
     * <p>Only sketches of the same type can be merged together. If incompatible sketches are
     * provided, a runtime error will occur.
     *
     * <p>If sketches of different {@code precision}s are merged, the merged sketch will get the
     * minimum precision encountered among all the input sketches.
     *
     * <p>An "empty sketch" represented by an byte array of length 0 is returned if the input {@code
     * PCollection} is empty.
     *
     * <p>Corresponds to the {@code HLL_COUNT.MERGE_PARTIAL(sketch)} function in <a
     * href="https://cloud.google.com/bigquery/docs/reference/standard-sql/hll_functions">BigQuery</a>.
     */
    public static final class MergePartial {
  
      // Cannot be instantiated. This class is intended to be a namespace only.
      private MergePartial() {}
  
      /**
       * Returns a {@link Combine.Globally} {@code PTransform} that takes an input {@code
       * PCollection<byte[]>} of HLL++ sketches and returns a {@code PCollection<byte[]>} of a new
       * sketch merged from the input sketches.
       *
       * <p>Only sketches of the same type can be merged together. If incompatible sketches are
       * provided, a runtime error will occur.
       *
       * <p>If sketches of different {@code precision}s are merged, the merged sketch will get the
       * minimum precision encountered among all the input sketches.
       *
       * <p>Returns a singleton {@code PCollection} with an "empty sketch" (byte array of length 0) if
       * the input {@code PCollection} is empty.
       */
      public static Combine.Globally<byte @Nullable [], byte[]> globally() {
        return Combine.globally(HllCountMergePartialFn.create());
      }
  
      /**
       * Returns a {@link Combine.PerKey} {@code PTransform} that takes an input {@code
       * PCollection<KV<K, byte[]>>} of (key, HLL++ sketch) pairs and returns a {@code
       * PCollection<KV<K, byte[]>>} of (key, new sketch merged from the input sketches under the
       * key).
       *
       * <p>If sketches of different {@code precision}s are merged, the merged sketch will get the
       * minimum precision encountered among all the input sketches.
       *
       * <p>Only sketches of the same type can be merged together. If incompatible sketches are
       * provided, a runtime error will occur.
       */
      public static <K> Combine.PerKey<K, byte @Nullable [], byte[]> perKey() {
        return Combine.perKey(HllCountMergePartialFn.create());
      }
    }
  
    /**
     * Provides {@code PTransform}s to extract the estimated count of distinct elements (as {@code
     * Long}s) from each HLL++ sketch.
     *
     * <p>When extracting from an "empty sketch" represented by an byte array of length 0, the result
     * returned is 0.
     *
     * <p>Corresponds to the {@code HLL_COUNT.EXTRACT(sketch)} function in <a
     * href="https://cloud.google.com/bigquery/docs/reference/standard-sql/hll_functions">BigQuery</a>.
     */
    public static final class Extract {
  
      // Cannot be instantiated. This class is intended to be a namespace only.
      private Extract() {}
  
      /**
       * Returns a {@code PTransform} that takes an input {@code PCollection<byte[]>} of HLL++
       * sketches and returns a {@code PCollection<Long>} of the estimated count of distinct elements
       * extracted from each sketch.
       *
       * <p>Returns 0 if the input element is an "empty sketch" (byte array of length 0).
       */
      public static PTransform<PCollection<byte[]>, PCollection<Long>> globally() {
        return new Globally();
      }
  
      /**
       * Returns a {@code PTransform} that takes an input {@code PCollection<KV<K, byte[]>>} of (key,
       * HLL++ sketch) pairs and returns a {@code PCollection<KV<K, Long>>} of (key, estimated count
       * of distinct elements extracted from each sketch).
       */
      public static <K> PTransform<PCollection<KV<K, byte[]>>, PCollection<KV<K, Long>>> perKey() {
        return new PerKey<K>();
      }
  
      private static final class Globally extends PTransform<PCollection<byte[]>, PCollection<Long>> {
  
        @Override
        public PCollection<Long> expand(PCollection<byte[]> input) {
          return input.apply(
              ParDo.of(
                  new DoFn<byte[], Long>() {
                    @ProcessElement
                    public void processElement(
                        @Element byte[] sketch, OutputReceiver<Long> receiver) {
                      if (sketch == null) {
                        LOG.warn(
                            "Received a null and treated it as an empty sketch. "
                                + "Consider replacing nulls with empty byte arrays (byte[0]) "
                                + "in upstream transforms for better space-efficiency and safety.");
                        receiver.output(0L);
                      } else if (sketch.length == 0) {
                        receiver.output(0L);
                      } else {
                        receiver.output(HyperLogLogPlusPlus.forProto(sketch).result());
                      }
                    }
                  }));
        }
      }
  
      private static final class PerKey<K>
          extends PTransform<PCollection<KV<K, byte[]>>, PCollection<KV<K, Long>>> {
  
        @Override
        public PCollection<KV<K, Long>> expand(PCollection<KV<K, byte[]>> input) {
          return input.apply(
              ParDo.of(
                  new DoFn<KV<K, byte[]>, KV<K, Long>>() {
                    @ProcessElement
                    public void processElement(
                        @Element KV<K, byte[]> kv, OutputReceiver<KV<K, Long>> receiver) {
                      byte[] sketch = kv.getValue();
                      if (sketch == null) {
                        LOG.warn(
                            "Received a null and treated it as an empty sketch. "
                                + "Consider replacing nulls with empty byte arrays (byte[0]) "
                                + "in upstream transforms for better space-efficiency and safety.");
                        receiver.output(KV.of(kv.getKey(), 0L));
                      } else if (sketch.length == 0) {
                        receiver.output(KV.of(kv.getKey(), 0L));
                      } else {
                        receiver.output(
                            KV.of(kv.getKey(), HyperLogLogPlusPlus.forProto(sketch).result()));
                      }
                    }
                  }));
        }
      }
    }
  }
  