package com.twitter.elephantbird.pig.util;

import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.hadoop.conf.Configuration;
import org.apache.pig.LoadFunc;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.DataByteArray;
import org.apache.pig.data.DataType;
import org.apache.pig.data.NonSpillableDataBag;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.logicalLayer.schema.Schema;
import org.apache.pig.impl.logicalLayer.schema.Schema.FieldSchema;
import org.apache.thrift.TBase;
import org.apache.thrift.protocol.TType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Lists;
import com.twitter.elephantbird.pig.load.ThriftPigLoader;
import com.twitter.elephantbird.thrift.TStructDescriptor;
import com.twitter.elephantbird.thrift.TStructDescriptor.Field;
import com.twitter.elephantbird.util.ThriftUtils;
import com.twitter.elephantbird.util.TypeRef;

/**
 * <li> converts a Thrift struct to a Pig tuple
 * <li> utilities to provide schema for Pig loaders and Pig scripts
 */
public class ThriftToPig<M extends TBase<?, ?>> {

  public static final Logger LOG = LoggerFactory.getLogger(ThriftToPig.class);

  static final String USE_ENUM_ID_CONF_KEY = "elephantbird.pig.thrift.load.enum.as.int";

  // static because it is used in toSchema and toPigTuple which are static
  private static Boolean useEnumId = false;
  private static TupleFactory tupleFactory  = TupleFactory.getInstance();

  private TStructDescriptor structDesc;

  public static <M extends TBase<?, ?>> ThriftToPig<M> newInstance(Class<M> tClass) {
    return new ThriftToPig<M>(tClass);
  }

  public static <M extends TBase<?, ?>> ThriftToPig<M> newInstance(TypeRef<M> typeRef) {
    return new ThriftToPig<M>(typeRef.getRawClass());
  }

  public ThriftToPig(Class<M> tClass) {
    structDesc = TStructDescriptor.getInstance(tClass);
  }

  public TStructDescriptor getTStructDescriptor() {
    return structDesc;
  }

  /**
   * Converts a thrift object to Pig tuple.
   * All the fields are deserialized.
   * It might be better to use getLazyTuple() if not all fields
   * are required.
   */
  public Tuple getPigTuple(M thriftObj) {
    return toTuple(structDesc, thriftObj);
  }

  /**
   * Similar to {@link #getPigTuple(TBase)}. This delays
   * serialization of tuple contents until they are requested.
   * @param thriftObj
   * @return
   */
  public Tuple getLazyTuple(M thriftObj) {
    return new LazyTuple(structDesc, thriftObj);
  }

  /**
   * Set the flags that can be used by the conversion.
   * @param conf usually the Hadoop job conf
   */
  public static void setConversionProperties(Configuration conf) {
    if (conf != null) {
      useEnumId = conf.getBoolean(USE_ENUM_ID_CONF_KEY, false);
      LOG.debug("useEnumId is set to " + useEnumId);
    }
  }

  @SuppressWarnings("rawtypes")
  private static <T extends TBase>Tuple toTuple(TStructDescriptor tDesc, T tObj) {
    int size = tDesc.getFields().size();
    Tuple tuple = tupleFactory.newTuple(size);
    for (int i=0; i<size; i++) {
      Field field = tDesc.getFieldAt(i);
      Object value = tDesc.getFieldValue(i, tObj);
      try {
        tuple.set(i, toPigObject(field, value, false));
      } catch (ExecException e) { // not expected
        throw new RuntimeException(e);
      }
    }
    return tuple;
  }

  @SuppressWarnings("unchecked")
  public static Object toPigObject(Field field, Object value, boolean lazy) {
    if (value == null) {
      return null;
    }

    switch (field.getType()) {
    case TType.BOOL:
      return Integer.valueOf((Boolean)value ? 1 : 0);
    case TType.BYTE :
      return Integer.valueOf((Byte)value);
    case TType.I16 :
      return Integer.valueOf((Short)value);
    case TType.STRING:
      return stringTypeToPig(value);
    case TType.STRUCT:
      if (lazy) {
        return new LazyTuple(field.gettStructDescriptor(), (TBase<?, ?>)value);
      } else {
        return toTuple(field.gettStructDescriptor(), (TBase<?, ?>)value);
      }
    case TType.MAP:
      return toPigMap(field, (Map<Object, Object>)value, lazy);
    case TType.SET:
      return toPigBag(field.getSetElemField(), (Collection<Object>)value, lazy);
    case TType.LIST:
      return toPigBag(field.getListElemField(), (Collection<Object>)value, lazy);
    case TType.ENUM:
      if (useEnumId) {
        return field.getEnumValueOf(value.toString()).getValue();
      } else {
        return value.toString();
      }
    default:
      // standard types : I32, I64, DOUBLE, etc.
      return value;
    }
  }

  /**
   * TType.STRING is a mess in Thrift. It could be byte[], ByteBuffer,
   * or even a String!.
   */
  private static Object stringTypeToPig(Object value) {
    if (value instanceof String) {
      return value;
    }
    if (value instanceof byte[]) {
      byte[] buf = (byte[])value;
      return new DataByteArray(Arrays.copyOf(buf, buf.length));
    }
    if (value instanceof ByteBuffer) {
      ByteBuffer bin = (ByteBuffer)value;
      byte[] buf = new byte[bin.remaining()];
      bin.mark();
      bin.get(buf);
      bin.reset();
      return new DataByteArray(buf);
    }
    return null;
  }

  private static Map<String, Object> toPigMap(Field field,
                                              Map<Object, Object> map,
                                              boolean lazy) {
    // PIG map's key always a String. just use toString() and hope
    // things would work out ok.
    HashMap<String, Object> out = new HashMap<String, Object>(map.size());
    Field valueField = field.getMapValueField();
    for(Entry<Object, Object> e : map.entrySet()) {
      String key = e.getKey() == null ? null : e.getKey().toString();
      Object prev = out.put(key, toPigObject(valueField, e.getValue(), lazy));
      if (prev != null) {
        String msg = "Duplicate keys while converting to String while "
          + " processing map " + field.getName() + " (key type : "
          + field.getMapKeyField().getType() + " value type : "
          + field.getMapValueField().getType() + ")";
        LOG.warn(msg);
        throw new RuntimeException(msg);
      }
    }
    return out;
  }

  private static DataBag toPigBag(Field field,
                                  Collection<Object> values,
                                  boolean lazy) {
    List<Tuple> tuples = Lists.newArrayListWithExpectedSize(values.size());
    for(Object value : values) {
      Object pValue = toPigObject(field, value, lazy);
      if (pValue instanceof Tuple) { // DataBag should contain Tuples
        tuples.add((Tuple)pValue);
      } else {
        tuples.add(tupleFactory.newTuple(pValue));
      }
    }
    return new NonSpillableDataBag(tuples);
  }

  @SuppressWarnings("serial")
  /**
   * Delays serialization of Thrift fields until they are requested.
   */
  private static class LazyTuple extends AbstractLazyTuple {
    /* NOTE : This is only a partial optimization. The other part
     * is to avoid deserialization of the Thrift fields from the
     * binary buffer.
     *
     * Currently TDeserializer allows deserializing just one field,
     * psuedo-skipping over the fields before it.
     * But if we are going deserialize 5 fields out of 20, we will be
     * skipping over same set of fields multiple times. OTOH this might
     * still be better than a full deserialization.
     *
     * We need to write our own version of TBinaryProtocol that truly skips.
     * Even TDeserializer 'skips'/ignores only after deserializing fields.
     * (e.g. Strings, Integers, buffers etc).
     */
    private TBase<?, ?> tObject;
    private TStructDescriptor desc;

    LazyTuple(TStructDescriptor desc, TBase<?, ?> tObject) {
      initRealTuple(desc.getFields().size());
      this.tObject = tObject;
      this.desc = desc;
    }

    @Override
    protected Object getObjectAt(int index) {
      Field field = desc.getFieldAt(index);
      return toPigObject(field, desc.getFieldValue(index, tObject), true);
    }
  }

  /**
   * Returns Pig schema for the Thrift struct.
   */
  public static Schema toSchema(Class<? extends TBase<?, ?>> tClass) {
    return toSchema(TStructDescriptor.getInstance(tClass));
  }

  public Schema toSchema() {
    return toSchema(structDesc);
  }

  public static Schema toSchema(TStructDescriptor tDesc) {
    Schema schema = new Schema();

    try {
      for(Field field : tDesc.getFields()) {
        schema.add(singleFieldToFieldSchema(field.getName(), field));
      }
    } catch (FrontendException t) {
      throw new RuntimeException(t);
    }

    return schema;
  }

  /**
   * return {@link FieldSchema} for a given field.
   */
  private static FieldSchema singleFieldToFieldSchema(String fieldName, Field field) throws FrontendException {
    //TODO we should probably implement better naming, the current system is pretty nonsensical now
    switch (field.getType()) {
      case TType.STRUCT:
        return new FieldSchema(fieldName, toSchema(field.gettStructDescriptor()), DataType.TUPLE);
      case TType.LIST:
        return new FieldSchema(fieldName, singleFieldToTupleSchema(fieldName + "_tuple", field.getListElemField()), DataType.BAG);
      case TType.SET:
        return new FieldSchema(fieldName, singleFieldToTupleSchema(fieldName + "_tuple", field.getSetElemField()), DataType.BAG);
      case TType.MAP:
        if (field.getMapKeyField().getType() != TType.STRING
            && field.getMapKeyField().getType() != TType.ENUM) {
          LOG.warn("Using a map with non-string key for field " + field.getName()
              + ". while converting to PIG Tuple, toString() is used for the key."
              + " It could result in incorrect maps.");
        }
        return new FieldSchema(fieldName, new Schema(singleFieldToFieldSchema(null, field.getMapValueField())), DataType.MAP);
      default:
        return new FieldSchema(fieldName, null, getPigDataType(field));
    }
  }

  /**
   * A helper function which wraps a Schema in a tuple (for Pig bags) if our version of pig makes it necessary
   */
  private static Schema wrapInTupleIfPig9(Schema schema) throws FrontendException {
      if (PigUtil.Pig9orNewer) {
          return new Schema(new FieldSchema("t",schema,DataType.TUPLE));
      } else {
          return schema;
      }
  }

  /**
   * Returns a schema with single tuple (for Pig bags).
   */
  private static Schema singleFieldToTupleSchema(String fieldName, Field field) throws FrontendException {
    switch (field.getType()) {
      case TType.STRUCT:
        return wrapInTupleIfPig9(toSchema(field.gettStructDescriptor()));
      case TType.LIST:
      case TType.SET:
      case TType.MAP:
        return wrapInTupleIfPig9(new Schema(singleFieldToFieldSchema(fieldName, field)));
      default:
        return wrapInTupleIfPig9(new Schema(new FieldSchema(fieldName, null, getPigDataType(field))));
    }
  }

  private static byte getPigDataType(Field field) {
    switch (field.getType()) {
      case TType.BOOL:
      case TType.BYTE:
      case TType.I16:
      case TType.I32:
        return DataType.INTEGER;
      case TType.ENUM:
        if (useEnumId) {
          return DataType.INTEGER;
        } else {
          return DataType.CHARARRAY;
        }
      case TType.I64:
        return DataType.LONG;
      case TType.DOUBLE:
        return DataType.DOUBLE;
      case TType.STRING:
        return field.isBuffer() ? DataType.BYTEARRAY : DataType.CHARARRAY;
      default:
        throw new IllegalArgumentException("Unexpected type where a simple type is expected : " + field.getType());
    }
  }

