package com.twitter.elephantbird.util;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.protobuf.ByteString;
import com.google.protobuf.DescriptorProtos;
import com.google.protobuf.DescriptorProtos.FieldDescriptorProto.Label;
import com.google.protobuf.DescriptorProtos.FieldDescriptorProto.Type;
import com.google.protobuf.Descriptors;
import com.google.protobuf.Descriptors.DescriptorValidationException;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.Message;

import org.apache.thrift.TBase;
import org.apache.thrift.TFieldIdEnum;
import org.apache.thrift.protocol.TType;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.twitter.elephantbird.thrift.TStructDescriptor;
import com.twitter.elephantbird.thrift.TStructDescriptor.Field;
import com.twitter.elephantbird.util.Pair;

/**
 * Translates a Thrift object into a Protocol Buffer message.
 *

 * For most data types, there is a 1:1 mapping between Thrift
 * and Protobufs, except for:

 *
 *
 enums: converted to their String representations.
 *
 lists: added to the enclosing struct as a repeated field.
 *
 sets: protobufs doesn't have native support for them so they're
 * treated as lists and we rely on thrift to enforce uniqueness.
 *
 maps: protobufs also doesn't have native support for these so we
 * create an intermediate Message type with key and value fields of the
 * appropriate types and add these as repeated fields of the enclosing
 * struct.

 *
 * @param <T> Source thrift class
 */
public class ThriftToDynamicProto<T extends TBase<?, ?>> {
  private static final Logger LOG = LoggerFactory.getLogger(ThriftToDynamicProto.class);
  private static final List<Pair<String, Type>> EMPTY_FIELDS = new ArrayList<Pair<String, Type>>();

  private static final String MAP_KEY_FIELD_NAME   = "key";
  private static final String MAP_VALUE_FIELD_NAME = "value";

  private boolean supportNestedObjects = false;
  private boolean ignoreUnsupportedTypes = false;

  private final Descriptors.FileDescriptor fileDescriptor;

  // a map of descriptors keyed to their typeName
  private Map<String, DescriptorProtos.DescriptorProto.Builder> descriptorBuilderMap
      = Maps.newHashMap();

  // a map of message builders keyed to their typeName
  private Map<String, DynamicMessage.Builder> messageBuilderMap = Maps.newHashMap();

  /**
   * Create a Dynamic Protocol Buffer message with fields equivalent to those in the provided
   * Thrift class.
   * @param thriftClass
   * @throws DescriptorValidationException
   */
  public ThriftToDynamicProto(Class<T> thriftClass) throws DescriptorValidationException {
    this(thriftClass, new ArrayList<Pair<String, Type>>());
  }

  public ThriftToDynamicProto(Class<T> thriftClass,
                              boolean supportNestedObjects,
                              boolean ignoreUnsupportedTypes) throws DescriptorValidationException {
    this(thriftClass, new ArrayList<Pair<String, Type>>(), supportNestedObjects,
        ignoreUnsupportedTypes);
  }

  /**
   * Create a converter that produces a Dynamic Protocol Buffer message with fields equivalent to
   * those in the provided Thrift class, plus additional fields as defined by extraFields. Nested
   * Thrift Structs will be ignored, but other unsupported types will throw an exception.
   * @param thriftClass class to use to generate the protobuf schema
   * @param extraFields a list of pairs of (fieldName, protobufType) that defines extra fields to
   * add to the generated message.
   * @throws DescriptorValidationException
   */
  public ThriftToDynamicProto(Class<T> thriftClass, List<Pair<String, Type>> extraFields)
    throws DescriptorValidationException {
    this(thriftClass, extraFields, false, false);
  }

  /**
   * Create a converter that produces a Dynamic Protocol Buffer message with fields equivalent to
   * those in the provided Thrift class, plus additional fields as defined by extraFields.
   * @param thriftClass class to use to generate the protobuf schema
   * @param extraFields a list of pairs of (fieldName, protobufType) that defines extra fields to
   * add to the generated message.
   * @param supportNestedObjects if true, builds a nested structure. Defaults is false, which will
   * silently ignore nested objects, which include structs, lists/sets of structs and maps of all
   * types.  Lists/sets of base types are not considered to be nested objects.
   * @param ignoreUnsupportedTypes if true, ignores types that aren't supported. If false an
   * exception will be thrown when an unsupported type is encountered. Default is false.
   * @throws DescriptorValidationException
   */
  public ThriftToDynamicProto(Class<T> thriftClass,
                              List<Pair<String, Type>> extraFields,
                              boolean supportNestedObjects,
                              boolean ignoreUnsupportedTypes) throws DescriptorValidationException {
    this.supportNestedObjects = supportNestedObjects;
    this.ignoreUnsupportedTypes = ignoreUnsupportedTypes;

    // setup the descriptor proto builder for the top-level class
    DescriptorProtos.DescriptorProto.Builder desBuilder
        = DescriptorProtos.DescriptorProto.newBuilder();
    desBuilder.setName(protoMessageType(thriftClass));
    descriptorBuilderMap.put(desBuilder.getName(), desBuilder);

    // convert the thrift schema to a proto schema
    thriftToProtoSchema(desBuilder, TStructDescriptor.getInstance(thriftClass), extraFields);

    // add all of the message types to the file descriptor proto builder
    DescriptorProtos.FileDescriptorProto.Builder fileDescProtoBuilder =
      DescriptorProtos.FileDescriptorProto.newBuilder();
    for (DescriptorProtos.DescriptorProto.Builder builder : descriptorBuilderMap.values()) {
      fileDescProtoBuilder.addMessageType(builder);
    }

    // create dynamic message builders to be cloned for all types
    Descriptors.FileDescriptor dynamicDescriptor = Descriptors.FileDescriptor.buildFrom(
      fileDescProtoBuilder.build(), new Descriptors.FileDescriptor[0]);

    for (String type : descriptorBuilderMap.keySet()) {
      Descriptors.Descriptor msgDescriptor = dynamicDescriptor.findMessageTypeByName(type);
      messageBuilderMap.put(type, DynamicMessage.newBuilder(msgDescriptor));
    }

    fileDescriptor = dynamicDescriptor;
  }

  /**
   * For the given thriftClass, return a Protobufs builder to build a similar protobuf class.
   * @param thriftClass The thrift class for which the builder is desired.
   * @return a protobuf message builder
   */
  public Message.Builder getBuilder(Class<? extends TBase<?, ?>> thriftClass) {
    return messageBuilderMap.get(protoMessageType(thriftClass)).clone();
  }

  /**
   * Return Protobufs builder for a Map field
   */
  private Message.Builder mapEntryProtoBuilder(TStructDescriptor descriptor, Field field) {
    return messageBuilderMap.get(mapProtoMessageType(descriptor, field)).clone();
  }

  private void thriftToProtoSchema(DescriptorProtos.DescriptorProto.Builder desBuilder,
      TStructDescriptor fieldDesc, List<Pair<String, Type>> extraFields)
      throws DescriptorValidationException {
    int maxThriftId = doSchemaMapping(desBuilder, fieldDesc);

    // handle extra fields if they exist. Only supported on the top level message
    int extraFieldIdx = maxThriftId + 1;
    for (Pair<String, Type> extraField : extraFields) {
      addProtoField(desBuilder,
                    extraField.getFirst(),
                    ++extraFieldIdx,
                    extraField.getSecond(),
                    false);
    }
  }

  private int doSchemaMapping(DescriptorProtos.DescriptorProto.Builder desBuilder,
      TStructDescriptor fieldDesc) throws DescriptorValidationException {
    int maxThriftId = 0;
    for (Field tField : fieldDesc.getFields()) {
      maxThriftId = Math.max(tField.getFieldId(), maxThriftId);

      if (supportNestedObjects && tField.isMap()) {
        String typeName = mapProtoMessageType(fieldDesc, tField);
        if (descriptorBuilderMap.get(typeName) == null) {
          DescriptorProtos.DescriptorProto.Builder mapBuilder =
            mapDescriptorProtoBuilder(tField, typeName);

          descriptorBuilderMap.put(typeName, mapBuilder);
          addProtoField(desBuilder, tField.getName(), tField.getFieldId() + 1, typeName, true);
        }
      } else {
        Field field = resolveField(tField);
        Type protoType = thriftTypeToProtoType(field);
        boolean isContainer = isContainer(tField);

        if (supportNestedObjects && protoType == Type.TYPE_MESSAGE) {
          String typeName = resolveMessageTypeName(field.gettStructDescriptor());

          // Protobuf field ids start at 1. Thrift starts at 0.
          addProtoField(desBuilder,
                        tField.getName(),
                        tField.getFieldId() + 1,
                        typeName,
                        isContainer);
        } else if (protoType != null) {
          if (supportNestedObjects
             || (!supportNestedObjects && !hasNestedObject(tField))) {
            addProtoField(desBuilder,
                          tField.getName(),
                          tField.getFieldId() + 1,
                          protoType,
                          isContainer);
          }
        }
      }
    }

    return maxThriftId;
  }

  // When dealing with a Set/List, we want to retrieve the relevant Field type
  private Field resolveField(Field inputField) {
    if (inputField.isList()) {
      return inputField.getListElemField();
    } else if (inputField.isSet()) {
      return inputField.getSetElemField();
    } else {
      return inputField;
    }
  }

  /**
   * Determines whether a field is considered to be a nested object based on:
   * - whether the field itself is a struct
   * - whether the field is a list/set of structs
   * - whether field is a Map
   */
  private boolean hasNestedObject(Field field) {
    return field.isStruct()
      || (field.isList() && field.getListElemField().isStruct())
      || (field.isSet() && field.getSetElemField().isStruct())
      || field.isMap();
  }

  /**
   * Generate a DescriptorProto.Builder for the Message type that will be used
   * to represent the entries of the input Map field.
   *
   * @param field a Map Field (field.isMap() == true)
   * @param typeName name of new message type
   */
  private DescriptorProtos.DescriptorProto.Builder mapDescriptorProtoBuilder(
    Field field, String typeName) throws DescriptorValidationException {
    DescriptorProtos.DescriptorProto.Builder mapBuilder =
      DescriptorProtos.DescriptorProto.newBuilder().setName(typeName);

    Field keyField = field.getMapKeyField();
    Field valueField = field.getMapValueField();

    DescriptorProtos.FieldDescriptorProto.Builder keyBuilder = mapKeyProtoBuilder();
    DescriptorProtos.FieldDescriptorProto.Builder valueBuilder = mapValueProtoBuilder();

    setBuilderTypeFromField(keyField, keyBuilder);
    setBuilderTypeFromField(valueField, valueBuilder);

    mapBuilder.addField(keyBuilder.build());
    mapBuilder.addField(valueBuilder.build());

    return mapBuilder;
  }

  // field descriptor for key field to be used in Message used to hold Map entries
  private DescriptorProtos.FieldDescriptorProto.Builder mapKeyProtoBuilder() {
    return fieldDescriptorProtoBuilder(MAP_KEY_FIELD_NAME, 1).setLabel(Label.LABEL_REQUIRED);
  }

  // field descriptor for value field to be used in Message used to hold Map entries
  private DescriptorProtos.FieldDescriptorProto.Builder mapValueProtoBuilder() {
    return fieldDescriptorProtoBuilder(MAP_VALUE_FIELD_NAME, 2).setLabel(Label.LABEL_REQUIRED);
  }

  private void setBuilderTypeFromField(
      Field field,
      DescriptorProtos.FieldDescriptorProto.Builder builder
   ) throws DescriptorValidationException {
    Type valueProtoType = thriftTypeToProtoType(field);
    if (valueProtoType == Type.TYPE_MESSAGE) {
      builder.setTypeName(resolveMessageTypeName(field.gettStructDescriptor()));
    } else if (valueProtoType != null) {
      builder.setType(valueProtoType);
    }
  }

