TaskOutput.java

/*
 * 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 com.aliyun.odps.mapred.unittest;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import com.aliyun.odps.data.Record;
import com.aliyun.odps.mapred.Mapper;
import com.aliyun.odps.mapred.Reducer;
import com.aliyun.odps.mapred.conf.JobConf;
import com.aliyun.odps.mapred.local.MapOutputBuffer;
import com.aliyun.odps.pipeline.Pipeline;
import com.aliyun.odps.counter.Counters;

/**
 * Mapper/Reducer的输出.
 */
public class TaskOutput extends MapOutputBuffer {

  public TaskOutput(JobConf conf, int reduceNum) {
    super(conf, reduceNum);
  }

  public TaskOutput(JobConf job, Pipeline pipeline, String taskId, int reduceCopyNum) {
    super(job, pipeline, taskId, reduceCopyNum);
  }

  private Counters counters;

  private Map<Integer, List<KeyValue<Record, Record>>> outputKeyValues =
      new LinkedHashMap<Integer, List<KeyValue<Record, Record>>>();
  private Map<String, List<Record>> outputs = new HashMap<String, List<Record>>();


  /**
   * 返回 {@link Mapper} 最终输出到 {@link Reducer} 的键值对列表.
   * 
   * @return {@link Mapper} 输出到 {@link Reducer} 的键值对列表
   */
  public List<KeyValue<Record, Record>> getOutputKeyValues() {
    List<KeyValue<Record, Record>> list = new ArrayList<KeyValue<Record, Record>>();
    for (Map.Entry<Integer, List<KeyValue<Record, Record>>> entry : outputKeyValues
        .entrySet()) {
      list.addAll(entry.getValue());
    }
    return list;
  }

  /**
   * 返回 {@link Mapper} 输出到给定 reduceId 的键值对列表.
   * 
   * @param reduceId
   *          {@link Reducer} 序号，从 0 开始计数。
   * @return 输出到给定 reduceId 的键值对列表
   */
  public List<KeyValue<Record, Record>> getOutputKeyValues(int reduceId) {
    return outputKeyValues.get(reduceId);
  }

  void setOutputKeyValues(int reduceId, List<KeyValue<Record, Record>> outputKeyValues) {
    this.outputKeyValues.put(reduceId, outputKeyValues);
  }

  /**
   * 获取默认输出结果，返回 List<{@link Record}>.
   * 
   * @return 默认输出结果
   */
  public List<Record> getOutputRecords() {
    return getOutputRecords("__default__");
  }

  /**
   * 获取给定标签的输出结果，返回 List<{@link Record}>.
   * 
   * @param label
   *          输出标签
   * @return 给定标签的输出结果
   */
  public List<Record> getOutputRecords(String label) {
    return getOutputRecords(label, true);
  }

  /**
   * 获取默认输出结果，返回 List<{@link Record}>.
   * 
   * @return 默认输出结果
   */
  public List<Record> getOutputRecords(boolean sort) {
    return getOutputRecords("__default__", sort);
  }

  /**
   * 获取给定标签的输出结果，返回 List<{@link Record}>.
   * 
   * @param label
   *          输出标签
   * @return 给定标签的输出结果
   */
  public List<Record> getOutputRecords(String label, boolean sort) {
    List<Record> records = outputs.get(label);
    if (records == null) {
      records = new ArrayList<Record>();
    }
    List<Record> sortRecords = new ArrayList<Record>(records);
    LocalRecordComparator comparator = new LocalRecordComparator();
    if (sort) {
      Collections.sort(sortRecords, comparator);
    }
    return sortRecords;
  }

  /**
   * 获取作业 {@link Counters}.
   * 
   * @return 作业 {@link Counters}.
   */
  public Counters getCounters() {
    return counters;
  }

  void setCounters(Counters counters) {
    this.counters = counters;
  }

  void setOutputRecords(String label, List<Record> records) {
    outputs.put(label, records);
  }

  @Override
  public void add(Record key, Record value) {
    int partition = getPartition(key);
    this.add(key, value, partition);
  }

  @Override
  public void add(Record key, Record value, int partition) {
    super.add(key, value, partition);
    KeyValue<Record, Record> kv = new KeyValue<Record, Record>(key.clone(), value.clone());

    List<KeyValue<Record, Record>> kvs = this.getOutputKeyValues(partition);
    if (kvs == null) {
      kvs = new ArrayList<KeyValue<Record, Record>>();
      outputKeyValues.put(partition, kvs);
    }
    kvs.add(kv);
  }

  @Override
  public void add(Record record, String label) {
    List<Record> records = outputs.get(label);
    if (records == null) {
      records = new ArrayList<Record>();
      outputs.put(label, records);
    }
    records.add(record.clone());
  }
  
  @Override
  public long getTotalRecordCount() {
    long numOutputRecords = 0;
    for (List<Record> records : outputs.values()) {
      numOutputRecords += records.size();
    }
    return numOutputRecords + super.getTotalRecordCount();
  }

}