  /**
   * Turn a Thrift Struct into a loading schema for a pig script.
   */
  public static String toPigScript(Class<? extends TBase<?, ?>> thriftClass,
                                   Class<? extends LoadFunc> pigLoader) {
    StringBuilder sb = new StringBuilder();
    /* we are commenting out explicit schema specification. The schema is
     * included mainly to help the readers of the pig script. Pig learns the
     * schema directly from the loader.
     * If explicit schema is not commented, we might have surprising results
     * when a Thrift class (possibly in control of another team) changes,
     * but the Pig script is not updated. Commenting it out avoids this.
     */
    StringBuilder prefix = new StringBuilder("       --  ");
    sb.append("raw_data = load '$INPUT_FILES' using ")
      .append(pigLoader.getName())
      .append("('")
      .append(thriftClass.getName())
      .append("');\n")
      .append(prefix)
      .append("as ");
    prefix.append("   ");

    try {
      stringifySchema(sb, toSchema(thriftClass), DataType.TUPLE, prefix);
    } catch (FrontendException e) {
      throw new RuntimeException(e);
    }

    sb.append("\n");
    return sb.toString();
  }

  /**
   * Print formatted schema. This is a modified version of
   * {@link Schema#stringifySchema(StringBuilder, Schema, byte)}
   * with support for (indented) pretty printing.
   */
  // This is used for building up output string
  // type can only be BAG or TUPLE
  public static void stringifySchema(StringBuilder sb,
                                     Schema schema,
                                     byte type,
                                     StringBuilder prefix)
                                          throws FrontendException{
      // this is a modified version of {@link Schema#stringifySchema(StringBuilder, Schema, byte)}
      if (type == DataType.TUPLE) {
          sb.append("(") ;
      }
      else if (type == DataType.BAG) {
          sb.append("{") ;
      }

      prefix.append("  ");
      sb.append("\n").append(prefix);

      if (schema == null) {
          sb.append("null") ;
      }
      else {
          boolean isFirst = true ;
          for (int i=0; i< schema.size() ;i++) {

              if (!isFirst) {
                  sb.append(",\n").append(prefix);
              }
              else {
                  isFirst = false ;
              }

              FieldSchema fs = schema.getField(i) ;

              if(fs == null) {
                  sb.append("null");
                  continue;
              }

              if (fs.alias != null) {
                  sb.append(fs.alias);
                  sb.append(": ");
              }

              if (DataType.isAtomic(fs.type)) {
                  sb.append(DataType.findTypeName(fs.type)) ;
              }
              else if ( (fs.type == DataType.TUPLE) ||
                        (fs.type == DataType.BAG) ) {
                  // safety net
                  if (schema != fs.schema) {
                      stringifySchema(sb, fs.schema, fs.type, prefix) ;
                  }
                  else {
                      throw new AssertionError("Schema refers to itself "
                                               + "as inner schema") ;
                  }
              } else if (fs.type == DataType.MAP) {
                sb.append(DataType.findTypeName(fs.type) + "[");
                if (fs.schema!=null)
                    stringifySchema(sb, fs.schema, fs.type, prefix);
                sb.append("]");
              } else {
                  sb.append(DataType.findTypeName(fs.type)) ;
              }
          }
      }

      prefix.setLength(prefix.length()-2);
      sb.append("\n").append(prefix);

      if (type == DataType.TUPLE) {
          sb.append(")") ;
      }
      else if (type == DataType.BAG) {
          sb.append("}") ;
      }
  }

  public static void main(String[] args) throws Exception {
    if (args.length > 0) {
      Class<? extends TBase<?, ?>> tClass = ThriftUtils.getTypeRef(args[0]).getRawClass();
      System.out.println(args[0] + " : " + toSchema(tClass).toString());
      System.out.println(toPigScript(tClass, ThriftPigLoader.class));
    }
  }
}


package com.twitter.elephantbird.mapreduce.input;

import java.io.IOException;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.ListIterator;

import com.google.common.collect.Lists;
import com.google.common.io.Closeables;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.slf4j.LoggerFactory;

import com.twitter.elephantbird.mapreduce.input.LuceneIndexInputFormat.LuceneIndexInputSplit;

/**
 * See {@link LuceneIndexInputFormat}
 *
 * @param <T> - the type that your lucene Documents will be converted to
 * @author Alex Levenson
 */
public abstract class LuceneIndexRecordReader<T extends Writable>
    extends RecordReader<IntWritable, T> {
  private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(LuceneIndexRecordReader.class);

  private IntWritable currentKey;
  private T currentValue;
  private Iterator<T> currentValueIter;

  private List<Query> queries;
  private ListIterator<Query> currentQueryIter;

  private ListIterator<Path> currentIndexPathIter;
  private IndexReader currentIndexReader;
  private int numIndexes;

  private IndexSearcher indexSearcher;
  private Configuration conf;

  protected TaskAttemptContext context;

  /**
   * Convert a query from its serialized (string) form to a {@link Query}
   * This can be done using a QueryParser if your serialized form is a lucene query string,
   * or with custom {@link Query} serialization / deserialization
   *
   * @param serializedQuery string representation of the query as retrieved from the job conf
   * @return a {@link Query} to use to search the index
   * @throws IOException
   */
  protected abstract Query deserializeQuery(String serializedQuery) throws IOException;

  /**
   * Open an index stored in path
   * Override if you want more control over how indexes are opened
   *
   * @param path path to the index
   * @param conf job conf
   * @return an IndexReader of the index in path
   * @throws IOException
   */
  protected IndexReader openIndex(Path path, Configuration conf) throws IOException {
    return DirectoryReader.open(new LuceneHdfsDirectory(path, path.getFileSystem(conf)));
  }

  /**
   * Given an {@link IndexReader} return an {@link IndexSearcher} for it
   * Override if you want more control over how the {@link IndexSearcher} is created
   * @param reader the IndexReader
   * @return an IndexSearcher
   */
  protected IndexSearcher createSearcher(IndexReader reader) {
    return new IndexSearcher(reader);
  }

  /**
   * Search the index and return an Iterator of values extracted from the index.
   * This should probably be a lazy iterator if there will be many returned values.
   *
   * @param searcher the index searcher to query
   * @param query the query to run
   * @return an iterator of values to be emitted as records (one by one) by this record reader
   * @throws IOException
   */
  protected abstract Iterator<T> search(final IndexSearcher searcher, Query query)
      throws IOException;

  @Override
  public void initialize(InputSplit genericSplit, TaskAttemptContext context) throws IOException,
      InterruptedException {

    conf = HadoopCompat.getConfiguration(context);
    this.context = context;

    // convert query strings in the job conf to lucene Query objects
    List<String> serializedQueries = LuceneIndexInputFormat.getQueries(conf);
    queries = Lists.newArrayListWithCapacity(serializedQueries.size());
    for (String serializedQuery : serializedQueries) {
      queries.add(deserializeQuery(serializedQuery));
    }

    // get the indexes this record reader needs to read
    List<Path> indexes = ((LuceneIndexInputSplit) genericSplit).getIndexDirs();
    numIndexes = indexes.size();
    // we will iterate over this just once
    currentIndexPathIter = indexes.listIterator();

    // set the value and query iterators to empty so that
    // nextKeyValue() will roll over to the next index
    currentValueIter = Collections.<T>emptyList().listIterator();
    currentQueryIter = Collections.<Query>emptyList().listIterator();
  }

  @Override
  public boolean nextKeyValue() throws IOException, InterruptedException {
    // first exhaust currentValues
    if (currentValueIter.hasNext()) {
      currentValue = currentValueIter.next();
      return true;
    }

    // no more currentValues, move to the next query
    // and try again
    if (currentQueryIter.hasNext()) {
      currentKey = new IntWritable(currentQueryIter.nextIndex());
      Query nextQuery = currentQueryIter.next();
      LOG.info("Running query " + nextQuery);
      Iterator<T> values = search(indexSearcher, nextQuery);
      currentValueIter = values;
      return nextKeyValue();
    }

    // no more currentValues nor queries in this index, move to the next index
    if (currentIndexPathIter.hasNext()) {
      closeIndexReader(currentIndexReader);
      Path next = currentIndexPathIter.next();
      LOG.info("Searching index: " + next);
      currentIndexReader = openIndex(next, conf);
      indexSearcher = createSearcher(currentIndexReader);
      currentQueryIter = queries.listIterator();
      return nextKeyValue();
    }

    // all done
    return false;
  }

  /**
   * Because {@link IndexReader#close()} is final, it is impossible to properly mock an IndexReader
   * in unit tests. Wrapping the call to close in this method allows us to override this method
   * during tests.
   *
   * NOTE: There's no good reason to override this method (unless you're creating mock IndexReaders)
   *
   * @param reader to close
   * @throws IOException
   */
  protected void closeIndexReader(IndexReader reader) throws IOException {
    Closeables.close(currentIndexReader, false);
  }

  @Override
  public IntWritable getCurrentKey() throws IOException, InterruptedException {
    return currentKey;
  }

  @Override
  public T getCurrentValue() throws IOException, InterruptedException {
    return currentValue;
  }

  @Override
  public void close() throws IOException {
    // nothing to do
  }

  /**
   * This is sort of an approximation of progress.
   * It splits the progress equally among all indexes, then among all queries for that index.
   * So it won't move linearly, since we don't know how many hits there will be per query
   */
  @Override
  public float getProgress() {
    if (numIndexes < 1) {
      return 1.0f;
    }

    float indexProgress = (float) currentIndexPathIter.previousIndex() / (float) numIndexes;

    float queriesProgress = 1.0f;
    if (queries.size() > 0) {
      queriesProgress = (float) currentQueryIter.previousIndex() / (float) queries.size();
    }

    queriesProgress *= 1.0f / numIndexes;

    return indexProgress + queriesProgress;
  }
}


package com.twitter.elephantbird.pig.util;

import java.util.Collections;
import java.util.List;

import com.google.common.base.Joiner;
import com.google.common.collect.Lists;
import com.google.protobuf.Descriptors.Descriptor;
import com.google.protobuf.Descriptors.EnumValueDescriptor;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.ByteString;
import com.google.protobuf.Message;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.DataByteArray;
import org.apache.pig.data.DataType;
import org.apache.pig.data.NonSpillableDataBag;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.logicalLayer.schema.Schema;
import org.apache.pig.impl.logicalLayer.schema.Schema.FieldSchema;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A class for turning codegen'd protos into Pig Tuples and Schemas
 * for custom Pig LoadFuncs.
 * @author Kevin Weil
 */
public class ProtobufToPig {
  private static final Logger LOG = LoggerFactory.getLogger(ProtobufToPig.class);

  private static final TupleFactory tupleFactory_ = TupleFactory.getInstance();

  public ProtobufToPig() { }

  /**
   * Turn a generic message into a Tuple.  Individual fields that are enums
   * are converted into their string equivalents.  Fields that are not filled
   * out in the protobuf are set to null, unless there is a default field value in
   * which case that is used instead.
   * @param msg the protobuf message
   * @return a pig tuple representing the message.
   */
  public Tuple toTuple(Message msg) {
    if (msg == null) {
      // Pig tuples deal gracefully with nulls.
      // Also, we can be called with null here in recursive calls.
      return null;
    }

    Descriptor msgDescriptor = msg.getDescriptorForType();
    Tuple tuple = tupleFactory_.newTuple(msgDescriptor.getFields().size());
    int curField = 0;
    try {
      // Walk through all the possible fields in the message.
      for (FieldDescriptor fieldDescriptor : msgDescriptor.getFields()) {
        // Get the set value, or the default value, or null.
        Object fieldValue = msg.getField(fieldDescriptor);

        if (fieldDescriptor.getType() == FieldDescriptor.Type.MESSAGE) {
          tuple.set(curField++, messageToTuple(fieldDescriptor, fieldValue));
        } else {
          tuple.set(curField++, singleFieldToTuple(fieldDescriptor, fieldValue));
        }
      }
    } catch (ExecException e) {
      LOG.warn("Could not convert msg " + msg + " to tuple", e);
    }

    return tuple;
  }