  /**
   * For a TStructDescriptor, resolves the typeName and optionally converts and memoizes it's
   * schema.
   */
  private String resolveMessageTypeName(TStructDescriptor descriptor)
    throws DescriptorValidationException {
    String typeName = protoMessageType(descriptor.getThriftClass());

    // Anytime we have a new message typeName, we make sure that we have a builder for it.
    // If not, we create one.
    DescriptorProtos.DescriptorProto.Builder builder = descriptorBuilderMap.get(typeName);
    if (builder == null) {
      builder = DescriptorProtos.DescriptorProto.newBuilder();
      builder.setName(typeName);
      descriptorBuilderMap.put(typeName, builder);
      doSchemaMapping(builder, descriptor);
    }

    return typeName;
  }

  private void addProtoField(DescriptorProtos.DescriptorProto.Builder builder, String name,
                             int fieldIdx, Type type, boolean isRepeated) {
    DescriptorProtos.FieldDescriptorProto.Builder fdBuilder
        = fieldDescriptorProtoBuilder(name, fieldIdx).setType(type);
    if (isRepeated) {
      fdBuilder.setLabel(Label.LABEL_REPEATED);
    }
    builder.addField(fdBuilder.build());
  }

  private void addProtoField(DescriptorProtos.DescriptorProto.Builder builder, String name,
                             int fieldIdx, String type, boolean isRepeated) {
    DescriptorProtos.FieldDescriptorProto.Builder fdBuilder
        = fieldDescriptorProtoBuilder(name, fieldIdx).setTypeName(type);
    if (isRepeated) {
      fdBuilder.setLabel(Label.LABEL_REPEATED);
    }
    builder.addField(fdBuilder.build());
  }

  private DescriptorProtos.FieldDescriptorProto.Builder fieldDescriptorProtoBuilder(
      String name, int fieldIdx) {
    DescriptorProtos.FieldDescriptorProto.Builder fdBuilder =
      DescriptorProtos.FieldDescriptorProto.newBuilder()
      .setName(name)
      .setNumber(fieldIdx);
    return fdBuilder;
  }

  /**
   * Convert a thrift object to a protobuf message.
   * @param thriftObj thrift object
   * @return protobuf protobuf message
   */
  @SuppressWarnings("unchecked")
  public Message convert(T thriftObj) {
    return doConvert((TBase<?, ?>)
                     Preconditions.checkNotNull(thriftObj, "Can not convert a null object"));
  }

  /**
   * conver TBase object to Message object
   * @param thriftObj
   */
  @SuppressWarnings("unchecked")
  public <F extends TFieldIdEnum> Message doConvert(TBase<?, F> thriftObj) {
    if (thriftObj == null) { return null; }

    Class<TBase<?, F>> clazz = (Class<TBase<?, F>>) thriftObj.getClass();
    checkState(clazz);

    Message.Builder builder = getBuilder(clazz);

    TStructDescriptor fieldDesc = TStructDescriptor.getInstance(clazz);
    int fieldId = 0;
    for (Field tField : fieldDesc.getFields()) {
      // don't want to carry over default values from unset fields
      if (!thriftObj.isSet((F) tField.getFieldIdEnum())
          || (!supportNestedObjects && hasNestedObject(tField))) {
        fieldId++;
        continue;
      }

      // recurse into the object if it's a struct, otherwise just add the field
      if (supportNestedObjects && tField.getType() == TType.STRUCT) {
        TBase<?, ?> fieldValue = (TBase<?, ?>) fieldDesc.getFieldValue(fieldId++, thriftObj);
        Message message = doConvert(fieldValue);
        if (message != null) {
          FieldDescriptor protoFieldDesc = builder.getDescriptorForType().findFieldByName(
              tField.getName());
          builder.setField(protoFieldDesc, message);
        }
      } else {
        fieldId = convertField(thriftObj, builder, fieldDesc, fieldId, tField);
      }
    }
    return builder.build();
  }

  private void checkState(Class<? extends TBase<?, ?>> thriftClass) {
    Preconditions.checkState(hasBuilder(thriftClass),
      "No message builder found for thrift class: " + thriftClass.getCanonicalName());
  }

  private boolean hasBuilder(Class<? extends TBase<?, ?>> thriftClass) {
    return messageBuilderMap.get(protoMessageType(thriftClass)) != null;
  }

  private Object sanitizeRawValue(Object value, Field tField) {
    Object returnValue = value;

    if (tField.isEnum()) {
      // TODO: proper enum handling
      returnValue = returnValue.toString();
    } else if (tField.isBuffer()) {
      returnValue = ByteString.copyFrom((byte[]) returnValue);
    }

    if (returnValue instanceof Byte) {
      returnValue = new Integer((Byte) returnValue);
    } else if (returnValue instanceof Short) {
      returnValue = new Integer((Short) returnValue);
    }

    return returnValue;
  }

  /*
   * Determines if the field in question is a Set or List of a Struct type
   */
  private boolean isStructContainer(Field tField) {
    return (tField.isList() && tField.getListElemField().isStruct())
      || (tField.isSet() && tField.getSetElemField().isStruct());
  }

  @SuppressWarnings("unchecked")
  private int convertField(TBase<?, ?> thriftObj, Message.Builder builder,
                           TStructDescriptor fieldDesc, int fieldId, Field tField) {
    int tmpFieldId = fieldId;
    FieldDescriptor protoFieldDesc = builder.getDescriptorForType().findFieldByName(
        tField.getName());

    if (protoFieldDesc == null) {
      // not finding a field might be ok if we're ignoring an unsupported types
      Type protoType = thriftTypeToProtoType(tField);
      if (protoType == null
          && (ignoreUnsupportedTypes
              || (!supportNestedObjects && hasNestedObject(tField)))) {
        return tmpFieldId; // no-op
      }
      throw new RuntimeException("Field " + tField.getName() + " not found in dynamic protobuf.");
    }

    Object fieldValue = fieldDesc.getFieldValue(tmpFieldId++, thriftObj);
    if (fieldValue == null) {
      return tmpFieldId;
    }

    try {
      // For non-Map containers that contain struct types,
      // we have to convert each struct into a Message.
      if (isStructContainer(tField)) {
        List<Message> convertedStructs = Lists.newLinkedList();

        Iterable<TBase<?, ?>> structIterator = (Iterable<TBase<?, ?>>) fieldValue;
        for (TBase<?, ?> struct : structIterator) {
          convertedStructs.add(doConvert(struct));
        }

        fieldValue = convertedStructs;
      } else if (tField.isMap()) {
        List<Message> convertedMapEntries = Lists.newLinkedList();

        Map<?, ?> rawMap = (Map) fieldValue;
        for (Map.Entry<?, ?> entry : rawMap.entrySet()) {
          Message.Builder mapBuilder = mapEntryProtoBuilder(fieldDesc, tField);
          Message msg = buildMapEntryMessage(mapBuilder, tField, entry.getKey(), entry.getValue());
          convertedMapEntries.add(msg);
        }

        fieldValue = convertedMapEntries;
      } else {
        // protobufs throws an exception if you try to set byte on an int32 field so we need to
        // convert it to an Integer before it gets set
        fieldValue = sanitizeRawValue(fieldValue, tField);
      }

      // Container types have to be added as repeated fields
      if (isContainer(tField)) {
        Iterable<?> container = (Iterable) fieldValue;
        for (Object obj : container) {
          builder.addRepeatedField(protoFieldDesc, obj);
        }
      } else {
        builder.setField(protoFieldDesc, fieldValue);
      }
    } catch (IllegalArgumentException e) {
      LOG.error(String.format("Could not set protoField(index=%d, name=%s, type=%s) with "
        + "thriftField(index=%d, name=%s, type=%d, value=%s)",
        protoFieldDesc.getIndex(), protoFieldDesc.getName(), protoFieldDesc.getType(),
          tmpFieldId - 1, tField.getName(), tField.getType(), fieldValue), e);
      throw e;
    }

    return tmpFieldId;
  }

  /**
   * Builds a Message that contains the key value pair of a Map entry
   */
  private Message buildMapEntryMessage(Message.Builder mapBuilder, Field field,
                                       Object mapKey, Object mapValue) {
    FieldDescriptor keyFieldDescriptor =
      mapBuilder.getDescriptorForType().findFieldByName(MAP_KEY_FIELD_NAME);
    FieldDescriptor valueFieldDescriptor =
      mapBuilder.getDescriptorForType().findFieldByName(MAP_VALUE_FIELD_NAME);

    boolean isKeyStruct = field.getMapKeyField().isStruct();
    boolean isValueStruct = field.getMapValueField().isStruct();

    Object convertedKey;
    if (isKeyStruct) {
      convertedKey = doConvert((TBase<?, ?>) mapKey);
    } else {
      convertedKey = sanitizeRawValue(mapKey, field.getMapKeyField());
    }

    Object convertedValue;
    if (isValueStruct) {
      convertedValue = doConvert((TBase<?, ?>) mapValue);
    } else {
      convertedValue = sanitizeRawValue(mapValue, field.getMapValueField());
    }

    mapBuilder.setField(keyFieldDescriptor, convertedKey);
    mapBuilder.setField(valueFieldDescriptor, convertedValue);

    return mapBuilder.build();
  }

  // Checks a field to determine whether it's a List/Set/Map
  private boolean isContainer(Field field) {
    return field.isSet() || field.isList() || field.isMap();
  }

  private Type thriftTypeToProtoType(Field tField) {
    byte thriftType = tField.getType();
    switch (thriftType)  {
      case TType.BOOL:
        return Type.TYPE_BOOL;
      case TType.BYTE:
        return Type.TYPE_INT32;
      case TType.DOUBLE:
        return Type.TYPE_DOUBLE;
      case TType.I16:
        return Type.TYPE_INT32;
      case TType.I32:
        return Type.TYPE_INT32;
      case TType.I64:
        return Type.TYPE_INT64;
      case TType.STRING:
        // Thrift thinks bytes and strings are interchangeable. Protocol buffers are not insane.
        return tField.isBuffer() ? Type.TYPE_BYTES : Type.TYPE_STRING;
      case TType.ENUM:
        // TODO: proper enum handling. For now, convert to strings.
        return Type.TYPE_STRING;
      case TType.STRUCT:
        if (supportNestedObjects) { return Type.TYPE_MESSAGE; }
        return null;
      case TType.MAP:
        return null;
      case TType.SET:
        return null;
      case TType.LIST:
        return null;
      default:
        if (ignoreUnsupportedTypes) {
          LOG.warn("Thrift type " + thriftType + " not supported for field "
             + tField.getName() + ". Ignoring");
          return null;
        }
        throw new IllegalArgumentException("Can't map Thrift type " + thriftType
          + " to a Protobuf type for field: " + tField.getName());
    }
  }

  // name the proto message type after the thrift class name. Dots are not permitted in protobuf
  // names
  private String protoMessageType(Class<? extends TBase<?, ?>> thriftClass) {
    return thriftClass.getCanonicalName().replace(".", "_");
  }

  /**
   * name the proto message used for Map types after the thrift class name of the enclosing
   * struct and the field name
   */
  private String mapProtoMessageType(TStructDescriptor descriptor, Field field) {
    return String.format("%s_%s", protoMessageType(descriptor.getThriftClass()), field.getName());
  }

  // Given the class name, finds the corresponding Descriptor and return the appropriate
  // FieldDescriptor
  public FieldDescriptor getFieldDescriptor(Class<? extends TBase<?, ?>> thriftClass,
                                            String fieldName) {
    checkState(thriftClass);
    Descriptors.Descriptor descriptor = getBuilder(thriftClass).getDescriptorForType();
    return descriptor.findFieldByName(fieldName);
  }

