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.hadoop.mapred;
  
  import java.io.IOException;
  
  import org.apache.hadoop.io.serializer.SerializationFactory;
  import org.apache.hadoop.mapreduce.JobContext;
  import org.apache.hadoop.mapreduce.TaskCounter;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.client.ShuffleClient;
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.identity.UserIdentifier;
  import org.apache.celeborn.util.HadoopUtils;
  
  public class CelebornMapOutputCollector<K extends Object, V extends Object>
      implements MapOutputCollector<K, V> {
    private static final Logger logger = LoggerFactory.getLogger(CelebornMapOutputCollector.class);
    private Class<K> keyClass;
    private Class<V> valClass;
    private Task.TaskReporter reporter;
    private CelebornSortBasedPusher<K, V> celebornSortBasedPusher;
    private int numReducers;
  
    @Override
    public void init(Context context) throws IOException {
      JobConf jobConf = context.getJobConf();
      reporter = context.getReporter();
      keyClass = (Class<K>) jobConf.getMapOutputKeyClass();
      valClass = (Class<V>) jobConf.getMapOutputValueClass();
      context.getMapTask().getTaskID().getId();
      numReducers = jobConf.getNumReduceTasks();
  
      int IOBufferSize = jobConf.getInt(JobContext.IO_SORT_MB, 100);
      // Java bytebuffer cannot be larger than Integer.MAX_VALUE
      if ((IOBufferSize & 0x7FF) != IOBufferSize) {
        throw new IOException("Invalid \"" + JobContext.IO_SORT_MB + "\": " + IOBufferSize);
      }
  
      CelebornConf celebornConf = HadoopUtils.fromYarnConf(jobConf);
      JobConf celebornAppendConf = new JobConf(HadoopUtils.MR_CELEBORN_CONF);
      String lmHost = celebornAppendConf.get(HadoopUtils.MR_CELEBORN_LM_HOST);
      int lmPort = Integer.parseInt(celebornAppendConf.get(HadoopUtils.MR_CELEBORN_LM_PORT));
      String applicationAttemptId = celebornAppendConf.get(HadoopUtils.MR_CELEBORN_APPLICATION_ID);
      logger.info(
          "Mapper initialized with celeborn {} {} {} {}",
          lmHost,
          lmPort,
          applicationAttemptId,
          IOBufferSize);
      UserIdentifier userIdentifier =
          new UserIdentifier(celebornConf.userSpecificTenant(), celebornConf.userSpecificUserName());
  
      final float spiller = jobConf.getFloat(JobContext.MAP_SORT_SPILL_PERCENT, (float) 0.8);
      int pushSize = (int) ((IOBufferSize << 20) * spiller);
  
      SerializationFactory serializationFactory = new SerializationFactory(jobConf);
      celebornSortBasedPusher =
          new CelebornSortBasedPusher<>(
              jobConf.getNumMapTasks(),
              jobConf.getNumReduceTasks(),
              // this is map id
              context.getMapTask().getTaskID().getTaskID().getId(),
              // this is attempt id
              context.getMapTask().getTaskID().getId(),
              serializationFactory.getSerializer(keyClass),
              serializationFactory.getSerializer(valClass),
              IOBufferSize << 20,
              pushSize,
              jobConf.getOutputKeyComparator(),
              reporter.getCounter(TaskCounter.MAP_OUTPUT_BYTES),
              reporter.getCounter(TaskCounter.MAP_OUTPUT_RECORDS),
              ShuffleClient.get(applicationAttemptId, lmHost, lmPort, celebornConf, userIdentifier),
              celebornConf);
    }
  
    @Override
    public void collect(K key, V value, int partition) throws IOException {
      reporter.progress();
      if (key.getClass() != keyClass) {
        throw new IOException(
            "Type mismatch in key from map: expected "
                + keyClass.getName()
                + ", received "
                + key.getClass().getName());
      }
      if (value.getClass() != valClass) {
        throw new IOException(
            "Type mismatch in value from map: expected "
                + valClass.getName()
                + ", received "
                + value.getClass().getName());
      }
      if (partition < 0 || partition >= numReducers) {
        throw new IOException("Illegal partition for " + key + " (" + partition + ")");
      }
      celebornSortBasedPusher.checkException();
      celebornSortBasedPusher.insert(key, value, partition);
    }
  
    @Override
    public void close() throws IOException {
      logger.info("Mapper collector close");
      reporter.progress();
      celebornSortBasedPusher.close();
      celebornSortBasedPusher.checkException();
    }
  
    @Override
    public void flush() throws IOException {
      logger.info("Mapper collector flush");
      celebornSortBasedPusher.flush();
      celebornSortBasedPusher.checkException();
      reporter.progress();
    }
  }
  