  /**
   * Returns either {@link #messageToTuple(FieldDescriptor, Object)}
   * or {@link #singleFieldToTuple(FieldDescriptor, Object)} depending
   * on whether the field is a Message or a simple field.
   */
  public Object fieldToPig(FieldDescriptor fieldDescriptor, Object fieldValue) {
    if (fieldValue == null) {
      // protobufs unofficially ensures values are not null. just in case:
      return null;
    }
    if (fieldDescriptor.getType() == FieldDescriptor.Type.MESSAGE) {
      return messageToTuple(fieldDescriptor, fieldValue);
    } else {
      return singleFieldToTuple(fieldDescriptor, fieldValue);
    }
  }

  /**
   * Translate a nested message to a tuple.  If the field is repeated, it walks the list and adds each to a bag.
   * Otherwise, it just adds the given one.
   * @param fieldDescriptor the descriptor object for the given field.
   * @param fieldValue the object representing the value of this field, possibly null.
   * @return the object representing fieldValue in Pig -- either a bag or a tuple.
   */
  @SuppressWarnings("unchecked")
  protected Object messageToTuple(FieldDescriptor fieldDescriptor, Object fieldValue) {
    if (fieldValue == null) {
      // protobufs unofficially ensures values are not null. just in case:
      return null;
    }
    assert fieldDescriptor.getType() == FieldDescriptor.Type.MESSAGE : "messageToTuple called with field of type " + fieldDescriptor.getType();

    if (fieldDescriptor.isRepeated()) {
      // The protobuf contract is that if the field is repeated, then the object returned is actually a List
      // of the underlying datatype, which in this case is a nested message.
      List<Message> messageList = (List<Message>) (fieldValue != null ? fieldValue : Lists.newArrayList());

      DataBag bag = new NonSpillableDataBag(messageList.size());
      for (Message m : messageList) {
        bag.add(new ProtobufTuple(m));
      }
      return bag;
    } else {
      return new ProtobufTuple((Message)fieldValue);
    }
  }

  /**
   * Translate a single field to a tuple.  If the field is repeated, it walks the list and adds each to a bag.
   * Otherwise, it just adds the given one.
   * @param fieldDescriptor the descriptor object for the given field.
   * @param fieldValue the object representing the value of this field, possibly null.
   * @return the object representing fieldValue in Pig -- either a bag or a single field.
   * @throws ExecException if Pig decides to.  Shouldn't happen because we won't walk off the end of a tuple's field set.
   */
  @SuppressWarnings("unchecked")
  protected Object singleFieldToTuple(FieldDescriptor fieldDescriptor, Object fieldValue) {
    assert fieldDescriptor.getType() != FieldDescriptor.Type.MESSAGE : "messageToFieldSchema called with field of type " + fieldDescriptor.getType();

    if (fieldDescriptor.isRepeated()) {
      // The protobuf contract is that if the field is repeated, then the object returned is actually a List
      // of the underlying datatype, which in this case is a "primitive" like int, float, String, etc.
      // We have to make a single-item tuple out of it to put it in the bag.
      List<Object> fieldValueList = (List<Object>) (fieldValue != null ? fieldValue : Collections.emptyList());
      DataBag bag = new NonSpillableDataBag(fieldValueList.size());
      for (Object singleFieldValue : fieldValueList) {
        Object nonEnumFieldValue = coerceToPigTypes(fieldDescriptor, singleFieldValue);
        Tuple innerTuple = tupleFactory_.newTuple(1);
        try {
          innerTuple.set(0, nonEnumFieldValue);
        } catch (ExecException e) { // not expected
          throw new RuntimeException(e);
        }
        bag.add(innerTuple);
      }
      return bag;
    } else {
      return coerceToPigTypes(fieldDescriptor, fieldValue);
    }
  }

  /**
   * If the given field value is an enum, translate it to the enum's name as a string, since Pig cannot handle enums.
   * Also, if the given field value is a bool, translate it to 0 or 1 to avoid Pig bools, which can be sketchy.
   * @param fieldDescriptor the descriptor object for the given field.
   * @param fieldValue the object representing the value of this field, possibly null.
   * @return the object, unless it was from an enum field, in which case we return the name of the enum field.
   */
  private Object coerceToPigTypes(FieldDescriptor fieldDescriptor, Object fieldValue) {
    if (fieldDescriptor.getType() == FieldDescriptor.Type.ENUM && fieldValue != null) {
      EnumValueDescriptor enumValueDescriptor = (EnumValueDescriptor)fieldValue;
      return enumValueDescriptor.getName();
    } else if (fieldDescriptor.getType() == FieldDescriptor.Type.BOOL && fieldValue != null) {
      Boolean boolValue = (Boolean)fieldValue;
      return new Integer(boolValue ? 1 : 0);
    } else if (fieldDescriptor.getType() == FieldDescriptor.Type.BYTES && fieldValue != null) {
      ByteString bsValue = (ByteString)fieldValue;
      return new DataByteArray(bsValue.toByteArray());
    }
    return fieldValue;
  }

  /**
   * Turn a generic message descriptor into a Schema.  Individual fields that are enums
   * are converted into their string equivalents.
   * @param msgDescriptor the descriptor for the given message type.
   * @return a pig schema representing the message.
   */
  public Schema toSchema(Descriptor msgDescriptor) {
    Schema schema = new Schema();

    try {
      // Walk through all the possible fields in the message.
      for (FieldDescriptor fieldDescriptor : msgDescriptor.getFields()) {
        if (fieldDescriptor.getType() == FieldDescriptor.Type.MESSAGE) {
          schema.add(messageToFieldSchema(fieldDescriptor));
        } else {
          schema.add(singleFieldToFieldSchema(fieldDescriptor));
        }
      }
    } catch (FrontendException e) {
      LOG.warn("Could not convert descriptor " + msgDescriptor + " to schema", e);
    }

    return schema;
  }

  /**
   * Turn a nested message into a Schema object.  For repeated nested messages, it generates a schema for a bag of
   * tuples.  For non-repeated nested messages, it just generates a schema for the tuple itself.
   * @param fieldDescriptor the descriptor object for the given field.
   * @return the Schema for the nested message.
   * @throws FrontendException if Pig decides to.
   */
  private FieldSchema messageToFieldSchema(FieldDescriptor fieldDescriptor) throws FrontendException {
    assert fieldDescriptor.getType() == FieldDescriptor.Type.MESSAGE : "messageToFieldSchema called with field of type " + fieldDescriptor.getType();

    Schema innerSchema = toSchema(fieldDescriptor.getMessageType());

    if (fieldDescriptor.isRepeated()) {
      Schema tupleSchema = new Schema();
      tupleSchema.add(new FieldSchema(fieldDescriptor.getName() + "_tuple", innerSchema, DataType.TUPLE));
      return new FieldSchema(fieldDescriptor.getName(), tupleSchema, DataType.BAG);
    } else {
      return new FieldSchema(fieldDescriptor.getName(), innerSchema, DataType.TUPLE);
    }
  }

  /**
   * Turn a single field into a Schema object.  For repeated single fields, it generates a schema for a bag of single-item tuples.
   * For non-repeated fields, it just generates a standard field schema.
   * @param fieldDescriptor the descriptor object for the given field.
   * @return the Schema for the nested message.
   * @throws FrontendException if Pig decides to.
   */
  private FieldSchema singleFieldToFieldSchema(FieldDescriptor fieldDescriptor) throws FrontendException {
    assert fieldDescriptor.getType() != FieldDescriptor.Type.MESSAGE : "singleFieldToFieldSchema called with field of type " + fieldDescriptor.getType();

    if (fieldDescriptor.isRepeated()) {
      Schema itemSchema = new Schema();
      itemSchema.add(new FieldSchema(fieldDescriptor.getName(), null, getPigDataType(fieldDescriptor)));
      Schema itemTupleSchema = new Schema();
      itemTupleSchema.add(new FieldSchema(fieldDescriptor.getName() + "_tuple", itemSchema, DataType.TUPLE));

      return new FieldSchema(fieldDescriptor.getName() + "_bag", itemTupleSchema, DataType.BAG);
    } else {
      return new FieldSchema(fieldDescriptor.getName(), null, getPigDataType(fieldDescriptor));
    }
  }

  /**
   * Translate between protobuf's datatype representation and Pig's datatype representation.
   * @param fieldDescriptor the descriptor object for the given field.
   * @return the byte representing the pig datatype for the given field type.
   */
  private byte getPigDataType(FieldDescriptor fieldDescriptor) {
    switch (fieldDescriptor.getType()) {
      case INT32:
      case UINT32:
      case SINT32:
      case FIXED32:
      case SFIXED32:
      case BOOL: // We convert booleans to ints for pig output.
        return DataType.INTEGER;
      case INT64:
      case UINT64:
      case SINT64:
      case FIXED64:
      case SFIXED64:
        return DataType.LONG;
      case FLOAT:
        return DataType.FLOAT;
      case DOUBLE:
        return DataType.DOUBLE;
      case STRING:
      case ENUM: // We convert enums to strings for pig output.
        return DataType.CHARARRAY;
      case BYTES:
        return DataType.BYTEARRAY;
      case MESSAGE:
        throw new IllegalArgumentException("getPigDataType called on field " + fieldDescriptor.getFullName() + " of type message.");
      default:
        throw new IllegalArgumentException("Unexpected field type. " + fieldDescriptor.toString() + " " + fieldDescriptor.getFullName() + " " + fieldDescriptor.getType());
    }
  }

  /**
   * Turn a generic message descriptor into a loading schema for a pig script.
   * @param msgDescriptor the descriptor for the given message type.
   * @param loaderClassName the fully qualified classname of the pig loader to use.  Not
   * passed a Class<? extends LoadFunc> because in many situations that class
   * is being generated as well, and so doesn't exist in compiled form.
   * @return a pig script that can load the given message.
   */
  public String toPigScript(Descriptor msgDescriptor, String loaderClassName) {
    StringBuffer sb = new StringBuffer();
    final int initialTabOffset = 3;

    sb.append("raw_data = load '$INPUT_FILES' using " + loaderClassName + "()").append("\n");
    sb.append(tabs(initialTabOffset)).append("as (").append("\n");
    sb.append(toPigScriptInternal(msgDescriptor, initialTabOffset));
    sb.append(tabs(initialTabOffset)).append(");").append("\n").append("\n");

    return sb.toString();
  }