  // Picks off the FileDescriptor for this instance
  public Descriptors.FileDescriptor getFileDescriptor() {
    return fileDescriptor;
  }
}


package com.twitter.elephantbird.pig.load;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.pig.Expression;
import org.apache.pig.LoadFunc;
import org.apache.pig.LoadMetadata;
import org.apache.pig.LoadPushDown;
import org.apache.pig.ResourceSchema;
import org.apache.pig.ResourceStatistics;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.PigSplit;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.util.ObjectSerializer;
import org.apache.pig.impl.util.Pair;
import org.apache.pig.impl.util.UDFContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.twitter.elephantbird.mapreduce.io.BinaryWritable;
import com.twitter.elephantbird.pig.util.PigCounterHelper;
import com.twitter.elephantbird.util.TypeRef;

/**
 * Provides common functionality & helpers required by all elephantbird loaders.
 * TODO: Rename "BaseLoadFunc" as this class has nothing to do with LZO.
 */
public abstract class LzoBaseLoadFunc extends LoadFunc implements LoadMetadata, LoadPushDown {
  private static final Logger LOG = LoggerFactory.getLogger(LzoBaseLoadFunc.class);

  @SuppressWarnings("rawtypes")
  protected RecordReader reader;

  // Making accessing Hadoop counters from Pig slightly more convenient.
  private final PigCounterHelper counterHelper_ = new PigCounterHelper();

  protected Configuration jobConf;
  protected String contextSignature;
  protected static final String projectionKey = "LzoBaseLoadFunc_projectedFields";

  protected RequiredFieldList requiredFieldList = null;

  /**
   * Construct a new load func.
   */
  public LzoBaseLoadFunc() {
  }

  /**
   * A convenience function for working with Hadoop counter objects from load functions.  The Hadoop
   * reporter object isn't always set up at first, so this class provides brief buffering to ensure
   * that counters are always recorded.
   */
  protected void incrCounter(String group, String counter, long incr) {
    counterHelper_.incrCounter(group, counter, incr);
  }

  /**
   * A convenience function for working with Hadoop counter objects from load functions.  The Hadoop
   * reporter object isn't always set up at first, so this class provides brief buffering to ensure
   * that counters are always recorded.
   */
  protected void incrCounter(Enum<?> key, long incr) {
    counterHelper_.incrCounter(key, incr);
  }

  /** same as incrCounter(pair.first, pair.second, incr). */
  protected void incrCounter(Pair<String, String> groupCounterPair, long incr) {
    counterHelper_.incrCounter(groupCounterPair.first, groupCounterPair.second, incr);
  }

  /** UDF properties for this class based on context signature */
  protected Properties getUDFProperties() {
    return UDFContext.getUDFContext()
        .getUDFProperties(this.getClass(), new String[] {contextSignature});
  }

  @Override
  public void setUDFContextSignature(String signature) {
    this.contextSignature = signature;
  }

  @Override
  public void setLocation(String location, Job job) throws IOException {
    FileInputFormat.setInputPaths(job, location);
    this.jobConf = HadoopCompat.getConfiguration(job);

    String projectedFields = getUDFProperties().getProperty(projectionKey);
    if (projectedFields != null) {
      requiredFieldList =
        (RequiredFieldList) ObjectSerializer.deserialize(projectedFields);
    }
  }

  /**
   * A utility method for loaders that read {@link BinaryWritable} values.
   */
  protected <M> M getNextBinaryValue(TypeRef<M> typeRef) throws IOException {
    //typeRef is just to help compiler resolve the type at compile time.
    try {
      if (reader != null && reader.nextKeyValue()) {
        @SuppressWarnings("unchecked")
        BinaryWritable<M> writable = (BinaryWritable<M>)reader.getCurrentValue();
        return writable.get();
      }
    } catch (InterruptedException e) {
      LOG.error("InterruptedException encountered, bailing.", e);
      throw new IOException(e);
    }

    return null;
  }

  // LoadPushDown implementation:

  @Override
  public List<OperatorSet> getFeatures() {
    return Arrays.asList(LoadPushDown.OperatorSet.PROJECTION);
  }

  @Override
  public RequiredFieldResponse pushProjection(
      RequiredFieldList requiredFieldList) throws FrontendException {
    // the default implementation disables the feature.
    // a projection needs to be explicitly supported by sub classes.
    return null;
  }

  /**
   * A helper method for implementing
   * {@link LoadPushDown#pushProjection(RequiredFieldList)}. <p>
   *
   * Stores requiredFieldList in context. The requiredFields are read from
   * context on the backend (in side {@link #setLocation(String, Job)}).
   */
  protected RequiredFieldResponse pushProjectionHelper(
                                          RequiredFieldList requiredFieldList)
                                          throws FrontendException {
    try {
      getUDFProperties().setProperty(projectionKey,
                                     ObjectSerializer.serialize(requiredFieldList));
    } catch (IOException e) { // not expected
      throw new FrontendException(e);
    }

    return new RequiredFieldResponse(true);
  }

  @Override
  public void prepareToRead(@SuppressWarnings("rawtypes") RecordReader reader, PigSplit split) {
      this.reader = reader;
  }

  @Override
  public ResourceSchema getSchema(String filename, Job job) throws IOException {
    // most loaders are expected to override this.
    return null;
  }

  /*
   * NOT IMPLEMENTED
   */
  @Override
  public String[] getPartitionKeys(String arg0, Job arg1) throws IOException {
    return null;
  }

  /*
   * NOT IMPLEMENTED
   */
  @Override
  public ResourceStatistics getStatistics(String arg0, Job arg1) throws IOException {
    return null;
  }

  /*
   * NOT IMPLEMENTED
   */
  @Override
  public void setPartitionFilter(Expression arg0) throws IOException {
  }
}


package com.twitter.elephantbird.pig.store;

import java.io.IOException;
import java.util.Properties;

import org.apache.hadoop.fs.Path;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.pig.StoreFunc;
import org.apache.pig.impl.util.Pair;
import org.apache.pig.impl.util.UDFContext;

import com.twitter.elephantbird.pig.util.PigCounterHelper;

/**
 * Base class for StoreFunc implementations. Implements some of the trivial interface methods.
 */
public abstract class BaseStoreFunc extends StoreFunc {

  private final PigCounterHelper counterHelper = new PigCounterHelper();
  @SuppressWarnings("rawtypes")
  protected RecordWriter writer;
  protected String contextSignature;

  /**
   * A convenience function for working with Hadoop counter objects from load functions. The Hadoop
   * reporter object isn't always set up at first, so this class provides brief buffering to ensure
   * that counters are always recorded.
   */
  protected void incrCounter(String group, String counter, long incr) {
    counterHelper.incrCounter(group, counter, incr);
  }

  /**
   * A convenience function for working with Hadoop counter objects from load functions. The Hadoop
   * reporter object isn't always set up at first, so this class provides brief buffering to ensure
   * that counters are always recorded.
   */
  protected void incrCounter(Enum<?> key, long incr) {
    counterHelper.incrCounter(key, incr);
  }

  /** same as incrCounter(pair.first, pair.second, incr). */
  protected void incrCounter(Pair<String, String> groupCounterPair, long incr) {
    counterHelper.incrCounter(groupCounterPair.first, groupCounterPair.second, incr);
  }

  @Override
  public void setStoreFuncUDFContextSignature(String signature) {
    this.contextSignature = signature;
  }

  /** UDF properties for this class based on context signature */
  protected Properties getUDFProperties() {
    return UDFContext.getUDFContext()
        .getUDFProperties(this.getClass(), new String[] { contextSignature });
  }

  @Override
  public void prepareToWrite(@SuppressWarnings("rawtypes") RecordWriter writer) throws IOException {
    this.writer = writer;
  }

  @Override
  public void setStoreLocation(String location, Job job) throws IOException {
    FileOutputFormat.setOutputPath(job, new Path(location));
  }

  /**
   * Invokes writer.write(key, value). <br>
   * The main purpose is to catch {@link InterruptedException} and covert it to an
   * {@link IOException}. Avoids unnecessary try-catch dance in putNext() implementations.
   */
  @SuppressWarnings("unchecked")
  final protected void writeRecord(Object key, Object value) throws IOException {
    try {
      writer.write(key, value);
    } catch (InterruptedException e) {
      throw new IOException(e);
    }
  }
}


package com.twitter.elephantbird.util;

import java.io.IOException;
import java.util.List;

import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.ContentSummary;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.fs.PathFilter;

/**
 * HDFS utilities
 */
public final class HdfsUtils {
  private HdfsUtils() { }

  /**
   * Converts a path to a qualified string
   */
  public static class PathToQualifiedString implements Function<Path, String> {
    private Configuration conf;

    public PathToQualifiedString(Configuration conf) {
      this.conf = Preconditions.checkNotNull(conf);
    }

    @Override
    public String apply(Path path) {
      try {
        return path.getFileSystem(conf).makeQualified(path).toString();
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
    }
  }

  /**
   * Used by {@link HdfsUtils#walkPath} to 'visit' or process a
   * path
   */
  public static interface PathVisitor {
    void visit(FileStatus fileStatus);
  }

  /**
   * Recursively walk a path applying visitor to each path accepted by
   * filter
   *
   * @param path root path to begin walking, will be visited if
   *             it passes the filter and directory flag
   * @param fs FileSystem for this path
   * @param filter filter to determine which paths to accept
   * @param visitor visitor to apply to each accepted path
   * @throws IOException
   */
  public static void walkPath(Path path,
                              FileSystem fs,
                              PathFilter filter,
                              PathVisitor visitor) throws IOException {

    FileStatus fileStatus = fs.getFileStatus(path);

    if (filter.accept(path)) {
      visitor.visit(fileStatus);
    }

    if (fileStatus.isDir()) {
      FileStatus[] children = fs.listStatus(path);
      for (FileStatus childStatus : children) {
        walkPath(childStatus.getPath(), fs, filter, visitor);
      }
    }
  }

  /**
   * Recursively walk a path, adding paths that are accepted by filter to accumulator
   *
   * @param path root path to begin walking, will be added to accumulator
   * @param fs FileSystem for this path
   * @param filter filter to determine which paths to accept
   * @param accumulator all paths accepted will be added to accumulator
   * @throws IOException
   */
  public static void collectPaths(Path path,
                                  FileSystem fs,
                                  PathFilter filter,
                                  final List<Path> accumulator) throws IOException {

    walkPath(path, fs, filter, new PathVisitor() {
      @Override
      public void visit(FileStatus fileStatus) {
        accumulator.add(fileStatus.getPath());
      }
    });
  }

  private static class PathSizeVisitor implements PathVisitor {
    private long size = 0;

    @Override
    public void visit(FileStatus fileStatus) {
      size += fileStatus.getLen();
    }

    public long getSize() {
      return size;
    }
  }

  /**
   * Calculates the total size of all the contents of a directory that are accepted
   * by filter. All subdirectories will be searched recursively and paths in subdirectories
   * that are accepted by filter will also be counted.
   *
   * Does not include the size of directories themselves
   * (which are 0 in HDFS but may not be 0 on local file systems)
   *
   * To get the size of a directory without filtering, use
   * {@link #getDirectorySize(Path, FileSystem)} which is much more efficient.
   *
   * @param path path to recursively walk
   * @param fs FileSystem for this path
   * @param filter path filter for which paths size's to include in the total
   *               NOTE: you do *not* need to filter out directories, this will be done for you
   * @return size of the directory in bytes
   * @throws IOException
   */
  public static long getDirectorySize(Path path, FileSystem fs, PathFilter filter)
      throws IOException {
    PathSizeVisitor visitor = new PathSizeVisitor();
    PathFilter composite = new PathFilters.CompositePathFilter(
      PathFilters.newExcludeDirectoriesFilter(fs.getConf()),
      filter);
    walkPath(path, fs, composite, visitor);
    return visitor.getSize();
  }

