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
   *
   * <p>http://www.apache.org/licenses/LICENSE-2.0
   *
   * <p>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.tez.runtime.library.input;
  
  import java.io.IOException;
  import java.util.*;
  
  import org.apache.hadoop.classification.InterfaceAudience.Public;
  import org.apache.hadoop.io.RawComparator;
  import org.apache.tez.runtime.api.Input;
  import org.apache.tez.runtime.api.MergedInputContext;
  import org.apache.tez.runtime.api.MergedLogicalInput;
  import org.apache.tez.runtime.api.ProgressFailedException;
  import org.apache.tez.runtime.library.api.KeyValuesReader;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * A {@link MergedLogicalInput} which merges multiple {@link OrderedGroupedKVInput}s and returns a
   * single view of these by merging values which belong to the same key.
   *
   * <p>Combiners and Secondary Sort are not implemented, so there is no guarantee on the order of
   * values.
   */
  @Public
  public class CelebornOrderedGroupedMergedKVInput extends MergedLogicalInput {
  
    private static final Logger LOG =
        LoggerFactory.getLogger(CelebornOrderedGroupedMergedKVInput.class);
    private final Set<Input> completedInputs =
        Collections.newSetFromMap(new IdentityHashMap<Input, Boolean>());
  
    public CelebornOrderedGroupedMergedKVInput(MergedInputContext context, List<Input> inputs) {
      super(context, inputs);
    }
  
    /** Provides an ordered {@link KeyValuesReader} */
    @Override
    public KeyValuesReader getReader() throws Exception {
      return new OrderedGroupedMergedKeyValuesReader(getInputs(), getContext());
    }
  
    @Override
    public void setConstituentInputIsReady(Input input) {
      synchronized (completedInputs) {
        completedInputs.add(input);
        if (completedInputs.size() == getInputs().size()) {
          informInputReady();
        }
      }
    }
  
    private static class OrderedGroupedMergedKeyValuesReader extends KeyValuesReader {
      private final PriorityQueue<KeyValuesReader> pQueue;
  
      @SuppressWarnings("rawtypes")
      private final RawComparator keyComparator;
  
      private final List<KeyValuesReader> finishedReaders;
      private final ValuesIterable currentValues;
      private KeyValuesReader nextKVReader;
      private Object currentKey;
      private final MergedInputContext context;
  
      public OrderedGroupedMergedKeyValuesReader(List<Input> inputs, MergedInputContext context)
          throws Exception {
        keyComparator = ((OrderedGroupedKVInput) inputs.get(0)).getInputKeyComparator();
        pQueue =
            new PriorityQueue<KeyValuesReader>(inputs.size(), new KVReaderComparator(keyComparator));
        finishedReaders = new ArrayList<KeyValuesReader>(inputs.size());
        for (Input input : inputs) {
          KeyValuesReader reader = (KeyValuesReader) input.getReader();
          if (reader.next()) {
            pQueue.add(reader);
          }
        }
        currentValues = new ValuesIterable();
        this.context = context;
      }
  
      private void advanceAndAddToQueue(KeyValuesReader kvsReadr) throws IOException {
        if (kvsReadr.next()) {
          pQueue.add(kvsReadr);
        }
      }
  
      private void addToQueue(KeyValuesReader kvsReadr) throws IOException {
        if (kvsReadr != null) {
          pQueue.add(kvsReadr);
        }
      }
  
      @Override
      public boolean next() throws IOException {
        // Skip values of current key if not consumed by the user
        currentValues.discardCurrent();
  
        for (KeyValuesReader reader : finishedReaders) {
          // add them back to queue
          advanceAndAddToQueue(reader);
        }
        finishedReaders.clear();
  
        nextKVReader = pQueue.poll();
        context.notifyProgress();
        if (nextKVReader != null) {
          currentKey = nextKVReader.getCurrentKey();
          currentValues.moveToNext();
          return true;
        } else {
          hasCompletedProcessing();
          completedProcessing = true;
        }
        return false;
      }
  
      @Override
      public Object getCurrentKey() throws IOException {
        return currentKey;
      }
  
      @Override
      public Iterable<Object> getCurrentValues() throws IOException {
        return currentValues;
      }
  
      private class ValuesIterable implements Iterable<Object> {
        private ValuesIterator iterator = new ValuesIterator();
  
        @Override
        public Iterator<Object> iterator() {
          return iterator;
        }
  
        public void discardCurrent() throws IOException {
          iterator.discardCurrent();
        }
  
        public void moveToNext() throws IOException {
          iterator.moveToNext();
        }
      }
  
      @SuppressWarnings("unchecked")
      private class ValuesIterator implements Iterator<Object> {
  
        private Iterator<Object> currentValuesIter;
  
        public void moveToNext() throws IOException {
          currentValuesIter = nextKVReader.getCurrentValues().iterator();
        }
  
        @Override
        public boolean hasNext() {
          if (currentValuesIter != null) { // No current key. next needs to be called.
            if (currentValuesIter.hasNext()) {
              return true;
            } else {
              finishedReaders.add(nextKVReader);
              nextKVReader = pQueue.poll();
              try {
                if (nextKVReader != null
                    && keyComparator.compare(currentKey, nextKVReader.getCurrentKey()) == 0) {
                  currentValuesIter = nextKVReader.getCurrentValues().iterator();
                  return true;
                } else { // key changed or no more data.
                  // Add the reader back to queue
                  addToQueue(nextKVReader);
                  currentValuesIter = null;
                  return false;
                }
              } catch (IOException e) {
                throw new RuntimeException(e);
              }
            }
          } else {
            return false;
          }
        }
  
        public void discardCurrent() throws IOException {
          if (currentValuesIter != null) {
            do {
              finishedReaders.add(nextKVReader);
              nextKVReader = pQueue.poll();
            } while (nextKVReader != null
                && keyComparator.compare(currentKey, nextKVReader.getCurrentKey()) == 0);
            addToQueue(nextKVReader);
            currentValuesIter = null;
          }
        }
  
        @Override
        public Object next() {
          return currentValuesIter.next();
        }
  
        @Override
        public void remove() {
          throw new UnsupportedOperationException();
        }
      }
  
      /** Comparator that compares KeyValuesReader on their current key */
      @SuppressWarnings({"rawtypes", "unchecked"})
      private static class KVReaderComparator implements Comparator<KeyValuesReader> {
  
        private RawComparator keyComparator;
  
        public KVReaderComparator(RawComparator keyComparator) {
          this.keyComparator = keyComparator;
        }
  
        @Override
        public int compare(KeyValuesReader o1, KeyValuesReader o2) {
          try {
            return keyComparator.compare(o1.getCurrentKey(), o2.getCurrentKey());
          } catch (IOException e) {
            LOG.error("Caught exception while comparing keys in shuffle input", e);
            throw new RuntimeException(e);
          }
        }
      }
    }
  
    public float getProgress() throws ProgressFailedException, InterruptedException {
      float totalProgress = 0.0f;
      for (Input input : getInputs()) {
        totalProgress += ((OrderedGroupedKVInput) input).getProgress();
      }
      return (1.0f) * totalProgress / getInputs().size();
    }
  }
  