  /**
   * Same as toPigScript(Descriptor, String) but allows parameters for the loader.
   *
   * @param msgDescriptor
   * @param loaderClassName
   * @param params
   * @return a pig script that can load the given message.
   */
  public String toPigScript(Descriptor msgDescriptor, String loaderClassName, String... params) {
    StringBuffer sb = new StringBuffer();
    final int initialTabOffset = 3;

    sb.append("raw_data = load '$INPUT_FILES' using ")
    .append(loaderClassName)
    .append("(");
    String paramString = "";
    if (params.length > 0) {
      paramString = "'" + Joiner.on(",'").join(params) + "'";
    }
    sb.append(paramString).append(")").append("\n");
    sb.append("/**\n");
    sb.append(tabs(initialTabOffset)).append("as (").append("\n");
    sb.append(toPigScriptInternal(msgDescriptor, initialTabOffset));
    sb.append(tabs(initialTabOffset)).append(")").append("\n").append("\n");
    sb.append("**/\n;\n");
    return sb.toString();
  }

  /**
   * Turn a generic message descriptor into a loading schema for a pig script.  Individual fields that are enums
   * are converted into their string equivalents.
   * @param msgDescriptor the descriptor for the given message type.
   * @param numTabs the tab depth at the current point in the recursion, for pretty printing.
   * @return a pig schema representing the message.
   */
  private StringBuffer toPigScriptInternal(Descriptor msgDescriptor, int numTabs) {
    StringBuffer sb = new StringBuffer();
    try {
      // Walk through all the possible fields in the message.
      for (FieldDescriptor fieldDescriptor : msgDescriptor.getFields()) {
        // We have to add a comma after every line EXCEPT for the last, or Pig gets mad.
        boolean isLast = (fieldDescriptor == msgDescriptor.getFields().get(msgDescriptor.getFields().size() - 1));
        if (fieldDescriptor.getType() == FieldDescriptor.Type.MESSAGE) {
          sb.append(messageToPigScript(fieldDescriptor, numTabs + 1, isLast));
        } else {
          sb.append(singleFieldToPigScript(fieldDescriptor, numTabs + 1, isLast));
        }
      }
    } catch (FrontendException e) {
      LOG.warn("Could not convert descriptor " + msgDescriptor + " to pig script", e);
    }

    return sb;
  }

  /**
   * Turn a nested message into a pig script load string.  For repeated nested messages, it generates a string for a bag of
   * tuples.  For non-repeated nested messages, it just generates a string for the tuple itself.
   * @param fieldDescriptor the descriptor object for the given field.
   * @param numTabs the tab depth at the current point in the recursion, for pretty printing.
   * @return the pig script load schema for the nested message.
   * @throws FrontendException if Pig decides to.
   */
  private StringBuffer messageToPigScript(FieldDescriptor fieldDescriptor, int numTabs, boolean isLast) throws FrontendException {
    assert fieldDescriptor.getType() == FieldDescriptor.Type.MESSAGE : "messageToPigScript called with field of type " + fieldDescriptor.getType();

    if (fieldDescriptor.isRepeated()) {
      return new StringBuffer().append(tabs(numTabs)).append(fieldDescriptor.getName()).append(": bag {").append("\n")
          .append(tabs(numTabs + 1)).append(fieldDescriptor.getName()).append("_tuple: tuple (").append("\n")
          .append(toPigScriptInternal(fieldDescriptor.getMessageType(), numTabs + 2))
          .append(tabs(numTabs + 1)).append(")").append("\n")
          .append(tabs(numTabs)).append("}").append(isLast ? "" : ",").append("\n");
    } else {
      return new StringBuffer().append(tabs(numTabs)).append(fieldDescriptor.getName()).append(": tuple (").append("\n")
          .append(toPigScriptInternal(fieldDescriptor.getMessageType(), numTabs + 1))
          .append(tabs(numTabs)).append(")").append(isLast ? "" : ",").append("\n");
    }
  }

  /**
   * Turn a single field into a pig script load string.  For repeated single fields, it generates a string for a bag of single-item tuples.
   * For non-repeated fields, it just generates a standard single-element string.
   * @param fieldDescriptor the descriptor object for the given field.
   * @param numTabs the tab depth at the current point in the recursion, for pretty printing.
   * @return the pig script load string for the nested message.
   * @throws FrontendException if Pig decides to.
   */
  private StringBuffer singleFieldToPigScript(FieldDescriptor fieldDescriptor, int numTabs, boolean isLast) throws FrontendException {
    assert fieldDescriptor.getType() != FieldDescriptor.Type.MESSAGE : "singleFieldToPigScript called with field of type " + fieldDescriptor.getType();

    if (fieldDescriptor.isRepeated()) {
      return new StringBuffer().append(tabs(numTabs)).append(fieldDescriptor.getName()).append("_bag: bag {").append("\n")
          .append(tabs(numTabs + 1)).append(fieldDescriptor.getName()).append("_tuple: tuple (").append("\n")
          .append(tabs(numTabs + 2)).append(fieldDescriptor.getName()).append(": ").append(getPigScriptDataType(fieldDescriptor)).append("\n")
          .append(tabs(numTabs + 1)).append(")").append("\n")
          .append(tabs(numTabs)).append("}").append(isLast ? "" : ",").append("\n");
    } else {
      return new StringBuffer().append(tabs(numTabs)).append(fieldDescriptor.getName()).append(": ")
          .append(getPigScriptDataType(fieldDescriptor)).append(isLast ? "" : ",").append("\n");
    }
  }

  /**
   * Translate between protobuf's datatype representation and Pig's datatype representation.
   * @param fieldDescriptor the descriptor object for the given field.
   * @return the byte representing the pig datatype for the given field type.
   */
  private String getPigScriptDataType(FieldDescriptor fieldDescriptor) {
    switch (fieldDescriptor.getType()) {
      case INT32:
      case UINT32:
      case SINT32:
      case FIXED32:
      case SFIXED32:
      case BOOL: // We convert booleans to ints for pig output.
        return "int";
      case INT64:
      case UINT64:
      case SINT64:
      case FIXED64:
      case SFIXED64:
        return "long";
      case FLOAT:
        return "float";
      case DOUBLE:
        return "double";
      case STRING:
      case ENUM: // We convert enums to strings for pig output.
        return "chararray";
      case BYTES:
        return "bytearray";
      case MESSAGE:
        throw new IllegalArgumentException("getPigScriptDataType called on field " + fieldDescriptor.getFullName() + " of type message.");
      default:
        throw new IllegalArgumentException("Unexpected field type. " + fieldDescriptor.toString() + " " + fieldDescriptor.getFullName() + " " + fieldDescriptor.getType());
    }
  }

  private StringBuffer tabs(int numTabs) {
    StringBuffer sb = new StringBuffer();
    for (int i = 0; i < numTabs; i++) {
      sb.append("  ");
    }
    return sb;
  }

}


package com.twitter.elephantbird.thrift;

import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.nio.ByteBuffer;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import org.apache.thrift.TBase;
import org.apache.thrift.TEnum;
import org.apache.thrift.TException;
import org.apache.thrift.TFieldIdEnum;
import org.apache.thrift.TUnion;
import org.apache.thrift.meta_data.FieldMetaData;

import org.apache.thrift.protocol.TType;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import com.twitter.elephantbird.util.ThriftUtils;

/**
 * Expanded metadata of a Thrift class. The main purpose is
 * help recursive traversals of Thrift structs or objects
 * with the following :
 * <ul>
 *  <li> build much more detailed information about fields so that
 *       other iterators don't need to.
 *  <li> avoids runtime type checking while processing many objects of the
 *       same class (common case).
 *  <li> Handles different Thrift quirks.
 *
 */
public class TStructDescriptor {

  private static final Map<Class<?>, TStructDescriptor> structMap = Maps.newHashMap();

  private List<Field> fields;
  private Class<? extends TBase<?, ?>> tClass;
  private boolean isUnion;

  public Class<? extends TBase<?, ?>> getThriftClass() {
    return tClass;
  }

  public TBase<?, ?> newThriftObject() throws TException {
    try {
      return tClass.newInstance();
    } catch (Exception e) { //not excpected
      throw new TException(e);
    }
  }
  /**
   * The list of fields returned is immutable.
   */
  public List<Field> getFields() {
    return fields;
  }

  public Field getFieldAt(int idx) {
    return fields.get(idx);
  }

  @SuppressWarnings("unchecked")
  public Object getFieldValue(int fieldIdx, TBase tObject) {
    /* Thrift 0.5 throws an NPE when the field is binary and
     * happens to be null. Otherwise user could just
     * invoke tObject.getFieldValue(field.getFieldIdEnum()).
     * Should revisit this once we move to a newer version.
     *
     * here, the assumption is the field is not null most of the
     * time. otherwise, rather than catching the exception, we
     * could cache 'BufferForFieldName()' method and invoke it.
     *
     * this also helps with unions.
     */
    Field field = fields.get(fieldIdx);
    try {
      if (isUnion && field.getFieldIdEnum() != ((TUnion<?, ?>)tObject).getSetField()) {
        return null;
      }
      return tObject.getFieldValue(field.getFieldIdEnum());
    } catch (NullPointerException e) {
      return null;
    }
  }

  /**
   * Creates a descriptor for a Thrift class
   */
  public static TStructDescriptor getInstance(Class<? extends TBase<?, ?>> tClass) {
    synchronized (structMap) {
      TStructDescriptor desc = structMap.get(tClass);
      if (desc == null) {
        desc = new TStructDescriptor();
        desc.tClass = tClass;
        structMap.put(tClass, desc);
        desc.build(tClass);
      }
      return desc;
    }
  }

  private TStructDescriptor() {
    // all the initialization is done in build().
  }

  private void build(Class<? extends TBase<?, ?>> tClass) {
    Map<? extends TFieldIdEnum, FieldMetaData> fieldMap = FieldMetaData.getStructMetaDataMap(tClass);
    Field[] arr = new Field[fieldMap.size()];

    isUnion = TUnion.class.isAssignableFrom(tClass);

    int idx = 0;
    for (Entry<? extends TFieldIdEnum, FieldMetaData> e : fieldMap.entrySet()) {
      String fieldName = e.getKey().getFieldName();
      arr[idx++] = new Field(e.getKey(),
                             e.getValue(),
                             fieldName,
                             ThriftUtils.getFieldType(tClass, fieldName));
    }
    // make it immutable since users have access.
    fields = ImmutableList.copyOf(arr);
  }

  /**
   * returns 'enum name -> enum object' mapping.
   * Currently used for converting Tuple to a Thrift object.
   */
  static private Map<String, TEnum> extractEnumMap(Class<? extends TEnum> enumClass) {
    ImmutableMap.Builder<String, TEnum> builder = ImmutableMap.builder();
    for(TEnum e : enumClass.getEnumConstants()) {
      builder.put(e.toString(), e);
    }
    return builder.build();
  }

  /**
   * Maintains all the relevant info for a field
   */
  public static class Field {

    private final TFieldIdEnum fieldIdEnum;
    private final FieldMetaData fieldMetaData;
    private final short fieldId;
    private final String fieldName;
    private final byte ttype;

    // following fields are set when they are relevant.
    // though these are not declared final, this class should be immutable.
    private Field listElemField                 = null;   // lists
    private Field setElemField                  = null;   // sets
    private Field mapKeyField                   = null;   // maps
    private Field mapValueField                 = null;   // maps
    private Class<? extends TEnum> enumClass    = null;   // enums
    private Map<String, TEnum> enumMap          = null;   // enums
    private Map<Integer, TEnum> enumIdMap       = null;   // enums
    private TStructDescriptor tStructDescriptor = null;   // Structs
    private boolean isBuffer                    = false;  // strings