  /**
   * Calculates the total size of all the contents of a directory,
   * including the contents of all of its subdirectories.
   * Does not include the size of directories themselves
   * (which are 0 in HDFS but may not be 0 on local file systems)
   *
   * @param path path to recursively walk
   * @param fs FileSystem for this path
   * @return size of the directory's contents in bytes
   * @throws IOException
   */
  public static long getDirectorySize(Path path, FileSystem fs) throws IOException {
    ContentSummary cs = fs.getContentSummary(path);
    return cs.getLength();
  }

  /**
   * Given a list of paths that (potentially) have glob syntax in them,
   * return a list of paths with all the globs expanded.
   *
   * @param pathsWithGlobs a list of paths that may or may not have glob syntax in them
   * @param conf job conf
   * @return an equivalent list of paths with no glob syntax in them
   * @throws IOException
   */
  public static List<Path> expandGlobs(List<String> pathsWithGlobs, Configuration conf)
    throws IOException {

    List<Path> paths = Lists.newLinkedList();
    for (String pathStr : pathsWithGlobs) {
      Path path = new Path(pathStr);
      FileSystem fs = path.getFileSystem(conf);
      FileStatus[] statuses = fs.globStatus(path);
      // some versions of hadoop return null for non-existent paths
      if (statuses != null) {
        for (FileStatus status : statuses) {
          paths.add(status.getPath());
        }
      }
    }
    return paths;
  }
}


package com.twitter.elephantbird.pig.load;

import java.io.IOException;
import java.util.Map;

import org.apache.hadoop.io.MapWritable;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.pig.Expression;
import org.apache.pig.LoadMetadata;
import org.apache.pig.PigException;
import org.apache.pig.ResourceSchema;
import org.apache.pig.ResourceStatistics;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.DataType;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
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;

import com.google.common.collect.Maps;
import com.twitter.elephantbird.mapreduce.input.LzoW3CLogInputFormat;
import com.twitter.elephantbird.mapreduce.input.LzoW3CLogRecordReader;

/**
 * Load w3c log LZO file line by line, passing each line as a single-field Tuple to Pig.
 */
public class LzoW3CLogLoader extends LzoBaseLoadFunc implements LoadMetadata {
  protected static final Logger LOG = LoggerFactory.getLogger(LzoW3CLogLoader.class);

  protected static final TupleFactory tupleFactory_ = TupleFactory.getInstance();
  protected final String fileURI;
  protected enum LzoW3CLogLoaderCounters { LinesW3CDecoded, UnparseableLines};

  /**
   * Constructor.
   * @param fileURI path to HDFS file that contains the CRC to column list mappings, one per line.
   * @throws IOException
   */
  public LzoW3CLogLoader(String fileURI) throws IOException {
    LOG.debug("Initialize LzoW3CLogLoader from " + fileURI);
    this.fileURI = fileURI;
  }

  /**
   * Return every non-null line as a single-element tuple to Pig.
   */
  @Override
  public Tuple getNext() throws IOException {
    LzoW3CLogRecordReader w3CLogRecordReader = (LzoW3CLogRecordReader) reader;
    if (w3CLogRecordReader == null) {
      return null;
    }
    MapWritable value_;
    try {
      if (w3CLogRecordReader.nextKeyValue()
          && (value_ = w3CLogRecordReader.getCurrentValue()) != null) {
        Map<String, String> values = Maps.newHashMap();

        for (Writable key : value_.keySet()) {
          Writable value = value_.get(key);
          values.put(key.toString(), value != null ? value.toString() : null);
        }
        incrCounter(LzoW3CLogLoaderCounters.LinesW3CDecoded, 1L);
        incrCounter(LzoW3CLogLoaderCounters.UnparseableLines,
            w3CLogRecordReader.getBadRecordsSkipped());
        return tupleFactory_.newTuple(values);
      }
    } catch (InterruptedException e) {
      int errCode = 6018;
      String errMsg = "Error while reading input";
      throw new ExecException(errMsg, errCode,
          PigException.REMOTE_ENVIRONMENT, e);
    }
    return null;
  }

  @Override
  public void setLocation(String location, Job job)
  throws IOException {
    FileInputFormat.setInputPaths(job, location);
  }

  @SuppressWarnings("rawtypes")
  @Override
  public InputFormat getInputFormat() {
    return LzoW3CLogInputFormat.newInstance(fileURI);
  }

  /**
   * NOT IMPLEMENTED
   */
  @Override
  public String[] getPartitionKeys(String arg0, Job arg1) throws IOException {
    return null;
  }

  @Override
  public ResourceSchema getSchema(String arg0, Job arg1) throws IOException {
    return new ResourceSchema(new Schema(new FieldSchema("data", DataType.MAP)));
  }

  /**
   * NOT IMPLEMENTED
   */
  @Override
  public ResourceStatistics getStatistics(String arg0, Job arg1) throws IOException {
    return null;
  }

  /**
   * NOT IMPLEMENTED
   */
  @Override
  public void setPartitionFilter(Expression arg0) throws IOException {

  }
}


package com.twitter.elephantbird.pig.util;

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

import com.hadoop.compression.lzo.LzopInputStream;
import org.apache.pig.impl.io.BufferedPositionedInputStream;

/**
 * A wrapper around Pig's BufferedPositionedInputStream that counts the number of
 * compressed bytes read rather than the number of decompressed bytes read.  Necessary
 * because the filesizes and splits are based off of compressed sizes, but the BufferedPositionedInputStream
 * only knows about uncompressed bytes read.  This screws up input splits.
 * Note that this class behaves identically to BufferedPositionedInputStream if the stream given in the
 * constructor is NOT an LzopInputStream.
 */
public class LzoBufferedPositionedInputStream extends BufferedPositionedInputStream {
  private LzopInputStream lzopIn_ = null;
  private long start_ = 0;

  /**
   * Store the input stream and start position if it's an LzopInputStream.
   * @param in the input stream
   * @param pos the starting position in the file.
   */
  public LzoBufferedPositionedInputStream(InputStream in, long pos) {
    super(in, pos);

    if (in instanceof LzopInputStream) {
      lzopIn_ = (LzopInputStream)in;
      start_ = pos;
    }
  }

  /**
   * If it's an LzopInputStream, return the compressed bytes read plus the offset.  Otherwise,
   * default to superclass behavior.
   */
  @Override
  public long getPosition() throws IOException {
    if (lzopIn_ != null) {
      return start_ + lzopIn_.getCompressedBytesRead();
    } else {
      return super.getPosition();
    }
  }

  /**
   * super.skip() skips uncompressed bytes, so we need to override it.
   * This allows us to specify skip size in a compressed stream.
   */
  @Override
  public long skip(long toSkip) throws IOException {
    long pos = getPosition();
    while (getPosition() - pos < toSkip) {
      if (super.skip(toSkip) == 0) break;
    }
    return getPosition() - pos;
  }

}


package com.twitter.elephantbird.mapreduce.output;

import java.io.IOException;
import java.util.List;

import com.google.common.collect.Lists;
import com.twitter.elephantbird.util.ColumnarMetadata;
import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hive.serde2.ByteStream;
import org.apache.hadoop.hive.serde2.columnar.BytesRefArrayWritable;
import org.apache.hadoop.hive.serde2.columnar.BytesRefWritable;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;

import com.google.protobuf.CodedOutputStream;
import com.google.protobuf.Message;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.Message.Builder;
import com.twitter.elephantbird.mapreduce.io.ProtobufWritable;
import com.twitter.elephantbird.util.Protobufs;
import com.twitter.elephantbird.util.TypeRef;

/**
 * OutputFormat for storing protobufs in RCFile.

 *
 * Each of the top level fields is stored in a separate column.
 * The protobuf field numbers are stored in RCFile metadata.

 *
 * A protobuf message can contain
 * "unknown fields". These fields are preserved and stored
 * in the last column. e.g. if protobuf A with 4 fields (a, b, c, d) is
 * serialized and when it is deserialized A has only 3 fields (a, c, d),
 * then 'b' is carried over as an unknown field.
 */
public class RCFileProtobufOutputFormat extends RCFileOutputFormat {

  /* typeRef is only required for setting metadata for the RCFile
   * if we delay file creation until the first row is written,
   * this info could be derived from protobuf being written.
   */
  private TypeRef<? extends Message> typeRef;
  private List<FieldDescriptor> msgFields;
  private int numColumns;

  private BytesRefArrayWritable rowWritable = new BytesRefArrayWritable();
  private BytesRefWritable[] colValRefs;
  private ByteStream.Output byteStream = new ByteStream.Output();
  private CodedOutputStream protoStream = CodedOutputStream.newInstance(byteStream);

  /** internal, for MR use only. */
  public RCFileProtobufOutputFormat() {
  }

  public RCFileProtobufOutputFormat(TypeRef<? extends Message> typeRef) { // for PigLoader etc.
    this.typeRef = typeRef;
    init();
  }

  private void init() {
    Builder msgBuilder = Protobufs.getMessageBuilder(typeRef.getRawClass());
    msgFields = msgBuilder.getDescriptorForType().getFields();
    numColumns = msgFields.size() + 1; // known fields + 1 for unknown fields
    colValRefs = new BytesRefWritable[numColumns];

    for (int i = 0; i < numColumns; i++) {
      colValRefs[i] = new BytesRefWritable();
      rowWritable.set(i, colValRefs[i]);
    }
  }

  protected ColumnarMetadata makeColumnarMetadata() {

    List<Integer> fieldIds = Lists.newArrayList();

    for(FieldDescriptor fd : msgFields) {
      fieldIds.add(fd.getNumber());
    }
    fieldIds.add(-1); // -1 for unknown fields

    return ColumnarMetadata.newInstance(typeRef.getRawClass().getName(), fieldIds);
  }

  private class ProtobufWriter extends RCFileOutputFormat.Writer {

    ProtobufWriter(TaskAttemptContext job) throws IOException {
      super(RCFileProtobufOutputFormat.this, job,
            Protobufs.toText(makeColumnarMetadata().getMessage()));
    }

    @Override
    public void write(NullWritable key, Writable value) throws IOException, InterruptedException {
      @SuppressWarnings("unchecked")
      Message msg = ((ProtobufWritable<Message>)value).get();

      protoStream.flush();
      byteStream.reset(); // reinitialize the byteStream if buffer is too large?
      int startPos = 0;

      // top level fields are split across the columns.
      for (int i=0; i < numColumns; i++) {

        if (i < (numColumns - 1)) {

          FieldDescriptor fd = msgFields.get(i);
          if (fd.isRepeated() || msg.hasField(fd)) {
            // match protobuf's serialization (write only if hasField() is true)
            Protobufs.writeFieldNoTag(protoStream, fd, msg.getField(fd));
          }

        } else { // last column : write unknown fields
          msg.getUnknownFields().writeTo(protoStream); // could be empty
        }

        protoStream.flush();
        colValRefs[i].set(byteStream.getData(),
                          startPos,
                          byteStream.getCount() - startPos);
        startPos = byteStream.getCount();
      }

      super.write(null, rowWritable);
    }
  }

  /**
   * Stores supplied class name in configuration. This configuration is
   * read on the remote tasks to initialize the output format correctly.
   */
  public static void setClassConf(Class<? extends Message> protoClass, Configuration conf) {
    Protobufs.setClassConf(conf, RCFileProtobufOutputFormat.class, protoClass);
  }

