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.join;
  
  import java.io.IOException;
  import java.io.InputStream;
  import java.io.OutputStream;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.Iterator;
  import java.util.List;
  import java.util.NoSuchElementException;
  import java.util.Objects;
  import java.util.function.Function;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.coders.CoderException;
  import org.apache.beam.sdk.coders.CustomCoder;
  import org.apache.beam.sdk.coders.IterableCoder;
  import org.apache.beam.sdk.metrics.Counter;
  import org.apache.beam.sdk.metrics.Metrics;
  import org.apache.beam.sdk.util.common.Reiterator;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.TupleTag;
  import org.apache.beam.sdk.values.TupleTagList;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.AbstractIterator;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.HashMultiset;
  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.Iterators;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * A row result of a {@link CoGroupByKey}. This is a tuple of {@link Iterable}s produced for a given
   * key, and these can be accessed in different ways.
   */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class CoGbkResult {
    /**
     * A map of integer union tags to a list of union objects. Note: the key and the embedded union
     * tag are the same, so it is redundant to store it multiple times, but for now it makes encoding
     * easier.
     */
    private final List<Iterable<?>> valueMap;
  
    private final CoGbkResultSchema schema;
  
    private static final int DEFAULT_IN_MEMORY_ELEMENT_COUNT = 10_000;
  
    /**
     * Always try to cache at least this many elements per tag, even if it requires caching more than
     * the total in memory count.
     */
    private static final int DEFAULT_MIN_ELEMENTS_PER_TAG = 100;
  
    private static final Logger LOG = LoggerFactory.getLogger(CoGbkResult.class);
  
    private Counter keyCount = Metrics.counter(CoGbkResult.class, "cogbk-keys");
  
    private Counter largeKeyCount = Metrics.counter(CoGbkResult.class, "cogbk-large-keys");
  
    /**
     * A row in the {@link PCollection} resulting from a {@link CoGroupByKey} transform. Currently,
     * this row must fit into memory.
     *
     * @param schema the set of tuple tags used to refer to input tables and result values
     * @param taggedValues the raw results from a group-by-key
     */
    public CoGbkResult(CoGbkResultSchema schema, Iterable<RawUnionValue> taggedValues) {
      this(schema, taggedValues, DEFAULT_IN_MEMORY_ELEMENT_COUNT, DEFAULT_MIN_ELEMENTS_PER_TAG);
    }
  
    @SuppressWarnings("unchecked")
    public CoGbkResult(
        CoGbkResultSchema schema,
        Iterable<RawUnionValue> taggedValues,
        int inMemoryElementCount,
        int minElementsPerTag) {
      this.schema = schema;
      keyCount.inc();
      List<List<Object>> valuesByTag = new ArrayList<>();
      for (int unionTag = 0; unionTag < schema.size(); unionTag++) {
        valuesByTag.add(new ArrayList<>());
      }
  
      // Demultiplex the first imMemoryElementCount tagged union values
      // according to their tag.
      final Iterator<RawUnionValue> taggedIter = taggedValues.iterator();
      int elementCount = 0;
      while (taggedIter.hasNext()) {
        if (elementCount++ >= inMemoryElementCount) {
          // Let the tails be lazy.
          largeKeyCount.inc();
          break;
        }
        RawUnionValue value = taggedIter.next();
        // Make sure the given union tag has a corresponding tuple tag in the
        // schema.
        int unionTag = value.getUnionTag();
        if (schema.size() <= unionTag) {
          throw new IllegalStateException(
              "union tag " + unionTag + " has no corresponding tuple tag in the result schema");
        }
        valuesByTag.get(unionTag).add(value.getValue());
      }
  
      if (!taggedIter.hasNext()) {
        // Everything fits into memory, just store it.
        valueMap = (List) valuesByTag;
        return;
      }
  
      // If we get here, there were more elements than we can afford to
      // keep in memory, so we copy the re-iterable of remaining items
      // and append filtered views to each of the sorted lists computed earlier.
      LOG.info(
          "CoGbkResult has more than {} elements, reiteration (which may be slow) is required.",
          inMemoryElementCount);
      valueMap = new ArrayList<>();
      Function<Integer, Iterable<Object>> makeIterable;
      if (taggedIter instanceof Reiterator) {
        final Reiterator<RawUnionValue> tail = (Reiterator<RawUnionValue>) taggedIter;
  
        // As we iterate over this re-iterable (e.g. while iterating for one tag) we populate values
        // for other observed tags, if any.
        ObservingReiterator<RawUnionValue> tip =
            new ObservingReiterator<>(
                tail,
                new ObservingReiterator.Observer<RawUnionValue>() {
                  @Override
                  public void observeAt(ObservingReiterator<RawUnionValue> reiterator) {
                    ((TagIterable<?>) valueMap.get(reiterator.peek().getUnionTag()))
                        .offer(reiterator);
                  }
  
                  @Override
                  public void done() {
                    // Inform all tags that we have reached the end of the iterable, so anything that
                    // can be observed has been observed.
                    for (Iterable<?> iter : valueMap) {
                      ((TagIterable<?>) iter).finish();
                    }
                  }
                });
        makeIterable =
            unionTag ->
                new TagIterable<Object>(valuesByTag.get(unionTag), unionTag, minElementsPerTag, tip);
      } else {
        // Not reiterable, we have to filter each time, but there are some optimizations we can do...
        boolean[] sharedSeenEnd = {false};
        makeIterable =
            unionTag ->
                recordingFilteringIterable(
                    taggedValues,
                    unionTag,
                    Math.max(
                        inMemoryElementCount / (schema.size() * schema.size()), minElementsPerTag),
                    valueMap,
                    sharedSeenEnd);
      }
  
      for (int unionTag = 0; unionTag < schema.size(); unionTag++) {
        valueMap.add(makeIterable.apply(unionTag));
      }
    }
  
    public boolean isEmpty() {
      for (Iterable<?> tagValues : valueMap) {
        if (tagValues.iterator().hasNext()) {
          return false;
        }
      }
      return true;
    }
  
    /** Returns the schema used by this {@link CoGbkResult}. */
    public CoGbkResultSchema getSchema() {
      return schema;
    }
  
    @Override
    public String toString() {
      return valueMap.toString();
    }
  
    /**
     * Returns the values from the table represented by the given {@code TupleTag<V>} as an {@code
     * Iterable<V>} (which may be empty if there are no results).
     *
     * <p>If tag was not part of the original {@link CoGroupByKey}, throws an
     * IllegalArgumentException.
     */
    public <V> Iterable<V> getAll(TupleTag<V> tag) {
      int index = schema.getIndex(tag);
      if (index < 0) {
        throw new IllegalArgumentException("TupleTag " + tag + " is not in the schema");
      }
      @SuppressWarnings("unchecked")
      Iterable<V> unions = (Iterable<V>) valueMap.get(index);
      return unions;
    }
  
    /** Like {@link #getAll(TupleTag)} but using a String instead of a {@link TupleTag}. */
    public <V> Iterable<V> getAll(String tag) {
      return getAll(new TupleTag<>(tag));
    }
  
    /**
     * If there is a singleton value for the given tag, returns it. Otherwise, throws an
     * IllegalArgumentException.
     *
     * <p>If tag was not part of the original {@link CoGroupByKey}, throws an
     * IllegalArgumentException.
     */
    public <V> V getOnly(TupleTag<V> tag) {
      return innerGetOnly(tag, null, false);
    }
  
    /** Like {@link #getOnly(TupleTag)} but using a String instead of a TupleTag. */
    @SuppressWarnings("TypeParameterUnusedInFormals")
    public <V> V getOnly(String tag) {
      return getOnly(new TupleTag<>(tag));
    }
  
    /**
     * If there is a singleton value for the given tag, returns it. If there is no value for the given
     * tag, returns the defaultValue.
     *
     * <p>If tag was not part of the original {@link CoGroupByKey}, throws an
     * IllegalArgumentException.
     */
    public @Nullable <V> V getOnly(TupleTag<V> tag, @Nullable V defaultValue) {
      return innerGetOnly(tag, defaultValue, true);
    }
  
    /** Like {@link #getOnly(TupleTag, Object)} but using a String instead of a TupleTag. */
    public @Nullable <V> V getOnly(String tag, @Nullable V defaultValue) {
      return getOnly(new TupleTag<>(tag), defaultValue);
    }
  
    /** A {@link Coder} for {@link CoGbkResult}s. */
    public static class CoGbkResultCoder extends CustomCoder<CoGbkResult> {
  
      private final CoGbkResultSchema schema;
      private final UnionCoder unionCoder;
  
      /** Returns a {@link CoGbkResultCoder} for the given schema and {@link UnionCoder}. */
      public static CoGbkResultCoder of(CoGbkResultSchema schema, UnionCoder unionCoder) {
        return new CoGbkResultCoder(schema, unionCoder);
      }
  
      private CoGbkResultCoder(CoGbkResultSchema tupleTags, UnionCoder unionCoder) {
        this.schema = tupleTags;
        this.unionCoder = unionCoder;
      }
  
      @Override
      public List<? extends Coder<?>> getCoderArguments() {
        return ImmutableList.of(unionCoder);
      }
  
      public CoGbkResultSchema getSchema() {
        return schema;
      }
  
      public UnionCoder getUnionCoder() {
        return unionCoder;
      }
  
      @Override
      @SuppressWarnings("unchecked")
      public void encode(CoGbkResult value, OutputStream outStream)
          throws CoderException, IOException {
        if (!schema.equals(value.getSchema())) {
          throw new CoderException("input schema does not match coder schema");
        }
        if (schema.size() == 0) {
          return;
        }
        for (int unionTag = 0; unionTag < schema.size(); unionTag++) {
          tagListCoder(unionTag).encode(value.valueMap.get(unionTag), outStream);
        }
      }
  
      @Override
      public CoGbkResult decode(InputStream inStream) throws CoderException, IOException {
        if (schema.size() == 0) {
          return new CoGbkResult(schema, ImmutableList.<Iterable<?>>of());
        }
        List<Iterable<?>> valueMap = Lists.newArrayListWithExpectedSize(schema.size());
        for (int unionTag = 0; unionTag < schema.size(); unionTag++) {
          valueMap.add(tagListCoder(unionTag).decode(inStream));
        }
        return new CoGbkResult(schema, valueMap);
      }
  
      @SuppressWarnings("rawtypes")
      private IterableCoder tagListCoder(int unionTag) {
        return IterableCoder.of(unionCoder.getElementCoders().get(unionTag));
      }
  
      @Override
      public boolean equals(@Nullable Object object) {
        if (this == object) {
          return true;
        }
        if (!(object instanceof CoGbkResultCoder)) {
          return false;
        }
        CoGbkResultCoder other = (CoGbkResultCoder) object;
        return schema.equals(other.schema) && unionCoder.equals(other.unionCoder);
      }
  
      @Override
      public int hashCode() {
        return Objects.hashCode(schema);
      }
  
      @Override
      public void verifyDeterministic() throws NonDeterministicException {
        verifyDeterministic(
            this, "CoGbkResult requires the union coder to be deterministic", unionCoder);
      }
    }
  
    //////////////////////////////////////////////////////////////////////////////
    // Methods for directly constructing a CoGbkResult
    //
    // (for example, creating test data for a transform that consumes a
    // CoGbkResult)
  
    /** Returns a new CoGbkResult that contains just the given tag and given data. */
    public static <V> CoGbkResult of(TupleTag<V> tag, List<V> data) {
      return CoGbkResult.empty().and(tag, data);
    }
  
    /**
     * Returns a new {@link CoGbkResult} based on this, with the given tag and given data added to it.
     */
    public <V> CoGbkResult and(TupleTag<V> tag, List<V> data) {
      if (nextTestUnionId != schema.size()) {
        throw new IllegalArgumentException(
            "Attempting to call and() on a CoGbkResult apparently not created by" + " of().");
      }
      List<Iterable<?>> valueMap = new ArrayList<>(this.valueMap);
      valueMap.add(data);
      return new CoGbkResult(
          new CoGbkResultSchema(schema.getTupleTagList().and(tag)), valueMap, nextTestUnionId + 1);
    }
  
    /** Returns an empty {@link CoGbkResult}. */
    public static CoGbkResult empty() {
      return new CoGbkResult(
          new CoGbkResultSchema(TupleTagList.empty()), new ArrayList<Iterable<?>>());
    }
  
    //////////////////////////////////////////////////////////////////////////////
  
    private int nextTestUnionId = 0;
  
    private CoGbkResult(CoGbkResultSchema schema, List<Iterable<?>> valueMap, int nextTestUnionId) {
      this(schema, valueMap);
      this.nextTestUnionId = nextTestUnionId;
    }
  
    private CoGbkResult(CoGbkResultSchema schema, List<Iterable<?>> valueMap) {
      this.schema = schema;
      this.valueMap = valueMap;
    }
  
    private @Nullable <V> V innerGetOnly(
        TupleTag<V> tag, @Nullable V defaultValue, boolean useDefault) {
      int index = schema.getIndex(tag);
      if (index < 0) {
        throw new IllegalArgumentException("TupleTag " + tag + " is not in the schema");
      }
      @SuppressWarnings("unchecked")
      Iterator<V> unions = (Iterator<V>) valueMap.get(index).iterator();
      if (!unions.hasNext()) {
        if (useDefault) {
          return defaultValue;
        } else {
          throw new IllegalArgumentException(
              "TupleTag " + tag + " corresponds to an empty result, and no default was provided");
        }
      }
      V value = unions.next();
      if (unions.hasNext()) {
        throw new IllegalArgumentException(
            "TupleTag " + tag + " corresponds to a non-singleton result");
      }
      return value;
    }
  
    /**
     * A re-iterable that notifies an observer at every advance, and upon finishing, but only once
     * across all copies.
     *
     * @param <T> The value type of the underlying iterable.
     */
    private static class ObservingReiterator<T> implements Reiterator<T> {
  
      public interface Observer<T> {
        /**
         * Called exactly once, across all copies before advancing this iterator.
         *
         * <p>The iterator rather than the element is given so that the callee can perform a copy if
         * desired. This class offers a peek method to get at the current element without disturbing
         * the state of this iterator.
         */
        void observeAt(ObservingReiterator<T> reiterator);
  
        /** Called exactly once, across all copies, once this iterator is exhausted. */
        void done();
      }
  
      private PeekingReiterator<IndexingReiterator.Indexed<T>> underlying;
      private Observer<T> observer;
  
      // Used to keep track of what has been observed so far.
      // These are arrays to facilitate sharing values among all copies of the same root Reiterator.
      private final int[] lastObserved;
      private final boolean[] doneHasRun;
      private final PeekingReiterator<IndexingReiterator.Indexed<T>>[] mostAdvanced;
  
      public ObservingReiterator(Reiterator<T> underlying, Observer<T> observer) {
        this(new PeekingReiterator<>(new IndexingReiterator<>(underlying)), observer);
      }
  
      @SuppressWarnings("rawtypes") // array creation
      public ObservingReiterator(
          PeekingReiterator<IndexingReiterator.Indexed<T>> underlying, Observer<T> observer) {
        this(
            underlying,
            observer,
            new int[] {-1},
            new boolean[] {false},
            new PeekingReiterator[] {underlying});
      }
  
      private ObservingReiterator(
          PeekingReiterator<IndexingReiterator.Indexed<T>> underlying,
          Observer<T> observer,
          int[] lastObserved,
          boolean[] doneHasRun,
          PeekingReiterator<IndexingReiterator.Indexed<T>>[] mostAdvanced) {
        this.underlying = underlying;
        this.observer = observer;
        this.lastObserved = lastObserved;
        this.doneHasRun = doneHasRun;
        this.mostAdvanced = mostAdvanced;
      }
  
      @Override
      public Reiterator<T> copy() {
        return new ObservingReiterator<T>(
            underlying.copy(), observer, lastObserved, doneHasRun, mostAdvanced);
      }
  
      @Override
      public boolean hasNext() {
        boolean hasNext = underlying.hasNext();
        if (!hasNext && !doneHasRun[0]) {
          mostAdvanced[0] = underlying;
          observer.done();
          doneHasRun[0] = true;
        }
        return hasNext;
      }
  
      @Override
      public T next() {
        peek(); // trigger observation *before* advancing
        return underlying.next().value;
      }
  
      public T peek() {
        IndexingReiterator.Indexed<T> next = underlying.peek();
        if (next.index > lastObserved[0]) {
          assert next.index == lastObserved[0] + 1;
          mostAdvanced[0] = underlying;
          lastObserved[0] = next.index;
          observer.observeAt(this);
        }
        return next.value;
      }
  
      public void fastForward() {
        if (underlying != mostAdvanced[0]) {
          underlying = mostAdvanced[0].copy();
        }
      }
    }
  
    /**
     * Assigns a monotonically increasing index to each item in the underling Reiterator.
     *
     * @param <T> The value type of the underlying iterable.
     */
    private static class IndexingReiterator<T> implements Reiterator<IndexingReiterator.Indexed<T>> {
  
      private Reiterator<T> underlying;
      private int index;
  
      public IndexingReiterator(Reiterator<T> underlying) {
        this(underlying, 0);
      }
  
      public IndexingReiterator(Reiterator<T> underlying, int start) {
        this.underlying = underlying;
        this.index = start;
      }
  
      @Override
      public IndexingReiterator<T> copy() {
        return new IndexingReiterator<>(underlying.copy(), index);
      }
  
      @Override
      public boolean hasNext() {
        return underlying.hasNext();
      }
  
      @Override
      public Indexed<T> next() {
        return new Indexed<T>(index++, underlying.next());
      }
  
      public static class Indexed<T> {
        public final int index;
        public final T value;
  
        public Indexed(int index, T value) {
          this.index = index;
          this.value = value;
        }
      }
    }
  
    /**
     * Adapts an Reiterator, giving it a peek() method that can be used to observe the next element
     * without consuming it.
     *
     * @param <T> The value type of the underlying iterable.
     */
    private static class PeekingReiterator<T> implements Reiterator<T> {
      private Reiterator<T> underlying;
      private T next;
      private boolean nextIsValid;
  
      public PeekingReiterator(Reiterator<T> underlying) {
        this(underlying, null, false);
      }
  
      private PeekingReiterator(Reiterator<T> underlying, T next, boolean nextIsValid) {
        this.underlying = underlying;
        this.next = next;
        this.nextIsValid = nextIsValid;
      }
  
      @Override
      public PeekingReiterator<T> copy() {
        return new PeekingReiterator<>(underlying.copy(), next, nextIsValid);
      }
  
      @Override
      public boolean hasNext() {
        return nextIsValid || underlying.hasNext();
      }
  
      @Override
      public T next() {
        if (nextIsValid) {
          nextIsValid = false;
          return next;
        } else {
          return underlying.next();
        }
      }
  
      public T peek() {
        if (!nextIsValid) {
          next = underlying.next();
          nextIsValid = true;
        }
        return next;
      }
    }
  
    /**
     * An Iterable corresponding to a single tag.
     *
     * <p>The values in this iterable are populated lazily via the offer method as tip advances for
     * any tag.
     *
     * @param <T> The value type of the corresponding tag.
     */
    private static class TagIterable<T> implements Iterable<T> {
      int tag;
      int cacheSize;
  
      ObservingReiterator<RawUnionValue> tip;
  
      List<T> head;
      Reiterator<RawUnionValue> tail;
      boolean finished;
  
      public TagIterable(
          List<T> head, int tag, int cacheSize, ObservingReiterator<RawUnionValue> tip) {
        this.tag = tag;
        this.cacheSize = cacheSize;
        this.head = head;
        this.tip = tip;
      }
  
      void offer(ObservingReiterator<RawUnionValue> tail) {
        assert !finished;
        assert tail.peek().getUnionTag() == tag;
        if (head.size() < cacheSize) {
          head.add((T) tail.peek().getValue());
        } else if (this.tail == null) {
          this.tail = tail.copy();
        }
      }
  
      void finish() {
        Metrics.counter(
                CoGbkResult.class,
                this.tail == null ? "cogbk-small-iterables" : "cogbk-large-iterables")
            .inc();
        finished = true;
      }
  
      void seek(int tag) {
        while (tip.hasNext() && tip.peek().getUnionTag() != tag) {
          tip.next();
        }
      }
  
      @Override
      public Iterator<T> iterator() {
        return new Iterator<T>() {
  
          boolean isDone;
          boolean advanced;
          T next;
  
          /** Keeps track of the index, in head, that this iterator points to. */
          int index = -1;
          /** If the index is beyond what was cached in head, this is this iterators view of tail. */
          Iterator<T> tailIter;
  
          @Override
          public boolean hasNext() {
            if (!advanced) {
              advance();
            }
            return !isDone;
          }
  
          @Override
          public T next() {
            if (!advanced) {
              advance();
            }
            if (isDone) {
              throw new NoSuchElementException();
            }
            advanced = false;
            return next;
          }
  
          // Invariant: After advanced is called, either isDone is true or next is valid.
          private void advance() {
            assert !advanced;
            assert !isDone;
            advanced = true;
  
            index++;
            if (maybeAdvance()) {
              return;
            }
  
            // We were unable to advance; advance tip to populate either head or tail.
            tip.fastForward();
            if (tip.hasNext()) {
              tip.next();
              seek(tag);
            }
  
            // A this point, either head or tail should be sufficient to advance.
            Preconditions.checkState(maybeAdvance());
          }
  
          // Invariant: If returns true, either isDone is true or next is valid.
          private boolean maybeAdvance() {
            if (index < head.size()) {
              // First consume head.
              assert tailIter == null;
              next = head.get(index);
              return true;
            } else if (tail != null) {
              // Next consume tail, if any.
              if (tailIter == null) {
                tailIter =
                    Iterators.transform(
                        Iterators.filter(
                            tail.copy(), taggedUnion -> taggedUnion.getUnionTag() == tag),
                        taggedUnion -> (T) taggedUnion.getValue());
              }
              if (tailIter.hasNext()) {
                next = tailIter.next();
              } else {
                isDone = true;
              }
              return true;
            } else if (finished) {
              // If there are no more elements in head, and tail was not populated, and we are
              // finished, this is the end of the iteration.
              isDone = true;
              return true;
            } else {
              // We need more elements in either head or tail.
              return false;
            }
          }
        };
      }
    }
  
    private Iterable<Object> recordingFilteringIterable(
        Iterable<RawUnionValue> taggedIteratable,
        int unionTag,
        int minElementsPerTag,
        List<Iterable<?>> sharedValueMap,
        boolean[] sharedSeenEnd) {
      return () ->
          new RecordingFilteringIterator(
              taggedIteratable, unionTag, minElementsPerTag, sharedValueMap, sharedSeenEnd);
    }
  
    /**
     * This iterator implements the optimization that if there are below a certain number of values
     * for a given tag we cache those values in memory rather than reiterating and filtering each
     * time.
     *
     * <p>This is done lazily by having each iterator keep track of what it has seen locally, and the
     * first one to reach the end updates the shared map. As an added optimization, this iterable is
     * also updated in place if it only had a small number of elements.
     *
     * <p>Unfortuantely iterators do not promise deterministic ordering, so we cannot share the work
     * computing the local maps until the iterator is entirely exhausted.
     */
    private static class RecordingFilteringIterator extends AbstractIterator<Object> {
      private final Iterable<RawUnionValue> taggedIterable;
      private final Iterator<RawUnionValue> taggedIterator;
      private final int unionTag;
  
      private final int minElementsPerTag;
  
      private final List<List<Object>> localValueMap;
      private final List<Iterable<?>> sharedValueMap;
      private boolean[] sharedSeenEnd;
  
      private enum RemainingStatus {
        UNCOMPUTED,
        UNCOMPUTABLE,
        COMPUTED
      }
  
      private RemainingStatus remainingStatus = RemainingStatus.UNCOMPUTED;
      private Iterator<Object> remaining;
  
      public RecordingFilteringIterator(
          Iterable<RawUnionValue> taggedIteratable,
          int unionTag,
          int minElementsPerTag,
          List<Iterable<?>> sharedValueMap,
          boolean[] sharedSeenEnd) {
        this.taggedIterable = taggedIteratable;
        this.taggedIterator = taggedIteratable.iterator();
        this.unionTag = unionTag;
        this.minElementsPerTag = minElementsPerTag;
        localValueMap = new ArrayList<>();
        for (int i = 0; i < sharedValueMap.size(); i++) {
          localValueMap.add(new ArrayList<>());
        }
        this.sharedValueMap = sharedValueMap;
        this.sharedSeenEnd = sharedSeenEnd;
      }
  
      @Override
      protected Object computeNext() {
        if (sharedSeenEnd[0]) {
          if (remainingStatus == RemainingStatus.UNCOMPUTED) {
            // We iterated all the way to the end (likely on another iterator) and may have only found
            // a small number of values associated with this tag. Update this iterator, if possible,
            // to return them directly.
            remainingStatus = computeRemaining(sharedValueMap.get(unionTag));
          }
  
          if (remainingStatus == RemainingStatus.COMPUTED) {
            if (remaining.hasNext()) {
              return remaining.next();
            } else {
              return endOfData();
            }
          }
        }
  
        // Look for the next value with this tag, keeping track of up to minElementsPerTag values of
        // all other iterables as we go.
        while (taggedIterator.hasNext()) {
          RawUnionValue unionValue = taggedIterator.next();
          if (!sharedSeenEnd[0]) {
            List<Object> valuesForTag = localValueMap.get(unionValue.getUnionTag());
            if (valuesForTag != null) {
              if (valuesForTag.size() < minElementsPerTag) {
                valuesForTag.add(unionValue.getValue());
              } else {
                localValueMap.set(unionValue.getUnionTag(), null);
              }
            }
          }
          if (unionValue.getUnionTag() == unionTag) {
            return unionValue.getValue();
          }
        }
  
        // We got to the end of the iterable, update the shared set of values with those sets that
        // were small enough to cache.
        if (!sharedSeenEnd[0]) {
          Counter smallIterablesCount = Metrics.counter(CoGbkResult.class, "cogbk-small-iterables");
          Counter largeIterablesCount = Metrics.counter(CoGbkResult.class, "cogbk-large-iterables");
          for (int i = 0; i < sharedValueMap.size(); i++) {
            List<Object> localValues = localValueMap.get(i);
            (localValues == null ? largeIterablesCount : smallIterablesCount).inc();
            sharedValueMap.set(
                i, localValues != null ? localValues : simpleFilteringIterable(taggedIterable, i));
          }
          sharedSeenEnd[0] = true;
        }
  
        return endOfData();
      }
  
      private RemainingStatus computeRemaining(Iterable<?> allValuesIter) {
        if (!(allValuesIter instanceof Collection)) {
          return RemainingStatus.UNCOMPUTABLE;
        }
        Collection<Object> allValues = (Collection<Object>) allValuesIter;
        List<Object> seenValues = localValueMap.get(unionTag);
        if (allValues.size() == seenValues.size()) {
          remaining = Collections.emptyIterator();
          return RemainingStatus.COMPUTED;
        } else if (seenValues.size() == 0) {
          remaining = allValues.iterator();
          return RemainingStatus.COMPUTED;
        } else if (seenValues.size() == 1) {
          // Optimize the very common case.
          Iterator<Object> iter = allValues.iterator();
          if (Objects.equals(iter.next(), seenValues.get(0))) {
            remaining = iter;
            return RemainingStatus.COMPUTED;
          } else {
            ArrayList<Object> allButOne = Lists.newArrayList(allValues);
            if (allButOne.remove(seenValues.get(0))) {
              remaining = allButOne.iterator();
              return RemainingStatus.COMPUTED;
            } else {
              return RemainingStatus.UNCOMPUTABLE;
            }
          }
        } else {
          try {
            HashMultiset<Object> seenValueSet = HashMultiset.create(seenValues);
            List<Object> unseenValues = new ArrayList<>();
            for (Object value : allValues) {
              if (!seenValueSet.remove(value)) {
                unseenValues.add(value);
              }
            }
            if (seenValueSet.isEmpty()) {
              remaining = unseenValues.iterator();
              return RemainingStatus.COMPUTED;
            } else {
              // Semantically equal values didn't hash or compare equal.
              return RemainingStatus.UNCOMPUTABLE;
            }
          } catch (Exception exn) {
            // There's no promise elements have correct hash semantics or properly handle nulls.
            return RemainingStatus.UNCOMPUTABLE;
          }
        }
      }
    }
  
    private static Iterable<Object> simpleFilteringIterable(
        Iterable<RawUnionValue> taggedIterable, int unionTag) {
      return () ->
          new AbstractIterator<Object>() {
            Iterator<RawUnionValue> taggedIterator = taggedIterable.iterator();
  
            @Override
            protected Object computeNext() {
              while (taggedIterator.hasNext()) {
                RawUnionValue unionValue = taggedIterator.next();
                if (unionValue.getUnionTag() == unionTag) {
                  return unionValue.getValue();
                }
              }
              return endOfData();
            }
          };
    }
  }
  