    @SuppressWarnings("unchecked")
    private Field(TFieldIdEnum fieldIdEnum, FieldMetaData fieldMetaData, String fieldName,
                  Type genericType) {
      this.fieldIdEnum = fieldIdEnum;
      this.fieldMetaData = fieldMetaData;
      this.fieldId = fieldIdEnum == null ? 1 : fieldIdEnum.getThriftFieldId();
      this.fieldName = fieldName;
      this.ttype = getTTypeFromJavaType(genericType);

      Type firstTypeArg = null;
      Type secondTypeArg = null;

      if (genericType instanceof ParameterizedType) {
        Type[] typeArgs = ((ParameterizedType) genericType).getActualTypeArguments();
        firstTypeArg = typeArgs.length > 0 ? typeArgs[0] : null;
        secondTypeArg = typeArgs.length > 1 ? typeArgs[1] : null;
      }

      switch (ttype) {

        case TType.LIST:
          listElemField = new Field(null, null, fieldName + "_list_elem", firstTypeArg);
          break;

        case TType.MAP:
          mapKeyField = new Field(null, null, fieldName + "_map_key", firstTypeArg);
          mapValueField = new Field(null, null, fieldName + "_map_value", secondTypeArg);
          break;

        case TType.SET:
          setElemField = new Field(null, null, fieldName + "_set_elem", firstTypeArg);
          break;

        case TType.ENUM:
          enumClass = (Class<? extends TEnum>)genericType;
          enumMap = extractEnumMap(enumClass);

          ImmutableMap.Builder<Integer, TEnum> builder = ImmutableMap.builder();
          for(TEnum e : enumMap.values()) {
            builder.put(e.getValue(), e);
          }
          enumIdMap = builder.build();
          break;

        case TType.STRUCT:
          tStructDescriptor = getInstance((Class<? extends TBase<?, ?>>) genericType);
          break;

        case TType.STRING:
          isBuffer = !genericType.equals(String.class);
          break;

        default:
          // this is ok. TType.INT, TType.BYTE, etc.
      }
    }

    public short getFieldId() {
      return fieldId;
    }

    public byte getType() {
      return ttype;
    }

    /**
     * This is valid only for fields of a Struct. It is null for other fields
     * in containers like List.
     */
    public TFieldIdEnum getFieldIdEnum() {
      return fieldIdEnum;
    }

    public boolean isBuffer() {
      return isBuffer;
    }

    public boolean isList() {
      return listElemField != null;
    }

    public Field getListElemField() {
      return listElemField;
    }

    public boolean isSet() {
      return setElemField != null;
    }

    public Field getSetElemField() {
      return setElemField;
    }

    public boolean isMap() {
      return mapKeyField != null;
    }

    public Field getMapKeyField() {
      return mapKeyField;
    }

    public Field getMapValueField() {
      return mapValueField;
    }

    public boolean isStruct() {
      return tStructDescriptor != null;
    }

    public TStructDescriptor gettStructDescriptor() {
      return tStructDescriptor;
    }

    /**
     *
     * @return true if the field is a known enum.
     */
    public boolean isEnum() {
      return enumClass != null;
    }

    public Class<? extends TEnum> getEnumClass() {
      return enumClass;
    }

    /**
     * @param name string value of the enum
     * @return the TEnum value of the named enum, or null
     * if the field is not a known enum.
     */
    public TEnum getEnumValueOf(String name) {
      return enumMap == null ? null : enumMap.get(name);
    }

    /**
     * @param id int value of the enum
     * @return the TEnum value of the named enum, or null
     * if the field is not a known enum.
     */
    public TEnum getEnumValueOf(int id) {
      return enumIdMap == null ? null : enumIdMap.get(id);
    }

    /**
     * @return the possible values of the enum field, or null
     * if the field is not a known enum.
     */
    public Collection<TEnum> getEnumValues() {
      return enumMap == null ? null : enumMap.values();
    }

    public String getName() {
      return fieldName;
    }

    public short getId() {
     return fieldId;
    }

    public FieldMetaData getFieldMetaData() {
      return fieldMetaData;
    }
  }

  /** Derives Thrift TType from java {@link Type} of a thrift field */
  private static byte getTTypeFromJavaType(Type jType) {

    // check primitive types and final classes
    if (jType == Boolean.TYPE || jType == Boolean.class)   return TType.BOOL;
    if (jType == Byte.TYPE    || jType == Byte.class)      return TType.BYTE;
    if (jType == Short.TYPE   || jType == Short.class)     return TType.I16;
    if (jType == Integer.TYPE || jType == Integer.class)   return TType.I32;
    if (jType == Long.TYPE    || jType == Long.class)      return TType.I64;
    if (jType == Double.TYPE  || jType == Double.class)    return TType.DOUBLE;
    if (jType == String.class)                             return TType.STRING;
    if (jType == byte[].class)                             return TType.STRING; // buffer
    if (jType == Void.class)                               return TType.VOID;

    // non-generic simple classes
    if (jType instanceof Class) {
      Class<?> klass = (Class<?>) jType;

      if (TEnum.class.isAssignableFrom(klass))      return TType.ENUM;
      if (ByteBuffer.class.isAssignableFrom(klass)) return TType.STRING; // buffer (ByteBuffer)
      if (TBase.class.isAssignableFrom(klass))      return TType.STRUCT;
    }

    // generic types
    if (jType instanceof ParameterizedType) {
      Class<?> klass = (Class<?>) ((ParameterizedType)jType).getRawType();

      if (Map.class.isAssignableFrom(klass))    return TType.MAP;
      if (Set.class.isAssignableFrom(klass))    return TType.SET;
      if (List.class.isAssignableFrom(klass))   return TType.LIST;
    }

    throw new IllegalArgumentException("cannot convert java type '" + jType + "'  to thrift type");
  }
}


package com.twitter.elephantbird.mapreduce.output;

import java.io.File;
import java.io.IOException;
import java.io.Reader;

import com.google.common.io.Files;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.FileUtil;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.fs.PathFilter;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.FileOutputCommitter;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.LogByteSizeMergePolicy;
import org.apache.lucene.index.SerialMergeScheduler;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.NoLockFactory;
import org.apache.lucene.store.SimpleFSDirectory;
import org.apache.lucene.util.Version;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.twitter.elephantbird.util.PathFilters;
import com.twitter.elephantbird.util.TaskHeartbeatThread;

/**
 * Base class for output formats that write lucene indexes
 *

 * Subclasses must specify how to convert a key value pair into a {@link Document}
 *

 * Subclasses may provide an {@link Analyzer} to use during index creation
 * (which may be used depending on how documents are created by the subclass)
 *
 * @author Alex Levenson, based on code written by Kyle Maxwell
 */
public abstract class LuceneIndexOutputFormat<K, V> extends FileOutputFormat<K, V> {
  private static final Logger LOG = LoggerFactory.getLogger(LuceneIndexOutputFormat.class);

  /**
   * Convert a record from the MR framework into a lucene {@link Document}
   * You may re-use the same {@link Document} instance for efficiency
   *
   * @param key the key written to this output format
   * @param value the value written to this output format
   * @return a lucene Document suitable for insertion into an {@link IndexWriter}
   */
  protected abstract Document buildDocument(K key, V value) throws IOException;

  /**
   * Override this method if you intend to use an {@link Analyzer}
   * during index creation. If you do not override this method, {@link NeverTokenizeAnalyzer} will
   * be used which will throw an exception if you create a {@link Document} using a method that
   * invokes tokenization.
   *
   * @param conf the job's configuration
   * @return an {@link Analyzer} suitable for use by an {@link IndexWriter}
   */
  protected Analyzer newAnalyzer(Configuration conf) {
    return new NeverTokenizeAnalyzer();
  }

  /**
   * Override to use a different {@link Directory} implementation
   *
   * You may want to use {@link org.apache.lucene.store.FSDirectory#open}
   * which is supposed to select an appropriate
   * local FS implementation based on the current OS. However, we have seen cases
   * where using this leads to an implementation that hits {@link java.lang.OutOfMemoryError}
   * when building large indexes.
   */
  protected Directory getDirectoryImplementation(File location) throws IOException {
    return new SimpleFSDirectory(location, NoLockFactory.getNoLockFactory());
  }

  public static IndexWriter createIndexWriter(Directory location, Analyzer analyzer) throws IOException {
    return createIndexWriter(location, analyzer, LogByteSizeMergePolicy.DEFAULT_MERGE_FACTOR);
  }

  public static IndexWriter createIndexWriter(Directory location, Analyzer analyzer, int mergeFactor)
      throws IOException {

    LOG.info("Creating IndexWriter with:\nDirectory: "
        + location
        + "\nAnalyzer: "
        + analyzer
        + "\nMerge Factor: " + mergeFactor);
    IndexWriterConfig idxConfig = new IndexWriterConfig(Version.LUCENE_40, analyzer);
    LogByteSizeMergePolicy mergePolicy = new LogByteSizeMergePolicy();
    mergePolicy.setMergeFactor(mergeFactor);
    mergePolicy.setUseCompoundFile(false);

    idxConfig.setMergePolicy(mergePolicy);

    idxConfig.setMergeScheduler(new SerialMergeScheduler());

    IndexWriter writer = new IndexWriter(location, idxConfig);
    return writer;
  }

  @Override
  public RecordWriter<K, V> getRecordWriter(TaskAttemptContext job) throws IOException {
    FileOutputCommitter committer = (FileOutputCommitter) this.getOutputCommitter(job);
    File tmpDirFile = Files.createTempDir();
    Directory directory = getDirectoryImplementation(tmpDirFile);
    IndexWriter writer = createIndexWriter(directory, newAnalyzer(HadoopCompat.getConfiguration(job)));
    return new IndexRecordWriter(writer, committer, tmpDirFile);
  }

  private class IndexRecordWriter extends RecordWriter<K, V> {
    private IndexWriter writer;
    private FileOutputCommitter committer;
    private File tmpDirFile;
    private long recordsProcessed = 0;

    private IndexRecordWriter(IndexWriter writer, FileOutputCommitter committer, File tmpDirFile) {
      this.writer = writer;
      this.committer = committer;
      this.tmpDirFile = tmpDirFile;
    }

    @Override
    public void write(K key, V value) throws IOException {
      recordsProcessed++;
      if (recordsProcessed % 1000000 == 0) {
        LOG.info("Processing record " + recordsProcessed);
      }

      writer.addDocument(buildDocument(key, value));
    }

    @Override
    public void close(final TaskAttemptContext context) throws IOException, InterruptedException {
      TaskHeartbeatThread heartBeat = new TaskHeartbeatThread(context) {
        @Override
        public void progress() {
          String[] filesLeft = tmpDirFile.list();
          if (filesLeft != null) {
            int remaining = filesLeft.length - 2;
            LOG.info("Optimizing " + remaining + " segments");
          } else {
            LOG.info("Done optimizing segments, heartbeat thread still alive");
          }
        }
      };

      try {
        LOG.info("Starting heartbeat thread");
        heartBeat.start();

        Path work = committer.getWorkPath();
        Path output = new Path(work, "index-"
            + String.valueOf(HadoopCompat.getTaskAttemptID(context).getTaskID().getId()));

        writer.forceMerge(1);
        writer.close();

        FileSystem fs = FileSystem.get(HadoopCompat.getConfiguration(context));
        LOG.info("Copying index to HDFS...");

        if (!FileUtil.copy(tmpDirFile, fs, output, true, HadoopCompat.getConfiguration(context))) {
          throw new IOException("Failed to copy local index to HDFS!");
        }

        LOG.info("Index written to: " + output);
      } catch (IOException e) {
        LOG.error("Error committing index", e);
        throw e;
      } finally {
        // all things must die, eventually
        LOG.info("Stopping heartbeat thread");
        heartBeat.stop();
      }
    }
  }