  @Override
  public RecordWriter<NullWritable, Writable>
    getRecordWriter(TaskAttemptContext job) throws IOException, InterruptedException {

    if (typeRef == null) {
      typeRef = Protobufs.getTypeRef(HadoopCompat.getConfiguration(job), RCFileProtobufOutputFormat.class);
      init();
    }

    RCFileOutputFormat.setColumnNumber(HadoopCompat.getConfiguration(job), numColumns);
    return new ProtobufWriter(job);
  }

}


package com.twitter.elephantbird.mapreduce.output;

import java.io.IOException;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.ql.io.RCFile;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.SequenceFile.Metadata;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.GzipCodec;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.util.ReflectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.twitter.elephantbird.util.RCFileUtil;

/**
 * Hive's {@link org.apache.hadoop.hive.ql.io.RCFileOutputFormat} is written for
 * deprecated OutputFormat. Pig requires newer OutputFormat.
 * In addition RCFileOutputFormat's functionality this class adds RCFile
 * metadata support.
 *
 * TODO: contribute this to PIG.
 */

public class RCFileOutputFormat extends FileOutputFormat<NullWritable, Writable> {

  private static final Logger LOG = LoggerFactory.getLogger(RCFileOutputFormat.class);

  // in case we need different compression from global default compression
  public static String COMPRESSION_CODEC_CONF = "elephantbird.rcfile.output.compression.codec";

  public static String DEFAULT_EXTENSION = ".rc";
  public static String EXTENSION_OVERRIDE_CONF = "elephantbird.refile.output.filename.extension"; // "none" disables it.

  /**
   * set number of columns into the given configuration.
   *
   * @param conf
   *          configuration instance which need to set the column number
   * @param columnNum
   *          column number for RCFile's Writer
   *
   */
  public static void setColumnNumber(Configuration conf, int columnNum) {
    assert columnNum > 0;
    conf.setInt(RCFile.COLUMN_NUMBER_CONF_STR, columnNum);
  }

  /**
   * Returns the number of columns set in the conf for writers.
   *
   * @param conf
   * @return number of columns for RCFile's writer
   */
  public static int getColumnNumber(Configuration conf) {
    return conf.getInt(RCFile.COLUMN_NUMBER_CONF_STR, 0);
  }

  protected RCFile.Writer createRCFileWriter(TaskAttemptContext job,
                                             Text columnMetadata)
                                             throws IOException {
    Configuration conf = HadoopCompat.getConfiguration(job);

    // override compression codec if set.
    String codecOverride = conf.get(COMPRESSION_CODEC_CONF);
    if (codecOverride != null) {
      conf.setBoolean("mapred.output.compress", true);
      conf.set("mapred.output.compression.codec", codecOverride);
    }

    CompressionCodec codec = null;
    if (getCompressOutput(job)) {
      Class<? extends CompressionCodec> codecClass = getOutputCompressorClass(job, GzipCodec.class);
      codec = (CompressionCodec) ReflectionUtils.newInstance(codecClass, conf);
    }

    Metadata metadata = null;
    if (columnMetadata != null) {
      metadata = new Metadata();
      metadata.set(new Text(RCFileUtil.COLUMN_METADATA_PROTOBUF_KEY), columnMetadata);
    }

    String ext = conf.get(EXTENSION_OVERRIDE_CONF, DEFAULT_EXTENSION);
    Path file = getDefaultWorkFile(job, ext.equalsIgnoreCase("none") ? null : ext);

    LOG.info("writing to rcfile " + file.toString());

    return new RCFile.Writer(file.getFileSystem(conf), conf, file, job, metadata, codec);
  }

  /**
   * RecordWriter wrapper around an RCFile.Writer
   */
  static protected class Writer extends RecordWriter<NullWritable, Writable> {

    private RCFile.Writer rcfile;

    protected Writer(RCFileOutputFormat outputFormat,
                     TaskAttemptContext job,
                     Text columnMetadata) throws IOException {
      rcfile = outputFormat.createRCFileWriter(job, columnMetadata);
    }

    @Override
    public void close(TaskAttemptContext context) throws IOException, InterruptedException {
      rcfile.close();
    }

    @Override
    public void write(NullWritable key, Writable value) throws IOException, InterruptedException {
      rcfile.append(value);
      // add counters
    }
  }

  @Override
  public RecordWriter<NullWritable, Writable> getRecordWriter(
      TaskAttemptContext job) throws IOException, InterruptedException {
    return new Writer(this, job, null);
  }
}


package com.twitter.elephantbird.pig.util;

import java.util.Map;

import com.google.common.collect.Maps;
import com.twitter.elephantbird.util.HadoopCompat;

import org.apache.hadoop.mapreduce.Counter;
import org.apache.pig.impl.util.Pair;
import org.apache.pig.tools.pigstats.PigStatusReporter;

/**
 * A helper class to deal with Hadoop counters in Pig.  They are stored within the singleton
 * PigStatusReporter instance, but are null for some period of time at job startup, even after
 * Pig has been invoked.  This class buffers counters, trying each time to get a valid Reporter and flushing
 * stored counters each time it does.
 */
public class PigCounterHelper {
  private final Map<Pair<String, String>, Long> counterStringMap_ = Maps.newHashMap();
  private final Map<Enum<?>, Long> counterEnumMap_ = Maps.newHashMap();

  /**
   * Mocks the Reporter.incrCounter, but adds buffering.
   * See org.apache.hadoop.mapred.Reporter's incrCounter.
   */
  public void incrCounter(String group, String counterName, long incr) {
    PigStatusReporter reporter = PigStatusReporter.getInstance();
    if (reporter != null) { // common case
      Counter counter = reporter.getCounter(group, counterName);
      if (counter != null) {
        HadoopCompat.incrementCounter(counter, incr);

        if (counterStringMap_.size() > 0) {
          for (Map.Entry<Pair<String, String>, Long> entry : counterStringMap_.entrySet()) {
            HadoopCompat.incrementCounter(
                reporter.getCounter(entry.getKey().first, entry.getKey().second),
                entry.getValue());
          }
          counterStringMap_.clear();
        }
        return;
      }
    }
    // In the case when reporter is not available, or we can't get the Counter,
    // store in the local map.
    Pair<String, String> key = new Pair<String, String>(group, counterName);
    Long currentValue = counterStringMap_.get(key);
    counterStringMap_.put(key, (currentValue == null ? 0 : currentValue) + incr);
  }

  /**
   * Mocks the Reporter.incrCounter, but adds buffering.
   * See org.apache.hadoop.mapred.Reporter's incrCounter.
   */
  public void incrCounter(Enum<?> key, long incr) {
    PigStatusReporter reporter = PigStatusReporter.getInstance();
    if (reporter != null && reporter.getCounter(key) != null) {
      HadoopCompat.incrementCounter(reporter.getCounter(key), incr);
      if (counterEnumMap_.size() > 0) {
        for (Map.Entry<Enum<?>, Long> entry : counterEnumMap_.entrySet()) {
          HadoopCompat.incrementCounter(
              reporter.getCounter(entry.getKey()), entry.getValue());
        }
        counterEnumMap_.clear();
      }
    } else { // buffer the increments
      Long currentValue = counterEnumMap_.get(key);
      counterEnumMap_.put(key, (currentValue == null ? 0 : currentValue) + incr);
    }
  }
}


package com.twitter.elephantbird.cascading2.scheme;

import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.mapred.OutputCollector;
import org.apache.hadoop.mapred.RecordReader;
import org.apache.hadoop.mapred.SequenceFileInputFormat;

import com.twitter.elephantbird.mapreduce.input.MapReduceInputFormatWrapper;
import com.twitter.elephantbird.mapreduce.input.combine.DelegateCombineFileInputFormat;

import cascading.flow.FlowProcess;
import cascading.scheme.hadoop.SequenceFile;
import cascading.tap.Tap;
import cascading.tuple.Fields;

/**
 * This scheme allows SequenceFile splits to be combined via the DelegateCombineFileInputFormat
 * before it is read. It can be used to combine inputs for intermediate MR jobs in Cascading.
 *
 * To enable, set cascading.flowconnector.intermediateschemeclass to this class in the Hadoop
 * configuration.
 *
 * @author Akihiro Matsukawa
 */
public class CombinedSequenceFile extends SequenceFile {

  private static final String MR_COMPRESS_ENABLE = "mapreduce.output.fileoutputformat.compress";
  public static final String COMPRESS_ENABLE = "elephantbird.cascading.combinedsequencefile.compress.enable";

  private static final String MR_COMPRESS_TYPE = "mapreduce.output.fileoutputformat.compress.type";
  public static final String COMPRESS_TYPE = "elephantbird.cascading.combinedsequencefile.compress.type";

  private static final String MR_COMPRESS_CODEC = "mapreduce.output.fileoutputformat.compress.codec";
  public static final String COMPRESS_CODEC = "elephantbird.cascading.combinedsequencefile.compress.codec";

  protected CombinedSequenceFile() { super(); }

  public CombinedSequenceFile(Fields fields) { super(fields); }

  // We can allow overriding the compression settings for just this scheme here
  private static void updateJobConfForLocalSettings(JobConf conf) {
    String localSetCompressionEnabled = conf.get(COMPRESS_ENABLE);
    if(localSetCompressionEnabled != null) {
      conf.set(MR_COMPRESS_ENABLE, localSetCompressionEnabled);
    }

    String localSetCompressionType = conf.get(COMPRESS_TYPE);
    if(localSetCompressionType != null) {
      conf.set(MR_COMPRESS_TYPE, localSetCompressionType);
    }

    String localSetCompressionCodec = conf.get(COMPRESS_CODEC);
    if(localSetCompressionCodec != null) {
      conf.set(MR_COMPRESS_CODEC, localSetCompressionCodec);
    }
  }

  public static void sourceConfInit(JobConf conf) {
    updateJobConfForLocalSettings(conf);

    // Since the EB combiner works over the mapreduce API while Cascading is on the mapred API,
    // in order to use the EB combiner we must wrap the mapred SequenceFileInputFormat
    // with the MapReduceInputFormatWrapper, then wrap the DelegateCombineFileInputFormat
    // again with DeprecatedInputFormatWrapper to make it compatible with Cascading.
    MapReduceInputFormatWrapper.setWrappedInputFormat(SequenceFileInputFormat.class, conf);
    DelegateCombineFileInputFormat.setDelegateInputFormat(conf, MapReduceInputFormatWrapper.class);
  }

  @Override
  public void sourceConfInit(
      FlowProcess<JobConf> flowProcess,
      Tap<JobConf, RecordReader, OutputCollector> tap,
      JobConf conf ) {
    super.sourceConfInit(flowProcess, tap, conf);

    sourceConfInit(conf);
  }

  @Override
  public void sinkConfInit( FlowProcess<JobConf> flowProcess, Tap<JobConf, RecordReader, OutputCollector> tap, JobConf conf )
  {
    super.sinkConfInit(flowProcess, tap, conf);

    updateJobConfForLocalSettings(conf);
  }

}


package com.twitter.elephantbird.cascading.protobuf;

import java.io.InputStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;

import com.google.protobuf.Message;

/**
 * Utility methods for reflection based protobuf deserialization
 * @author Ning Liang
 */
public final class ProtobufReflectionUtil {
  private ProtobufReflectionUtil() { }

  /**
   * Parse the method for a message
   * @param klass the class containing the message
   * @return the parsed method
   */
  public static Method parseMethodFor(Class<Message> klass) {
    try {
      return klass.getMethod("parseDelimitedFrom", new Class[] {InputStream.class });
    } catch (SecurityException e) {
      throw new RuntimeException(e);
    } catch (NoSuchMethodException e) {
      throw new RuntimeException(e);
    }
  }