  /**
   * An analyzer that always throws {@link UnsupportedOperationException}
   * when {@link #createComponents(String, java.io.Reader)} is called
   *

   * Useful if you don't intend to use an {@link Analyzer} for tokenization
   * but are required to provide one to an {@link org.apache.lucene.index.IndexWriter}
   */
  public static class NeverTokenizeAnalyzer extends Analyzer {
    @Override
    protected TokenStreamComponents createComponents(String fieldName, Reader reader) {
      throw new UnsupportedOperationException();
    }
  }

  /**
   * Creates a path filter that accepts non-hidden directories that start with "index-"
   * This is what the indexes created by this output format look like,
   * so this is useful for finding them when traversing the file system
   */
  public static PathFilter newIndexDirFilter(Configuration conf) {
    return new PathFilters.CompositePathFilter(
      PathFilters.newExcludeFilesFilter(conf),
      PathFilters.EXCLUDE_HIDDEN_PATHS_FILTER,
      new PathFilter() {
        @Override
        public boolean accept(Path path) {
          return path.getName().startsWith("index-");
        }
      }
    );
  }
}


package com.twitter.elephantbird.util;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.List;
import java.util.Map;

import com.google.common.base.Function;
import com.google.common.base.Predicates;
import com.google.common.collect.Lists;
import com.google.protobuf.DescriptorProtos.DescriptorProto;
import com.google.protobuf.Descriptors.Descriptor;
import com.google.protobuf.Descriptors.DescriptorValidationException;
import com.google.protobuf.Descriptors.EnumValueDescriptor;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.Descriptors.FieldDescriptor.Type;
import com.google.protobuf.Message.Builder;
import com.google.protobuf.ByteString;
import com.google.protobuf.CodedInputStream;
import com.google.protobuf.CodedOutputStream;
import com.google.protobuf.DescriptorProtos;
import com.google.protobuf.Descriptors;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
import com.google.protobuf.ProtocolMessageEnum;
import com.google.protobuf.UninitializedMessageException;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Text;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class Protobufs {
  private static final Logger LOG = LoggerFactory.getLogger(Protobufs.class);

  public static final byte[] KNOWN_GOOD_POSITION_MARKER = new byte[] { 0x29, (byte)0xd8, (byte)0xd5, 0x06, 0x58,
                                                                         (byte)0xcd, 0x4c, 0x29, (byte)0xb2,
                                                                         (byte)0xbc, 0x57, (byte)0x99, 0x21, 0x71,
                                                                         (byte)0xbd, (byte)0xff };

  public static final byte NEWLINE_UTF8_BYTE = '\n';
  public static final byte[] NEWLINE_UTF8_BYTES = new byte[]{NEWLINE_UTF8_BYTE};

  public static final String IGNORE_KEY = "IGNORE";

  private static final String CLASS_CONF_PREFIX = "elephantbird.protobuf.class.for.";

  /**
   * Returns Protobuf class. The class name could be either normal name or
   * its canonical name (canonical name does not have a $ for inner classes).
   */
  public static Class<? extends Message> getProtobufClass(String protoClassName) {
    return getProtobufClass(null, protoClassName);
  }

  private static Class<? extends Message> getProtobufClass(Configuration conf, String protoClassName) {
    // Try both normal name and canonical name of the class.
    Class<?> protoClass = null;
    try {
      if (conf == null) {
        protoClass = Class.forName(protoClassName);
      } else {
        protoClass = conf.getClassByName(protoClassName);
      }
    } catch (ClassNotFoundException e) {
      // the class name might be canonical name.
      protoClass = getInnerProtobufClass(protoClassName);
    }

    return protoClass.asSubclass(Message.class);
  }

  /**
   * For a configured protoClass, should the message be dynamic or is it a pre-generated Message class? If protoClass is
   * null or set to DynamicMessage.class, then the configurer intends for a dynamically generated protobuf to be used.
   */
  public static boolean useDynamicProtoMessage(Class<?> protoClass) {
    return protoClass == null || protoClass.getCanonicalName().equals(DynamicMessage.class.getCanonicalName());
  }

  public static Class<? extends Message> getInnerProtobufClass(String canonicalClassName) {
    // is an inner class and is not visible from the outside.  We have to instantiate
    String parentClass = canonicalClassName.substring(0, canonicalClassName.lastIndexOf("."));
    String subclass = canonicalClassName.substring(canonicalClassName.lastIndexOf(".") + 1);
    return getInnerClass(parentClass, subclass);
  }

  public static Class<? extends Message> getInnerClass(String canonicalParentName, String subclassName) {
    try {
      Class<?> outerClass = Class.forName(canonicalParentName);
      for (Class<?> innerClass: outerClass.getDeclaredClasses()) {
        if (innerClass.getSimpleName().equals(subclassName)) {
          return innerClass.asSubclass(Message.class);
        }
      }
    } catch (ClassNotFoundException e) {
      LOG.error("Could not find class with parent " + canonicalParentName + " and inner class " + subclassName, e);
      throw new IllegalArgumentException(e);
    }
    return null;
  }

  public static Message.Builder getMessageBuilder(Class<? extends Message> protoClass) {
    try {
      Method newBuilder = protoClass.getMethod("newBuilder", new Class[] {});
      return (Message.Builder) newBuilder.invoke(null, new Object[] {});
    } catch (NoSuchMethodException e) {
      LOG.error("Could not find method newBuilder in class " + protoClass, e);
      throw new IllegalArgumentException(e);
    } catch (IllegalAccessException e) {
      LOG.error("Could not access method newBuilder in class " + protoClass, e);
      throw new IllegalArgumentException(e);
    } catch (InvocationTargetException e) {
      LOG.error("Error invoking method newBuilder in class " + protoClass, e);
      throw new IllegalArgumentException(e);
    }
  }

  public static Descriptor getMessageDescriptor(Class<? extends Message> protoClass) {
    try {
      Method getDescriptor = protoClass.getMethod("getDescriptor", new Class[] {});
      return (Descriptor)getDescriptor.invoke(null, new Object[] {});
    } catch (NoSuchMethodException e) {
      LOG.error("Could not find method getDescriptor in class " + protoClass, e);
      throw new IllegalArgumentException(e);
    } catch (IllegalAccessException e) {
      LOG.error("Could not access method getDescriptor in class " + protoClass, e);
      throw new IllegalArgumentException(e);
    } catch (InvocationTargetException e) {
      LOG.error("Error invoking method getDescriptor in class " + protoClass, e);
    }

    return null;
  }

  public static List<String> getMessageFieldNames(Class<? extends Message> protoClass) {
    return Lists.transform(getMessageDescriptor(protoClass).getFields(), new Function<FieldDescriptor, String>() {
      public String apply(FieldDescriptor f) {
        return f.getName();
      }
    });
  }

  /**
   * Returns a Message {@link Descriptor} for a dynamically generated
   * DescriptorProto.
   *
   * @param descProto
   * @throws DescriptorValidationException
   */
  public static Descriptor makeMessageDescriptor(DescriptorProto descProto)
                                      throws DescriptorValidationException {

    DescriptorProtos.FileDescriptorProto fileDescP =
      DescriptorProtos.FileDescriptorProto.newBuilder().addMessageType(descProto).build();

    Descriptors.FileDescriptor[] fileDescs = new Descriptors.FileDescriptor[0];
    Descriptors.FileDescriptor dynamicDescriptor = Descriptors.FileDescriptor.buildFrom(fileDescP, fileDescs);
    return dynamicDescriptor.findMessageTypeByName(descProto.getName());
  }

   // Translates from protobuf field names to other field names, stripping out any that have value IGNORE.
   // i.e. if you have a protobuf with field names a, b, c and your fieldNameTranslations map looks like
   // { "a" => "A", "c" => "IGNORE" }, then you'll end up with A, b
  public static List<String> getMessageFieldNames(Class<? extends Message> protoClass, Map<String, String> fieldNameTranslations) {
    return getMessageFieldNames(getMessageDescriptor(protoClass), fieldNameTranslations);
  }

   public static List<String> getMessageFieldNames(Descriptor descriptor, Map<String, String> fieldNameTranslations) {
     Function<FieldDescriptor, String> fieldTransformer = getFieldTransformerFor(fieldNameTranslations);
     return ListHelper.filter(Lists.transform(descriptor.getFields(), fieldTransformer), Predicates.<String>notNull());
   }

   public static Function<FieldDescriptor, String> getFieldTransformerFor(final Map<String, String> fieldNameTranslations) {
     return new Function<FieldDescriptor, String>() {
       public String apply(FieldDescriptor f) {
         String name = f.getName();
         if (fieldNameTranslations != null && fieldNameTranslations.containsKey(name)) {
           name = fieldNameTranslations.get(name);
         }
         if (IGNORE_KEY.equals(name)) {
           name = null;
         }
         return name;
       }
     };
   }

  @SuppressWarnings("unchecked")
  public static <M extends Message> M parseFrom(Class<M> protoClass, byte[] messageBytes) {
    try {
      Method parseFrom = protoClass.getMethod("parseFrom", new Class[] { byte[].class });
      return (M)parseFrom.invoke(null, new Object[] { messageBytes });
    } catch (NoSuchMethodException e) {
      LOG.error("Could not find method parseFrom in class " + protoClass, e);
      throw new IllegalArgumentException(e);
    } catch (IllegalAccessException e) {
      LOG.error("Could not access method parseFrom in class " + protoClass, e);
      throw new IllegalArgumentException(e);
    } catch (InvocationTargetException e) {
      LOG.error("Error invoking method parseFrom in class " + protoClass, e);
    }

    return null;
  }

  public static DynamicMessage parseDynamicFrom(Class<? extends Message> protoClass, byte[] messageBytes) {
    try {
      return DynamicMessage.parseFrom(getMessageDescriptor(protoClass), messageBytes);
    } catch (InvalidProtocolBufferException e) {
      LOG.error("Protocol buffer parsing error in parseDynamicFrom", e);
    } catch(UninitializedMessageException ume) {
      LOG.error("Uninitialized Message error in parseDynamicFrom " + protoClass.getName(), ume);
    }

    return null;
  }

  public static Message instantiateFromClassName(String canonicalClassName) {
    Class<? extends Message> protoClass = getInnerProtobufClass(canonicalClassName);
    Message.Builder builder = getMessageBuilder(protoClass);
    return builder.build();
  }

  public static <M extends Message> M instantiateFromClass(Class<M> protoClass) {
    return Protobufs.parseFrom(protoClass, new byte[] {});
  }

  public static Message addField(Message m, String name, Object value) {
    Message.Builder builder = m.toBuilder();
    setFieldByName(builder, name, value);
    return builder.build();
  }

  public static void setFieldByName(Message.Builder builder, String name, Object value) {
    FieldDescriptor fieldDescriptor = builder.getDescriptorForType().findFieldByName(name);
    if (value == null) {
      builder.clearField(fieldDescriptor);
    } else {
      builder.setField(fieldDescriptor, value);
    }
  }

  public static boolean isFieldSetByName(Message message, String name) {
    return message.hasField(message.getDescriptorForType().findFieldByName(name));
  }

  public static Object getFieldByName(Message message, String name) {
    return message.getField(message.getDescriptorForType().findFieldByName(name));
  }

  public static boolean hasFieldByName(Message message, String name) {
    return message.getDescriptorForType().findFieldByName(name) != null;
  }

  public static Type getTypeByName(Message message, String name) {
    return message.getDescriptorForType().findFieldByName(name).getType();
  }

  /**
   * Returns typeref for a Protobuf class
   */
  public static<M extends Message> TypeRef<M> getTypeRef(String protoClassName) {
    return new TypeRef<M>(getProtobufClass(protoClassName)){};
  }

  /**
   * Returns TypeRef for the Protobuf class that was set using setClass(jobConf);
   */
  public static<M extends Message> TypeRef<M> getTypeRef(Configuration jobConf, Class<?> genericClass) {
    String className = jobConf.get(CLASS_CONF_PREFIX + genericClass.getName());
    if (className == null) {
      throw new RuntimeException(CLASS_CONF_PREFIX + genericClass.getName() + " is not set");
    }
    return new TypeRef<M>(getProtobufClass(jobConf, className)){};
  }

  public static void setClassConf(Configuration jobConf, Class<?> genericClass,
      Class<? extends Message> protoClass) {
    HadoopUtils.setClassConf(jobConf,
           CLASS_CONF_PREFIX + genericClass.getName(),
           protoClass);
  }

  public static Text toText(Message message) {
    return new Text(message.toByteArray());
  }

  public static  B mergeFromText(B builder, Text bytes)
                                    throws InvalidProtocolBufferException {
    @SuppressWarnings("unchecked")
    B b = (B) builder.mergeFrom(bytes.getBytes());
    return b;
  }

  /**
   * Serializes a single field. All the native fields are serialized using
   * "NoTag" methods in {@link CodedOutputStream}
   * e.g. writeInt32NoTag(). The field index is not written.
   *
   * @param output
   * @param fd
   * @param value
   * @throws IOException
   */
  public static void writeFieldNoTag(CodedOutputStream   output,
                                     FieldDescriptor     fd,
                                     Object              value)
                                     throws IOException {
    if (value == null) {
      return;
    }

    if (fd.isRepeated()) {
      @SuppressWarnings("unchecked")
      List<Object> values = (List<Object>) value;
      for(Object obj : values) {
        writeSingleFieldNoTag(output, fd, obj);
      }
    } else {
      writeSingleFieldNoTag(output, fd, value);
    }
  }

  private static void writeSingleFieldNoTag(CodedOutputStream   output,
                                            FieldDescriptor     fd,
                                            Object              value)
                                            throws IOException {
    switch (fd.getType()) {
    case DOUBLE:
      output.writeDoubleNoTag((Double) value);    break;
    case FLOAT:
      output.writeFloatNoTag((Float) value);      break;
    case INT64:
    case UINT64:
      output.writeInt64NoTag((Long) value);       break;
    case INT32:
      output.writeInt32NoTag((Integer) value);    break;
    case FIXED64:
      output.writeFixed64NoTag((Long) value);     break;
    case FIXED32:
      output.writeFixed32NoTag((Integer) value);  break;
    case BOOL:
      output.writeBoolNoTag((Boolean) value);     break;
    case STRING:
      output.writeStringNoTag((String) value);    break;
    case GROUP:
    case MESSAGE:
      output.writeMessageNoTag((Message) value);  break;
    case BYTES:
      output.writeBytesNoTag((ByteString) value); break;
    case UINT32:
      output.writeUInt32NoTag((Integer) value);   break;
    case ENUM:
      output.writeEnumNoTag(((ProtocolMessageEnum) value).getNumber()); break;
    case SFIXED32:
      output.writeSFixed32NoTag((Integer) value); break;
    case SFIXED64:
      output.writeSFixed64NoTag((Long) value);    break;
    case SINT32:
      output.writeSInt32NoTag((Integer) value);   break;
    case SINT64:
      output.writeSInt64NoTag((Integer) value);   break;

    default:
      throw new IllegalArgumentException("Unknown type " + fd.getType()
                                         + " for " + fd.getFullName());
    }
  }

  /**
   *  Wrapper around {@link #readFieldNoTag(CodedInputStream, FieldDescriptor, Builder)}.

   *  same as

   *  builder.setField(fd, readFieldNoTag(input, fd, builder));
   */
  public static void setFieldValue(CodedInputStream input,
                              FieldDescriptor  fd,
                              Builder          builder)
                              throws IOException {
    builder.setField(fd, readFieldNoTag(input, fd, builder));
  }

  /**
   * Deserializer for protobuf objects written with
   * {@link #writeFieldNoTag(CodedOutputStream, FieldDescriptor, Object)}.
   *
   * @param input
   * @param fd
   * @param enclosingBuilder required to create a builder when the field
   *                         is a message
   * @throws IOException
   */
  public static Object readFieldNoTag(CodedInputStream input,
                                      FieldDescriptor  fd,
                                      Builder          enclosingBuilder)
                                      throws IOException {
    if (!fd.isRepeated()) {
      return readSingleFieldNoTag(input, fd, enclosingBuilder);
    }

    // repeated field

    List<Object> values = Lists.newArrayList();
    while (!input.isAtEnd()) {
      values.add(readSingleFieldNoTag(input, fd, enclosingBuilder));
    }
    return values;
  }

  private static Object readSingleFieldNoTag(CodedInputStream input,
                                             FieldDescriptor  fd,
                                             Builder          enclosingBuilder)
                                             throws IOException {
    switch (fd.getType()) {
    case DOUBLE:
      return input.readDouble();
    case FLOAT:
      return input.readFloat();
    case INT64:
    case UINT64:
      return input.readInt64();
    case INT32:
      return input.readInt32();
    case FIXED64:
      return input.readFixed64();
    case FIXED32:
      return input.readFixed32();
    case BOOL:
      return input.readBool();
    case STRING:
      return input.readString();
    case GROUP:
    case MESSAGE:
      Builder fieldBuilder = enclosingBuilder.newBuilderForField(fd);
      input.readMessage(fieldBuilder, null);
      return fieldBuilder.build();
    case BYTES:
      return input.readBytes();
    case UINT32:
      return input.readUInt32();
    case ENUM:
      EnumValueDescriptor eVal = fd.getEnumType().findValueByNumber(input.readEnum());
      // ideally if a given enum does not exist, we should search
      // unknown fields. but we don't have access to that here. return default
      return eVal != null ? eVal : fd.getDefaultValue();
    case SFIXED32:
      return input.readSFixed32();
    case SFIXED64:
      return input.readSFixed64();
    case SINT32:
      return input.readSInt32();
    case SINT64:
      return input.readSInt64();

    default:
      throw new IllegalArgumentException("Unknown type " + fd.getType()
          + " for " + fd.getFullName());
    }
  }

}


package com.twitter.elephantbird.mapreduce.input;

import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.CompressionCodecFactory;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.input.FileSplit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * An abstract base class that handles setting up all LZO-based RecordReaders.

 *
 * Error handling:

 * A small fraction of bad records are tolerated. When deserialization
 * of a record results in an exception or a null object, an a warning
 * is logged. If the rate of errors crosses a threshold
 * (default is 0.0001 or 0.01%) a RuntimeException is thrown.

 * By default, the threshold is checked at the end of reading.
 * The threshold can be set with configuration variable
 * elephantbird.mapred.input.bad.record.threshold.
 * A value of 0 disables error handling.

 */
public abstract class LzoRecordReader<K, V> extends RecordReader<K, V> {
  private static final Logger LOG = LoggerFactory.getLogger(LzoRecordReader.class);

  public static final String BAD_RECORD_THRESHOLD_CONF_KEY = "elephantbird.mapred.input.bad.record.threshold";
  /* Error out only after threshold rate is reached and we have see this many errors */
  public static final String BAD_RECORD_MIN_COUNT_CONF_KEY = "elephantbird.mapred.input.bad.record.min";
  public static final String BAD_RECORD_CHECK_AT_CLOSE = // long name is ok, not usually explicitly set.
      "elephantbird.mapred.input.bad.record.check.only.in.close";

  protected long start_;
  protected long pos_;
  protected long end_;
  private FSDataInputStream fileIn_;

  protected InputErrorTracker errorTracker;

  /**
   * Get the progress within the split.
   */
  @Override
  public float getProgress() {
    if (start_ == end_) {
      return 0.0f;
    } else {
      return Math.min(1.0f, (pos_ - start_) / (float) (end_ - start_));
    }
  }

  public synchronized long getPos() throws IOException {
    return pos_;
  }

  public long getLzoFilePos() throws IOException {
    return fileIn_.getPos();
  }

  @Override
  public void initialize(InputSplit genericSplit, TaskAttemptContext context) throws IOException, InterruptedException {
    FileSplit split = (FileSplit) genericSplit;
    start_ = split.getStart();
    end_ = start_ + split.getLength();
    final Path file = split.getPath();
    Configuration job = HadoopCompat.getConfiguration(context);

    errorTracker = new InputErrorTracker(job);

    LOG.info("input split: " + file + " " + start_ + ":" + end_);

    FileSystem fs = file.getFileSystem(job);
    CompressionCodecFactory compressionCodecs = new CompressionCodecFactory(job);
    final CompressionCodec codec = compressionCodecs.getCodec(file);
    if (codec == null) {
      throw new IOException("No codec for file " + file + " found, cannot run");
    }

    // Open the file and seek to the start of the split.
    fileIn_ = fs.open(split.getPath());

    // Creates input stream and also reads the file header.
    createInputReader(codec.createInputStream(fileIn_), job);

    if (start_ != 0) {
      fileIn_.seek(start_);
      skipToNextSyncPoint(false);
      start_ = fileIn_.getPos();
      LOG.info("Start is now " + start_);
    } else {
      skipToNextSyncPoint(true);
    }
    pos_ = start_;
  }

  @Override
  public void close() throws IOException {
    if (errorTracker != null) {
      errorTracker.close();
    }
  }

  protected abstract void createInputReader(InputStream input, Configuration conf) throws IOException;
  protected abstract void skipToNextSyncPoint(boolean atFirstRecord) throws IOException;

  /**
   * Tracks number of of errors in input and throws a Runtime exception
   * if the rate of errors crosses a limit.

   *
   * The intention is to skip over very rare file corruption or incorrect
   * input, but catch programmer errors (incorrect format, or incorrect
   * deserializers etc).
   */
  static class InputErrorTracker implements Closeable {
    long numRecords;
    long numErrors;

    final double errorThreshold; // max fraction of errors allowed
    final long minErrors; // throw error only after this many errors

    final boolean checkOnlyInClose;

    InputErrorTracker(Configuration conf) {
      //default threshold : 0.01%
      errorThreshold = conf.getFloat(BAD_RECORD_THRESHOLD_CONF_KEY, 0.0001f);
      minErrors = conf.getLong(BAD_RECORD_MIN_COUNT_CONF_KEY, 2);
      checkOnlyInClose = conf.getBoolean(BAD_RECORD_CHECK_AT_CLOSE, true);
      numRecords = 0;
      numErrors = 0;
    }

    void incRecords() {
      numRecords++;
    }

    @Override
    public void close() throws IOException {
      // should throw RTE or an IOE?
      // some applications might ignore IOE from close().
      try {
        if (numErrors > 0) {
          checkErrorThreshold(null);
        }
      } catch (RuntimeException e) {
        throw new IOException(e);
      }
    }

    void incErrors(Throwable cause) {
      numErrors++;
      if (numErrors > numRecords) {
        // incorrect use of this class
        throw new RuntimeException("Forgot to invoke incRecords()?");
      }

      LOG.warn("Error while reading an input record ("
          + numErrors + " out of " + numRecords + " so far ): ", cause);

      if (!checkOnlyInClose) {
        checkErrorThreshold(cause);
      }
    }