  /**
   * Parse the message in a given InputStream using scpecified Method
   * @param parseMethod the method used for parsing
   * @param in the input stream
   * @return the parsed message
   */
  public static Message parseMessage(Method parseMethod, InputStream in) {
    try {
      return (Message) parseMethod.invoke(null, in);
    } catch (IllegalArgumentException e) {
      throw new RuntimeException(e);
    } catch (IllegalAccessException e) {
      throw new RuntimeException(e);
    } catch (InvocationTargetException e) {
      throw new RuntimeException(e);
    }
  }

  /**
   * Parse the message in a given Message container
   * @param klass the class containing the message
   * @param in the input stream
   * @return the parsed Message
   */
  public static Message parseMessage(Class<Message> klass, InputStream in) {
    Method parseMethod = parseMethodFor(klass);
    return parseMessage(parseMethod, in);
  }
}


package com.twitter.elephantbird.mapreduce.input;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.serde2.ColumnProjectionUtils;
import org.apache.hadoop.hive.serde2.columnar.BytesRefArrayWritable;
import org.apache.hadoop.hive.serde2.columnar.BytesRefWritable;
import org.apache.hadoop.io.LongWritable;
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.hadoop.mapreduce.lib.input.FileSplit;
import org.apache.thrift.TBase;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.transport.TMemoryInputTransport;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Function;
import com.google.common.collect.Lists;
import com.twitter.elephantbird.mapreduce.io.ThriftWritable;
import com.twitter.elephantbird.util.ColumnarMetadata;
import com.twitter.elephantbird.util.RCFileUtil;
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;

public class RCFileThriftInputFormat extends RCFileBaseInputFormat {

  private static final Logger LOG = LoggerFactory.getLogger(RCFileThriftInputFormat.class);

  private TypeRef<? extends TBase<?, ?>> typeRef;

  // for MR
  public RCFileThriftInputFormat() {}

  public RCFileThriftInputFormat(TypeRef<? extends TBase<?, ?>> typeRef) {
    this.typeRef = typeRef;
  }

  /**
   * Stores supplied class name in configuration. This configuration is
   * read on the remote tasks to initialize the input format correctly.
   */
  public static void setClassConf(Class<? extends TBase<?, ?> > thriftClass, Configuration conf) {
    ThriftUtils.setClassConf(conf, RCFileThriftInputFormat.class, thriftClass);
  }

  @Override
  public RecordReader<LongWritable, Writable>
  createRecordReader(InputSplit split, TaskAttemptContext taskAttempt)
                                    throws IOException, InterruptedException {
    if (typeRef == null) {
      typeRef = ThriftUtils.getTypeRef(HadoopCompat.getConfiguration(taskAttempt),
                                       RCFileThriftInputFormat.class);
    }
    return new ThriftReader(createUnwrappedRecordReader(split, taskAttempt));
  }

  public class ThriftReader extends FilterRecordReader<LongWritable, Writable> {

    protected TStructDescriptor     tDesc;
    protected boolean               readUnknownsColumn = false;
    protected List<Field>           knownRequiredFields = Lists.newArrayList();
    protected ArrayList<Integer>    columnsBeingRead = Lists.newArrayList();

    protected TMemoryInputTransport memTransport = new TMemoryInputTransport();
    protected TBinaryProtocol tProto = new TBinaryProtocol(memTransport);

    protected ThriftWritable<TBase<?, ?>> thriftWritable;

    /**
     * The reader is expected to be a
     * RecordReader< LongWritable, BytesRefArrayWritable >
     */
    @SuppressWarnings({ "unchecked", "rawtypes" })
    public ThriftReader(RecordReader reader) {
      super(reader);
    }

    /** is valid only after initialize() is called */
    public boolean isReadingUnknonwsColumn() {
      return readUnknownsColumn;
    }

    @Override @SuppressWarnings("unchecked")
    public void initialize(InputSplit split, TaskAttemptContext ctx)
                           throws IOException, InterruptedException {
      // set up columns that needs to read from the RCFile.

      tDesc = TStructDescriptor.getInstance(typeRef.getRawClass());
      thriftWritable = ThriftWritable.newInstance((Class<TBase<?, ?>>)typeRef.getRawClass());
      final List<Field> tFields = tDesc.getFields();

      FileSplit fsplit = (FileSplit)split;
      Path file = fsplit.getPath();

      LOG.info(String.format("reading %s from %s:%d:%d"
          , typeRef.getRawClass().getName()
          , file.toString()
          , fsplit.getStart()
          , fsplit.getStart() + fsplit.getLength()));

      Configuration conf = HadoopCompat.getConfiguration(ctx);
      ColumnarMetadata storedInfo = RCFileUtil.readMetadata(conf, file);

      // list of field numbers
      List<Integer> tFieldIds = Lists.transform(tFields,
                                    new Function<Field, Integer>() {
                                       public Integer apply(Field fd) {
                                         return Integer.valueOf(fd.getFieldId());
                                       }
                                    });

      columnsBeingRead = RCFileUtil.findColumnsToRead(conf, tFieldIds, storedInfo);

      for(int idx : columnsBeingRead) {
        int fid = storedInfo.getFieldId(idx);
        if (fid >= 0) {
          knownRequiredFields.add(tFields.get(tFieldIds.indexOf(fid)));
        } else {
          readUnknownsColumn = true;
        }
      }

      ColumnProjectionUtils.setReadColumnIDs(conf, columnsBeingRead);

      // finally!
      super.initialize(split, ctx);
    }

    @Override @SuppressWarnings({ "unchecked", "rawtypes" })
    public Writable getCurrentValue() throws IOException, InterruptedException {
      try {
        thriftWritable.set(getCurrentThriftValue());
        return thriftWritable;
      } catch (TException e) {
        //TODO : add error tracking
        throw new IOException(e);
      }
    }

    /** returns super.getCurrentValue() */
    public BytesRefArrayWritable getCurrentBytesRefArrayWritable() throws IOException, InterruptedException {
      return (BytesRefArrayWritable) super.getCurrentValue();
    }

    /**
     * Builds Thrift object from the raw bytes returned by RCFile reader.
     * @throws TException
     */
    @SuppressWarnings({ "unchecked", "rawtypes" })
    public TBase<?, ?> getCurrentThriftValue() throws IOException, InterruptedException, TException {

      BytesRefArrayWritable byteRefs = getCurrentBytesRefArrayWritable();

      if (byteRefs == null) {
        return null;
      }

      TBase tObj = tDesc.newThriftObject();

      for (int i=0; i < knownRequiredFields.size(); i++) {
        BytesRefWritable buf = byteRefs.get(columnsBeingRead.get(i));
        if (buf.getLength() > 0) {
          memTransport.reset(buf.getData(), buf.getStart(), buf.getLength());

          Field field = knownRequiredFields.get(i);
          tObj.setFieldValue(field.getFieldIdEnum(),
                             ThriftUtils.readFieldNoTag(tProto, field));
        }
        // else no need to set default value since any default value
        // would have been serialized when this record was written.
      }

      // parse unknowns column if required
      if (readUnknownsColumn) {
        int last = columnsBeingRead.get(columnsBeingRead.size() - 1);
        BytesRefWritable buf = byteRefs.get(last);
        if (buf.getLength() > 0) {
          memTransport.reset(buf.getData(), buf.getStart(), buf.getLength());
          tObj.read(tProto);
        }
      }

      return tObj;
    }

  }
}


package com.twitter.elephantbird.pig.store;

import java.io.IOException;
import java.util.Map;
import java.util.Set;

import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.OutputFormat;
import org.apache.pig.data.Tuple;
import org.json.simple.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Joiner;
import com.google.common.collect.Sets;
import com.twitter.elephantbird.mapreduce.output.LzoTextOutputFormat;

/**
 * A storage class to store the ouput of each tuple in a delimited file
 * like PigStorage, but LZO compressed.
 */
public class LzoJsonStorage extends BaseStoreFunc {
  private static final Logger LOG = LoggerFactory.getLogger(LzoJsonStorage.class);

  // If null, keep all keys.
  private final Set<String> keysToKeep_;
  private final JSONObject json = new JSONObject();

  public LzoJsonStorage() {
    keysToKeep_ = null;
    LOG.info("Initialized LzoJsonStorage. Keeping all keys.");
  }

  public LzoJsonStorage(String...keysToKeep) {
    keysToKeep_ = Sets.newHashSet(keysToKeep);
    LOG.info("Initialized LzoJsonStorage. Keeping keys " + Joiner.on(", ").join(keysToKeep_) + ".");
  }

  /**
   * The first element is expected to be a map, or null. Anything else causes an error.
   * @param tuple the tuple to write.
   */
  @Override
  @SuppressWarnings("unchecked")
  public void putNext(Tuple tuple) throws IOException {
    json.clear();
    if (tuple != null && tuple.size() >= 1) {
      Map<String, Object> map = (Map<String, Object>) tuple.get(0);
      if (keysToKeep_ == null) {
        json.putAll(map);
      } else {
        for (Map.Entry<String, Object> entry : map.entrySet()) {
          if (keysToKeep_.contains(entry.getKey())) {
            json.put(entry.getKey(), entry.getValue());
          }
        }
      }
    }
    try {
      writer.write(null, new Text(json.toString()));
    } catch (InterruptedException e) {
      // Under what circumstances does this happen?
      throw new IOException(e);
    }
  }

  @Override
  public OutputFormat<NullWritable, Text> getOutputFormat() {
    return new LzoTextOutputFormat();
  }

}


package com.twitter.elephantbird.mapreduce.output;

import java.io.IOException;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;

import com.google.protobuf.Message;
import com.twitter.elephantbird.mapreduce.io.ProtobufConverter;
import com.twitter.elephantbird.mapreduce.io.ProtobufWritable;
import com.twitter.elephantbird.util.Protobufs;
import com.twitter.elephantbird.util.TypeRef;

/**
 * This is the class for all base64 encoded, line-oriented protocol buffer based output formats.
 * Data is written as one base64 encoded serialized protocol buffer per line.

 *
 * Do not forget to set Protobuf class using setClassConf().
 */

public class LzoProtobufB64LineOutputFormat<M extends Message> extends LzoOutputFormat<M, ProtobufWritable<M>> {
  protected TypeRef<M> typeRef_;

  protected void setTypeRef(TypeRef<M> typeRef) {
    typeRef_ = typeRef;
  }

  public LzoProtobufB64LineOutputFormat() {}

  public LzoProtobufB64LineOutputFormat(TypeRef<M> typeRef) {
    typeRef_ = typeRef;
  }

  /**
   * Sets an internal configuration in jobConf so that remote Tasks
   * instantiate appropriate object for this generic class based on protoClass
   */
  public static <M extends Message>
  void setClassConf(Class<M> protoClass, Configuration jobConf) {
    Protobufs.setClassConf(jobConf,
                           LzoProtobufB64LineOutputFormat.class,
                           protoClass);
  }

  @Override
  public RecordWriter<M, ProtobufWritable<M>> getRecordWriter(TaskAttemptContext job)
  throws IOException, InterruptedException {
    if (typeRef_ == null) {
      typeRef_ = Protobufs.getTypeRef(HadoopCompat.getConfiguration(job), LzoProtobufB64LineOutputFormat.class);
    }
    return new LzoBinaryB64LineRecordWriter<M, ProtobufWritable<M>>(ProtobufConverter.newInstance(typeRef_), getOutputStream(job));
  }
}


package com.twitter.elephantbird.mapreduce.input;

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.Path;
import org.apache.hadoop.io.IOUtils;
import org.apache.hadoop.io.LongWritable;
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.Job;
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;

import com.twitter.elephantbird.mapreduce.io.BinaryBlockReader;
import com.twitter.elephantbird.mapreduce.io.BinaryWritable;
import com.twitter.elephantbird.mapreduce.io.IdentityBinaryConverter;
import com.twitter.elephantbird.mapreduce.io.BinaryConverter;
import com.twitter.elephantbird.mapreduce.io.RawBytesWritable;
import com.twitter.elephantbird.util.HadoopUtils;
import com.twitter.elephantbird.util.Protobufs;
import com.twitter.elephantbird.util.TypeRef;

/**
 * The input could consist of heterogeneous mix of formats storing
 * compatible objects. Currently supported formats:
 *

 *
 Lzo Block storage of Thrift and Protobuf objects
 *
 Lzo B64Line storage of Thrift and Protobuf objects
 *

 *
 *

 * A small fraction of bad records are tolerated. See {@link LzoRecordReader}
 * for more information on error handling.
 */
public class MultiInputFormat<M>
                extends LzoInputFormat<LongWritable, BinaryWritable<M>> {

  private static final Logger LOG = LoggerFactory.getLogger(MultiInputFormat.class);

  // TODO need handle multiple input formats in a job better.
  //      might be better to store classname in the input split rather than in config.
  private static String CLASS_CONF_KEY = "elephantbird.class.for.MultiInputFormat";
  private static String GENERIC_DECODER_KEY = "elephantbird.decoder.class.for.MultiInputFormat";

  private TypeRef<M> typeRef;

  public MultiInputFormat() {}

  public MultiInputFormat(TypeRef<M> typeRef) {
    this.typeRef = typeRef;
  }

  private static enum Format {
    LZO_BLOCK,
    LZO_B64LINE;
  };

  /**
   * Sets jobs input format to {@link MultiInputFormat} and stores
   * supplied clazz's name in job configuration. This configuration is
   * read on the remote tasks to initialize the input format correctly.
   */
  public static void setInputFormatClass(Class<?> clazz, Job job) {
    job.setInputFormatClass(MultiInputFormat.class);
    setClassConf(clazz, HadoopCompat.getConfiguration(job));
  }

  /**
   * Stores supplied class name in configuration. This configuration is
   * read on the remote tasks to initialize the input format correctly.
   */
  public static void setClassConf(Class<?> clazz, Configuration conf) {
    HadoopUtils.setClassConf(conf, CLASS_CONF_KEY, clazz);
  }

  public static void setGenericConverterClassConf(Class<?> clazz, Configuration conf) {
    HadoopUtils.setClassConf(conf, GENERIC_DECODER_KEY, clazz);
  }

  @SuppressWarnings({ "unchecked", "rawtypes" }) // return type is runtime dependent
  @Override
  public RecordReader<LongWritable, BinaryWritable<M>>
  createRecordReader(InputSplit split, TaskAttemptContext taskAttempt)
                     throws IOException, InterruptedException {
    Configuration conf = HadoopCompat.getConfiguration(taskAttempt);
    if (typeRef == null) {
      setTypeRef(conf);
    }
    Class<?> recordClass = typeRef.getRawClass();

    final Format fileFormat;
    try {
      LOG.info("Determining format of file in split: " + split);
      fileFormat = determineFileFormat(split, conf);
    } catch (IOException e) {
      throw new IOException("Could not determine format of file in split: " + split, e);
    } catch (RuntimeException e) {
      throw new RuntimeException("Could not determine format of file in split: " + split, e);
    }

    // Explicit class names for Message and TBase are used so that
    // these classes need not be present when not required.

    // TODO: merge all the switch (fileFormat) cases below
    // and set the right converter/writable for each

    if (isSubclass(recordClass, "com.google.protobuf.Message")) {
      switch (fileFormat) {
      case LZO_BLOCK:
        return new LzoProtobufBlockRecordReader(typeRef);
      case LZO_B64LINE:
        return new LzoProtobufB64LineRecordReader(typeRef);
      }
    }

    if (isSubclass(recordClass, "org.apache.thrift.TBase")) {
      switch (fileFormat) {
      case LZO_BLOCK:
        return new LzoThriftBlockRecordReader(typeRef);
      case LZO_B64LINE:
        return new LzoThriftB64LineRecordReader(typeRef);
      }
    }

    if (recordClass.equals(byte[].class)) {
      switch (fileFormat) {
      case LZO_BLOCK:
        return new LzoBinaryBlockRecordReader(typeRef,
            new BinaryBlockReader<byte[]>(null, new IdentityBinaryConverter()){},
            new RawBytesWritable()){};
      case LZO_B64LINE:
        return new LzoBinaryB64LineRecordReader(typeRef,
            new RawBytesWritable(), new IdentityBinaryConverter());
      }
    }

    String genericDecoder = conf.get(GENERIC_DECODER_KEY);
    if(genericDecoder != null) {
      Class<?> decoderClazz = null;
      BinaryConverterProvider<?> converterProvider = null;
      try {
        decoderClazz = conf.getClassByName(genericDecoder);
        converterProvider = (BinaryConverterProvider<?>)decoderClazz.newInstance();
      } catch (Exception e) {
        throw new RuntimeException("failed to instantiate class '" + genericDecoder + "'", e);
      }

      BinaryConverter<?> converter = converterProvider.getConverter(conf);
      switch (fileFormat) {
        case LZO_BLOCK:
          return new LzoGenericBlockRecordReader(typeRef, converter);
        case LZO_B64LINE:
          return new LzoGenericB64LineRecordReader(typeRef, converter);
      }
    }

    throw new IOException( "could not determine reader for "
        + ((FileSplit)split).getPath() + " with class " + recordClass.getName());
  }

  /** set typeRef from conf */
  private void setTypeRef(Configuration conf) {
    String className = conf.get(CLASS_CONF_KEY);

    if (className == null) {
      throw new RuntimeException(CLASS_CONF_KEY + " is not set");
    }

    Class<?> clazz = null;
    try {
      clazz = conf.getClassByName(className);
    } catch (ClassNotFoundException e) {
      throw new RuntimeException("failed to instantiate class '" + className + "'", e);
    }

    typeRef = new TypeRef<M>(clazz){};
  }

  private static boolean isSubclass(Class<?> subClass, String superClassName) {
    try {
      Class<?> superClass = Class.forName(superClassName, true, subClass.getClassLoader());
      return superClass.isAssignableFrom(subClass);
    } catch (Exception e) {
      // the superClassName may not be present in classpath when not required.
      return false;
    }
  }

  /**
   * Checks to see if the input records are stored as SerializedBlock.
   * The block format starts with {@link Protobufs#KNOWN_GOOD_POSITION_MARKER}.
   * Otherwise the input is assumed to be Base64 encoded lines.
   */
  private static Format determineFileFormat(InputSplit split,
                                            Configuration conf)
                                            throws IOException {
    FileSplit fileSplit = (FileSplit) split;

    Path file = fileSplit.getPath();

    /* we could have a an optional configuration that maps a regex on a
     * file name to a format. E.g. ".*-block.lzo" to LZO_BLOCK file.
     */

    // most of the cost is opening the file and
    // reading first lzo block (about 256k of uncompressed data)

    CompressionCodec codec = new CompressionCodecFactory(conf).getCodec(file);
    if (codec == null) {
      throw new IOException("No codec for file " + file + " found");
    }

    InputStream in = file.getFileSystem(conf).open(file);
    InputStream lzoIn = null;

    // check if the file starts with magic bytes for Block storage format.
    try {
      lzoIn = codec.createInputStream(in);

      for(byte magic : Protobufs.KNOWN_GOOD_POSITION_MARKER) {
        int b = lzoIn.read();
        if (b < 0 || (byte)b != magic) {
          return Format.LZO_B64LINE;
        }
      }
    } finally {
      IOUtils.closeStream(lzoIn);
      IOUtils.closeStream(in);
    }

    // the check passed
    return Format.LZO_BLOCK;
  }
}


package com.twitter.elephantbird.pig.util;

import java.io.IOException;
import java.util.Map;

import com.google.common.base.Preconditions;

import org.apache.hadoop.io.Writable;
import org.apache.pig.ResourceSchema.ResourceFieldSchema;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.DataByteArray;
import org.apache.pig.data.DataType;
import org.apache.pig.data.Tuple;

/**
 * Base class providing default implementations for WritableConverter methods.
 *
 * @author Andy Schlaikjer
 */
public abstract class AbstractWritableConverter<W extends Writable> extends WritableStoreCaster<W>
    implements WritableConverter<W> {
  /**
   * Constructs new instance of this class using the given Writable instance for data conversion.
   * The class of the given Writable is used to initialize {@link #writableClass}.
   *
   * @param writable
   */
  public AbstractWritableConverter(W writable) {
    super(writable);
  }

  /**
   * Default constructor.
   */
  public AbstractWritableConverter() {
    super();
  }

  /**
   * Default implementation tests that class of internal Writable instance, if defined, is
   * assignable from the given class. If so, the internal Writable instance is replaced with a new
   * instance of the given class. If the given class is {@code null}, nothing is done.
   *
   * @see WritableConverter#initialize(java.lang.Class)
   */
  @Override
  public void initialize(Class<? extends W> writableClass) throws IOException {
    if (writableClass == null) {
      return;
    }
    if (writable != null) {
      Class<?> existingWritableClass = writable.getClass();
      Preconditions.checkArgument(existingWritableClass.isAssignableFrom(writableClass),
          "Existing Writable implementation '%s' is not assignable from class '%s'",
          existingWritableClass.getName(), writableClass.getName());
    }
    try {
      writable = writableClass.newInstance();
    } catch (Exception e) {
      throw new IOException(String.format("Failed to create instance of Writable type '%s'",
          writableClass.getName()), e);
    }
  }

  /**
   * Default implementation returns {@code null}.
   *
   * @see com.twitter.elephantbird.pig.util.WritableConverter#getLoadSchema()
   */
  @Override
  public ResourceFieldSchema getLoadSchema() throws IOException {
    return null;
  }

  /**
   * Default implementation passes raw bytes through to Pig without deserializing Writable data.
   *
   * @see WritableConverter#bytesToObject(org.apache.pig.data.DataByteArray)
   */
  @Override
  public Object bytesToObject(DataByteArray dataByteArray) throws IOException {
    return dataByteArray;
  }

  /**
   * Default implementation returns value of {@link #writableClass}.
   *
   * @see com.twitter.elephantbird.pig.util.WritableConverter#getWritableClass()
   */
  @Override
  @SuppressWarnings("unchecked")
  public Class<W> getWritableClass() throws IOException {
    if (writable == null) {
      return null;
    }
    return (Class<W>) writable.getClass();
  }

  /**
   * Default implementation does nothing.
   *
   * @see WritableConverter#checkStoreSchema(org.apache.pig.ResourceSchema.ResourceFieldSchema)
   */
  @Override
  public void checkStoreSchema(ResourceFieldSchema schema) throws IOException {
  }

  /**
   * Default implementation routes to one of the other {@code toWritable} methods based on the type
   * of the input Pig value.
   *
   * @see WritableConverter#toWritable(java.lang.Object, org.apache.hadoop.io.Writable)
   */
  @Override
  @SuppressWarnings("unchecked")
  public W toWritable(Object value) throws IOException {
    if (value == null) {
      return null;
    }
    // route to appropriate method using Pig data type
    switch (DataType.findType(value)) {
      case DataType.BYTEARRAY:
        return toWritable((DataByteArray) value);
      case DataType.CHARARRAY:
        return toWritable((String) value);
      case DataType.INTEGER:
        return toWritable((Integer) value);
      case DataType.LONG:
        return toWritable((Long) value);
      case DataType.FLOAT:
        return toWritable((Float) value);
      case DataType.DOUBLE:
        return toWritable((Double) value);
      case DataType.MAP:
        return toWritable((Map<String, Object>) value);
      case DataType.TUPLE:
        return toWritable((Tuple) value);
      case DataType.BAG:
        return toWritable((DataBag) value);
      default:
        throw new IOException("Pig value class '" + value.getClass().getName() + "' is unsupported");
    }
  }
}


package com.twitter.elephantbird.thrift;

import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocolException;
import org.apache.thrift.transport.TTransport;

/**
 * A shim on top of thrift to allow for thrift 0.7/0.9 compatibility.
 *
 * This one is designed for thrift 0.7
 *
 */
abstract class AbstractThriftBinaryProtocol extends TBinaryProtocol {
  public AbstractThriftBinaryProtocol(TTransport trans) {
    super(trans);
  }