    private void checkErrorThreshold(Throwable cause) {
      if (numErrors > 0 && errorThreshold <= 0) { // no errors are tolerated
        throw new RuntimeException("error while reading input records", cause);
      }

      double errRate = numErrors/(double)numRecords;

      // will always excuse the first error. We can decide if single
      // error crosses threshold inside close() if we want to.
      if (numErrors >= minErrors  && errRate > errorThreshold) {
        LOG.error(numErrors + " out of " + numRecords
            + " crosses configured threshold (" + errorThreshold + ")");
        throw new RuntimeException("error rate while reading input records crossed threshold", cause);
      }
    }
  }
}


package com.twitter.elephantbird.mapreduce.input;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.fs.PathFilter;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.mapreduce.lib.input.FileSplit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Lists;
import com.hadoop.compression.lzo.LzoIndex;

/**
 * An {@link org.apache.hadoop.mapreduce.InputFormat} for lzop compressed files. This class
 * handles the nudging the input splits onto LZO boundaries using the existing LZO index files.
 * Subclass and implement getRecordReader to define custom LZO-based input formats.

 * Note: unlike the stock FileInputFormat, this recursively examines directories for matching files.
 */
public abstract class LzoInputFormat<K, V> extends FileInputFormat<K, V> {
  private static final Logger LOG = LoggerFactory.getLogger(LzoInputFormat.class);

  private final PathFilter hiddenPathFilter = new PathFilter() {
    // avoid hidden files and directories.
    @Override
    public boolean accept(Path path) {
      String name = path.getName();
      return !name.startsWith(".") &&
             !name.startsWith("_");
    }
  };

  private final PathFilter visibleLzoFilter = new PathFilter() {
    //applies to lzo files
    @Override
    public boolean accept(Path path) {
      String name = path.getName();
      return hiddenPathFilter.accept(path) &&
             name.endsWith(".lzo");
    }};

  private final PathFilter lzoIndexFilter = new PathFilter() {
    @Override
    public boolean accept(Path path) {
      final String name = path.getName();
      return hiddenPathFilter.accept(path) &&
             name.endsWith(LzoIndex.LZO_INDEX_SUFFIX);
    }
  };

  private static class LzoSplitStatus {
    private FileStatus lzoFileStatus;
    private FileStatus lzoIndexFileStatus;

    @Override
    public String toString() {
      return LzoSplitStatus.class.getName() + "[ lzo=" + lzoFileStatus
          + " index=" + lzoIndexFileStatus + " ]";
    }
  }

  private final Map<Path,LzoSplitStatus> splitStatusMap
      = new HashMap<Path,LzoSplitStatus>();

  @Override
  protected List<FileStatus> listStatus(JobContext job) throws IOException {
    // The list of files is no different.
    List<FileStatus> files = super.listStatus(job);
    List<FileStatus> results = Lists.newArrayList();
    boolean recursive = HadoopCompat.getConfiguration(job).getBoolean("mapred.input.dir.recursive", false);
    Iterator<FileStatus> it = files.iterator();
    while (it.hasNext()) {
      FileStatus fileStatus = it.next();
      FileSystem fs = fileStatus.getPath().getFileSystem(HadoopCompat.getConfiguration(job));
      addInputPath(results, fs, fileStatus, recursive);
    }
    if (LOG.isDebugEnabled()) {
      LOG.debug("Total lzo input paths to process : " + results.size());
    }
    return results;
  }

  //MAPREDUCE-1501
  /**
   * Add lzo file(s). If recursive is set, traverses the directories.
   * @param results
   *          The List to store all files.
   * @param fs
   *          The FileSystem.
   * @param pathStat
   *          The input path.
   * @param recursive
   *          Traverse in to directory
   * @throws IOException
   */
  protected void addInputPath(List<FileStatus> results, FileSystem fs,
                 FileStatus pathStat, boolean recursive) throws IOException {
    Path path = pathStat.getPath();
    if (pathStat.isDir()) {
      if (recursive) {
        for(FileStatus stat: fs.listStatus(path, hiddenPathFilter)) {
          addInputPath(results, fs, stat, recursive);
        }
      }
    } else if ( visibleLzoFilter.accept(path) ) {
      results.add(pathStat);
      lzoSplitStatus(path).lzoFileStatus = pathStat;
    } else if (lzoIndexFilter.accept(path)) {
      final String lzoIndexName = path.getName();
      final String lzoName = lzoIndexName.substring(0, lzoIndexName.length() -
          LzoIndex.LZO_INDEX_SUFFIX.length());
      final Path lzoPath = new Path(path.getParent(), lzoName);
      lzoSplitStatus(lzoPath).lzoIndexFileStatus = pathStat;
    }
  }

  private LzoSplitStatus lzoSplitStatus(Path path) {
    LzoSplitStatus lzoSplitStatus = splitStatusMap.get(path);
    if (lzoSplitStatus == null) {
      lzoSplitStatus = new LzoSplitStatus();
      splitStatusMap.put(path, lzoSplitStatus);
    }
    return lzoSplitStatus;
  }

  @Override
  protected boolean isSplitable(JobContext context, Path filename) {
    /* This should ideally return 'false'
     * and splitting should be handled completely in
     * this.getSplit(). Right now, FileInputFormat splits across the
     * blocks and this.getSplits() adjusts the positions.
     */
    LzoSplitStatus lzoSplitStatus = splitStatusMap.get(filename);
    return lzoSplitStatus != null && lzoSplitStatus.lzoIndexFileStatus != null;
  }

  private List<InputSplit> getSplitsInternal(JobContext job)
      throws IOException {
    List<InputSplit> defaultSplits = super.getSplits(job);

    // Find new starts and ends of the file splits that align with the lzo blocks.
    List<InputSplit> result = new ArrayList<InputSplit>();

    Path prevFile = null;
    LzoIndex prevIndex = null;

    for (InputSplit genericSplit : defaultSplits) {
      // Load the index.
      FileSplit fileSplit = (FileSplit)genericSplit;
      Path file = fileSplit.getPath();

      LzoIndex index; // reuse index for files with multiple blocks.
      final LzoSplitStatus lzoSplitStatus = splitStatusMap.get(file);
      if ( file.equals(prevFile) ) {
        index = prevIndex;
      } else {
        if (LOG.isDebugEnabled()) {
          LOG.debug("Cached file status " + lzoSplitStatus);
        }
        final FileStatus indexFileStatus = lzoSplitStatus.lzoIndexFileStatus;
        index = indexFileStatus == null
            ? new LzoIndex()
            : LzoIndex.readIndex(
                  file.getFileSystem(HadoopCompat.getConfiguration(job)), file);
        prevFile = file;
        prevIndex = index;
      }

      if (index == null) {
        // In listStatus above, a (possibly empty, but non-null) index was put in for every split.
        throw new IOException("Index not found for " + file);
      }

      if (index.isEmpty()) {
        // Empty index, so leave the default split.
        // split's start position should be 0.
        result.add(fileSplit);
        continue;
      }

      long start = fileSplit.getStart();
      long end = start + fileSplit.getLength();

      long lzoStart = index.alignSliceStartToIndex(start, end);
      long lzoEnd = index.alignSliceEndToIndex(end,
          lzoSplitStatus.lzoFileStatus.getLen());

      if (lzoStart != LzoIndex.NOT_FOUND  && lzoEnd != LzoIndex.NOT_FOUND) {
        result.add(new FileSplit(file, lzoStart, lzoEnd - lzoStart, fileSplit.getLocations()));
        if (LOG.isDebugEnabled()) {
          LOG.debug("Added LZO split for " + file + "[start=" + lzoStart + ", length=" + (lzoEnd - lzoStart) + "]");
        }
      }
      // else ignore the data?
      // should handle splitting the entire file here so that
      // such errors can be handled better.
    }

    return result;
  }

  @Override
  public List<InputSplit> getSplits(JobContext job) throws IOException {
    try {
      return getSplitsInternal(job);
    } finally {
      splitStatusMap.clear(); // no use beyond getSplits
    }
  }
}


package com.twitter.elephantbird.util;

import java.io.DataOutputStream;
import java.io.IOException;
import java.io.OutputStream;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.hadoop.compression.lzo.LzoIndex;
import com.hadoop.compression.lzo.LzoIndexer;
import com.hadoop.compression.lzo.LzopCodec;

/**
 * Miscellaneous lzo related utilities.
 */
public class LzoUtils {

  public static final Logger LOG = LoggerFactory.getLogger(LzoUtils.class);
  public static final String LZO_OUTPUT_INDEXABLE_MINSIZE =
      "elephantbird.lzo.output.indexable.minsize";
  public static final String LZO_OUTPUT_INDEX = "elephantbird.lzo.output.index";

  /**
   * A work-around to support environments with older versions of LzopCodec.
   * It might not be feasible for to select right version of hadoop-lzo
   * in some cases. This should be removed latest by EB-3.0.
   */
  private static boolean isLzopIndexSupported = false;
  static {
    try {
      isLzopIndexSupported =
        null != LzopCodec.class.getMethod("createIndexedOutputStream",
                                          OutputStream.class,
                                          DataOutputStream.class);
    } catch (Exception e) {
      // older version of hadoop-lzo.
    }
  }

  /**
   * Creates an lzop output stream. The index for the lzop is
   * also written to another at the same time if
   * elephantbird.lzo.output.index is set in configuration.

   *
   * If the file size at closing is not larger than a single block,
   * the index file is deleted (in line with {@link LzoIndexer} behavior).
   */
  public static DataOutputStream
  getIndexedLzoOutputStream(Configuration conf, Path path) throws IOException {

    LzopCodec codec = new LzopCodec();
    codec.setConf(conf);

    final Path file = path;
    final FileSystem fs = file.getFileSystem(conf);
    FSDataOutputStream fileOut = fs.create(file, false);

    FSDataOutputStream indexOut = null;
    if (conf.getBoolean(LZO_OUTPUT_INDEX, false)) {
      if ( isLzopIndexSupported ) {
        Path indexPath = file.suffix(LzoIndex.LZO_TMP_INDEX_SUFFIX);
        indexOut = fs.create(indexPath, false);
      } else {
        LOG.warn("elephantbird.lzo.output.index is enabled, but LzopCodec "
            + "does not have createIndexedOutputStream method. "
            + "Please upgrade hadoop-lzo.");
      }
    }

    final boolean isIndexed = indexOut != null;
    final long minIndexableSize = conf.getLong(LZO_OUTPUT_INDEXABLE_MINSIZE,
        -1L);

    OutputStream out = ( isIndexed ?
        codec.createIndexedOutputStream(fileOut, indexOut) :
        codec.createOutputStream(fileOut) );

    return new DataOutputStream(out) {
      // override close() to handle renaming index file.

      public void close() throws IOException {
        super.close();

        if ( isIndexed ) {
          // rename or remove the index file based on file size.

          Path tmpPath = file.suffix(LzoIndex.LZO_TMP_INDEX_SUFFIX);
          FileStatus stat = fs.getFileStatus(file);
          final long minSizeToIndex = minIndexableSize < 0
              ? stat.getBlockSize()
              : minIndexableSize;

          if (stat.getLen() <= minSizeToIndex) {
            fs.delete(tmpPath, false);
          } else {
            fs.rename(tmpPath, file.suffix(LzoIndex.LZO_INDEX_SUFFIX));
          }
        }
      }
    };
  }

}