  public AbstractThriftBinaryProtocol(TTransport trans, boolean strictRead, boolean strictWrite) {
    super(trans, strictRead, strictWrite);
  }

  @SuppressWarnings("unused")
  public AbstractThriftBinaryProtocol(TTransport trans, long stringLengthLimit, long containerLengthLimit) {
    super(trans);
  }

  protected void resetAndInitialize(TBinaryProtocol protocol, int newLength) {
    protocol.reset();
    protocol.setReadLength(newLength);
  }

  /**
   * Check if the container size is valid.
   *
   * NOTE: This assumes that the elements are one byte each. So this does not
   * catch all cases, but does increase the chances of handling malformed
   * lengths when the number of remaining bytes in the underlying Transport is
   * clearly less than the container size that the Transport provides.
   */
  protected void checkContainerSize(int size) throws TProtocolException {
    if (size < 0) {
      throw new TProtocolException("Negative container size: " + size);
    }

    if (checkReadLength_ && (readLength_ - size) < 0) {
      throw new TProtocolException("Remaining message length is " + readLength_
              + " but container size in underlying TTransport is set to at least: " + size);
    }
  }
}


package com.twitter.elephantbird.util;

import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;

/**
 * An efficient stream searching class based on the Knuth-Morris-Pratt algorithm.
 * For more on the algorithm works see: http://www.inf.fh-flensburg.de/lang/algorithmen/pattern/kmpen.htm.
 */
public class StreamSearcher {

  protected byte[] pattern_;
  protected int[] borders_;

  // An upper bound on pattern length for searching. Results are undefined for longer patterns.
  public static final int MAX_PATTERN_LENGTH = 1024;

  public StreamSearcher(byte[] pattern) {
    setPattern(pattern);
  }

  /**
   * Sets a new pattern for this StreamSearcher to use.
   * @param pattern
   *          the pattern the StreamSearcher will look for in future calls to search(...)
   */
  public void setPattern(byte[] pattern) {
    pattern_ = Arrays.copyOf(pattern, pattern.length);
    borders_ = new int[pattern_.length + 1];
    preProcess();
  }

  /**
   * Searches for the next occurrence of the pattern in the stream, starting from the current stream position. Note
   * that the position of the stream is changed. If a match is found, the stream points to the end of the match -- i.e. the
   * byte AFTER the pattern. Else, the stream is entirely consumed. The latter is because InputStream semantics make it difficult to have
   * another reasonable default, i.e. leave the stream unchanged.
   *
   * @return bytes consumed if found, -1 otherwise.
   * @throws IOException
   */
  public long search(InputStream stream) throws IOException {
    long bytesRead = 0;

    int b;
    int j = 0;

    while ((b = stream.read()) != -1) {
      bytesRead++;

      while (j >= 0 && (byte)b != pattern_[j]) {
        j = borders_[j];
      }
      // Move to the next character in the pattern.
      ++j;

      // If we've matched up to the full pattern length, we found it.  Return,
      // which will automatically save our position in the InputStream at the point immediately
      // following the pattern match.
      if (j == pattern_.length) {
        return bytesRead;
      }
    }

    // No dice, Note that the stream is now completely consumed.
    return -1;
  }

  /**
   * Builds up a table of longest "borders" for each prefix of the pattern to find. This table is stored internally
   * and aids in implementation of the Knuth-Moore-Pratt string search.
   * <p>
   * For more information, see: http://www.inf.fh-flensburg.de/lang/algorithmen/pattern/kmpen.htm.
   */
  protected void preProcess() {
    int i = 0;
    int j = -1;
    borders_[i] = j;
    while (i < pattern_.length) {
      while (j >= 0 && pattern_[i] != pattern_[j]) {
        j = borders_[j];
      }
      borders_[++i] = ++j;
    }
  }
}


package com.twitter.elephantbird.util;

import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Wrappers around the singularization, pluaralization, etc in Inflection.
 *
 * Largely taken from the defunct rogueweb project at
 * <link>http://code.google.com/p/rogueweb/</link>
 * Original author Anthony Eden.
 */
public class Strings {

  public static String tableize(String word) {
    return pluralize(underscore(depackage(word)));
  }

  /**
   * Pluralize the given word.
   * @param word The word
   * @return The pluralized version of the word
   */
  public static String pluralize(String word) {
    return Inflection.pluralize(word);
  }

  public static String singularize(String word) {
    return Inflection.singularize(word);
  }

  /**
   * Underscore the given word.
   * @param word The word
   * @return The underscored version of the word
   */
  public static String underscore(String word) {
    String firstPattern = "([A-Z]+)([A-Z][a-z])";
    String secondPattern = "([a-z\\d])([A-Z])";
    String replacementPattern = "$1_$2";
    // Replace package separator with slash.
    word = word.replaceAll("\\.", "/");
    // Replace $ with two underscores for inner classes.
    word = word.replaceAll("\\$", "__");
    // Replace capital letter with _ plus lowercase letter.
    word = word.replaceAll(firstPattern, replacementPattern);
    word = word.replaceAll(secondPattern, replacementPattern);
    word = word.replace('-', '_');
    word = word.toLowerCase();
    return word;
  }

  public static String camelize(String word) {
    return camelize(word, false);
  }

  public static String camelize(String word, boolean lowercaseFirstLetter) {
    // Replace all slashes with dots (package separator)
    Pattern p = Pattern.compile("\\/(.?)");
    Matcher m = p.matcher(word);
    while (m.find()) {
      word = m.replaceFirst("." + m.group(1)/*.toUpperCase()*/);
      m = p.matcher(word);
    }

    // Uppercase the class name.
    p = Pattern.compile("(\\.?)(\\w)([^\\.]*)$");
    m = p.matcher(word);
    if (m.find()) {
      String rep = m.group(1) + m.group(2).toUpperCase() + m.group(3);
      rep = rep.replaceAll("\\$", "\\\\\\$");
      word = m.replaceAll(rep);
    }

    // Replace two underscores with $ to support inner classes.
    p = Pattern.compile("(__)(.)");
    m = p.matcher(word);
    while (m.find()) {
      word = m.replaceFirst("\\$" + m.group(2).toUpperCase());
      m = p.matcher(word);
    }

    // Remove all underscores
    p = Pattern.compile("(_)(.)");
    m = p.matcher(word);
    while (m.find()) {
      word = m.replaceFirst(m.group(2).toUpperCase());
      m = p.matcher(word);
    }

    if (lowercaseFirstLetter) {
      word = word.substring(0, 1).toLowerCase() + word.substring(1);
    }

    return word;
  }

  public static String depackage(String word) {
    return word.replaceAll("^.*\\.", "");
  }

  public static String foreignKey(String className) {
    return foreignKey(className, true);
  }

  public static String foreignKey(String className, boolean separateIdWithUnderscore) {
    return underscore(depackage(className)) + (separateIdWithUnderscore ? "_id" : "id");
  }

  public static String titleize(String word) {
    word = humanize(underscore(word));
    Pattern p = Pattern.compile("\\b([a-z])");
    Matcher m = p.matcher(word);
    while (m.find()) {
      word = m.replaceFirst(capitalize(m.group(1)));
      m = p.matcher(word);
    }
    return word;
  }

  /**
   * Converts a table name into a class name.
   * @param tableName The table name
   * @return The class name
   */
  public static String classify(String tableName) {
    return camelize(singularize(tableName.replaceAll(".*\\.", "")));
  }

  public static String humanize(String word) {
    word = word.replaceAll("_id", "");
    return capitalize(word.replaceAll("_", " "));
  }

  /**
   * Return the ordinal for the given number.
   * @param n The number
   * @return The ordinal String
   */
  public static String ordinalize(int n) {
    int mod100 = n % 100;
    if (mod100 == 11 || mod100 == 12 || mod100 == 13) {
      return String.valueOf(n) + "th";
    }
    switch (n % 10) {
      case 1:
        return String.valueOf(n) + "st";
      case 2:
        return String.valueOf(n) + "nd";
      case 3:
        return String.valueOf(n) + "rd";
      default:
        return String.valueOf(n) + "th";
    }
  }

  /**
   * Replace underscores with dashes.
   * @param word The word
   * @return The word with underscores converted to dashes
   */
  public static String dasherize(String word) {
    return word.replaceAll("_", "-");
  }

  public static String capitalize(String word) {
    word = word.toLowerCase();
    word = word.substring(0, 1).toUpperCase() + word.substring(1);
    return word;
  }
}


package com.twitter.elephantbird.mapreduce.io;

import com.google.protobuf.Message;
import com.twitter.elephantbird.util.TypeRef;

/**
 * A Hadoop Writable wrapper around a protocol buffer of type M.
 */

public class ProtobufWritable<M extends Message> extends BinaryWritable<M> {

  public ProtobufWritable() {
    super(null, null);
  }

  public ProtobufWritable(TypeRef<M> typeRef) {
    this(null, typeRef);
  }

  public ProtobufWritable(M message, TypeRef<M> typeRef) {
    super(message, new ProtobufConverter<M>(typeRef));
  }

  /**
   * Returns a ProtobufWritable for a given Protobuf class.
   */
  public static <M extends Message> ProtobufWritable<M> newInstance(Class<M> tClass) {
    return new ProtobufWritable<M>(new TypeRef<M>(tClass){});
  }

  @Override
  protected BinaryConverter<M> getConverterFor(Class<M> clazz) {
    return ProtobufConverter.newInstance(clazz);
  }
}


package com.twitter.elephantbird.util;

public class Utils {
  /**
   * returns Class.forName(className, true, classLoader).

   * Throws a RuntimeExcepiton if the class is not found.
   *
   * @see {@link Class#forName(String, boolean, ClassLoader)
   */
  public static Class<?> classForName(String className, ClassLoader classLoader) {
    try {
      return Class.forName(className, true, classLoader);
    } catch (ClassNotFoundException e) {
      throw new RuntimeException("failed to load class " + className, e);
    }
  }

  /**
   * Ensures the classLoader is 'consistent' with the original
   * class loader that created existingClass. Asserts

   * classLoader.loadClass(existingClass.getName()) == existingClass.
   *

   *
   * If classLoader fails to load the class, this returns silently.

   * Throws a RuntimeException with detailed message if the consistency
   * check fails.
   *
   * @param existingClass
   * @param classLoader
   */
  public static void ensureClassLoaderConsistency(Class<?> existingClass,
                                                  ClassLoader classLoader) {
    Class<?> loadedClass;
    try {
      loadedClass = Class.forName(existingClass.getName(), true, classLoader);
    } catch (ClassNotFoundException e) {
      return; // let class loading fail some where else.
    }

    if (!loadedClass.equals(existingClass)) {
      throw new RuntimeException("The class loader is incosistent with the "
              + "class loader that initially loaded "
              + existingClass.getClass()
              + ". This can lead to various unexpected side effects.");

    }
  }
}