package com.twitter.elephantbird.mapreduce.io;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.util.Arrays;

import org.apache.hadoop.io.BytesWritable;
import org.apache.hadoop.io.WritableComparable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A Hadoop Writable wrapper around a serialized messages like Protocol buffers.
 */
public abstract class BinaryWritable<M> implements WritableComparable<BinaryWritable<M>> {
  private static final Logger LOG = LoggerFactory.getLogger(BinaryWritable.class);

  // NOTE: only one of message and messageBytes is non-null at any time so that
  // message and messageBytes don't go out of sync (user could modify message).
  private M message;
  private byte[] messageBytes;
  private BinaryConverter<M> converter;

  protected BinaryWritable(M message, BinaryConverter<M> converter) {
    this.message = message;
    this.converter = converter;
  }

  /** throws an exception if the converter is not set */
  private void checkConverter() {
    if (converter == null) {
      throw new IllegalStateException("Runtime parameterized Protobuf/Thrift class is unkonwn. " +
                                      "This object was probably created with default constructor. " +
                                      "Please use setConverter(Class).");
    }
  }

  protected abstract BinaryConverter<M> getConverterFor(Class<M> clazz);

  /**
   * Sets the handler for serialization and deserialization based on the class.
   * This converter is often set in constructor. But some times it might be
   * impossible to know the the actual class during construction.

   *
   * E.g. when this writable is used as output value for a Mapper,
   * MR creates writable on the Reducer using the default constructor,
   * and there is no way for us to know the parameterized class.
   * In this case, user invokes setConverter() before
   * calling get() to supply parameterized class.

   *
   * The class name could be written as part of writable serialization, but we
   * don't yet see a need to do that as it has many other disadvantages.
   */
  public void setConverter(Class<M> clazz) {
    converter = getConverterFor(clazz);
  }

  /**
   * sets converter. useful for reusing existing converter.
   */
  public void setConverter(BinaryConverter<M> converter) {
    this.converter = converter;
  }

  /**
   * Returns the current object. Subsequent calls to get() may not return the
   * same object, but in stead might return a new object deserialized from same
   * set of bytes. As a result, multiple calls to get() should be avoided, and
   * modifications to an object returned by get() may not
   * reflect even if this writable is serialized later.

   * Please use set() to be certain of what object is serialized.

   *
   * The deserialization of the actual Protobuf/Thrift object is often delayed
   * till the first call to this method.

   * In some cases the the parameterized proto class may not be known yet
   * ( in case of default construction. see {@link #setConverter(Class)} ),
   * and this will throw an {@link IllegalStateException}.
   */
  public M get() {
    // may be we should rename this method. the contract would be less
    // confusing with a different name.
    if (message == null && messageBytes != null) {
      checkConverter();
      try {
        return converter.fromBytes(messageBytes);
      } catch (DecodeException e) {
        throw new IllegalStateException("Converter failed to deserialize", e);
      }
    }
    return message;
  }

  public void clear() {
    message = null;
    messageBytes = null;
  }

  public void set(M message) {
    this.message = message;
    this.messageBytes = null;
    // should we serialize the object to messageBytes instead?
    // that is the only way we can guarantee any subsequent modifications to
    // message by the user don't affect serialization. Unlike Protobuf objects
    // Thrift objects are mutable. For now we will delay deserialization until
    // it is required.
  }

  @Override
  public void write(DataOutput out) throws IOException {
    byte[] bytes = serialize();

    if (bytes != null) {
      out.writeInt(bytes.length);
      out.write(bytes, 0, bytes.length);
    } else {
      out.writeInt(0);
    }
  }

  /**
   * Converts the message to raw bytes, and caches the converted value.
   * @return converted value, which may be null in case of null message or error.
   */
  private byte[] serialize() {
    if (messageBytes == null && message != null) {
      checkConverter();
      messageBytes = converter.toBytes(message);
      if (messageBytes == null) {
        // should we throw an IOException instead?
        LOG.warn("Could not serialize " + message.getClass());
      } else {
        message = null; // so that message and messageBytes don't go out of
                        // sync.
      }
    }
    return messageBytes;
  }

  @Override
  public void readFields(DataInput in) throws IOException {
    message = null;
    messageBytes = null;
    int size = in.readInt();
    if (size > 0) {
      byte[] buf = new byte[size];
      in.readFully(buf, 0, size);
      messageBytes = buf;
      // messageBytes is deserialized in get()
    }
  }

  @Override
  public int compareTo(BinaryWritable<M> other) {
    byte[] thisBytes = serialize();
    byte[] otherBytes = other.serialize();
    int thisLen = thisBytes == null ? 0 : thisBytes.length;
    int otherLen = otherBytes == null ? 0 : otherBytes.length;
    return BytesWritable.Comparator.compareBytes(thisBytes, 0, thisLen,
        otherBytes, 0, otherLen);
  }

  @SuppressWarnings("unchecked")
  @Override
  public boolean equals(Object obj) {
    if (obj == null) {
      return false;
    }

    BinaryWritable<M> other;
    try {
      other = (BinaryWritable<M>)obj;
    } catch (ClassCastException e) {
      return false;
    }

    return compareTo(other) == 0;
  }

  /**
   *
Returns a hashCode that is based on the serialized bytes.
   * This makes the hash stable across multiple instances of JVMs.
   * (hashCode() is not required to return the same value in
   * different instances of the same applications in Java, just in a
   * single instance of the application; Hadoop imposes a more strict requirement.)
   *

   * In addition, it may not be feasible to create a deserialized object from
   * the serialized bytes (see {@link #setConverter(Class)})
   */
  @Override
  public int hashCode() {
    byte[] bytes = serialize();
    return (bytes == null) ? 31 : Arrays.hashCode(bytes);
  }

  @Override
  public String toString() {
    M msgObj = null;
    try {
      msgObj = get();
    } catch (IllegalStateException e) {
      // It is ok. might not be able to avoid this case in some situations.
      return super.toString() + "{could not be deserialized}";
    }
    if (msgObj == null) {
      return super.toString();
    }
    return msgObj.toString();
  }
}


package com.twitter.elephantbird.mapred.input;

import java.io.*;
import java.util.List;

import com.twitter.elephantbird.util.HadoopCompat;
import com.twitter.elephantbird.util.SplitUtil;
import org.apache.hadoop.conf.Configurable;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.serializer.Deserializer;
import org.apache.hadoop.io.serializer.SerializationFactory;
import org.apache.hadoop.io.serializer.Serializer;
import org.apache.hadoop.mapred.Counters;
import org.apache.hadoop.mapred.FileSplit;
import org.apache.hadoop.mapred.InputSplit;
import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.mapred.RecordReader;
import org.apache.hadoop.mapred.Reporter;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.StatusReporter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.TaskAttemptID;
import org.apache.hadoop.util.ReflectionUtils;

import com.twitter.elephantbird.mapreduce.input.MapredInputFormatCompatible;
import com.twitter.elephantbird.mapred.output.DeprecatedOutputFormatWrapper;
import com.twitter.elephantbird.util.HadoopUtils;

/**
 * The wrapper enables an {@link InputFormat} written for new
 * mapreduce interface to be used unmodified in contexts where
 * a {@link org.apache.hadoop.mapred.InputFormat} with old mapred
 * interface is required.

 *
 * Current restrictions on InputFormat:

 *
 the record reader should reuse key and value objects
 *    or implement {@link com.twitter.elephantbird.mapred.input.MapredInputFormatCompatible}

 *

 *
 * While this restriction is satisfied by most input formats,
 * it could be removed with a configuration option.
 *

 *
 * Usage:

 *    // set InputFormat class using a mapreduce InputFormat
 *    DeprecatedInputFormatWrapper.setInputFormat(org.apache.hadoop.mapreduce.lib.input.TextInputFormat.class, jobConf);
 *    jobConf.setOutputFormat(org.apache.hadoop.mapred.TextOutputFormat.class);
 *    // ...
 *

 *
 * @see DeprecatedOutputFormatWrapper
 *
 * @author Raghu Angadi
 *
 */
@SuppressWarnings("deprecation")
public class DeprecatedInputFormatWrapper<K, V> implements org.apache.hadoop.mapred.InputFormat<K, V> {

  public static final String CLASS_CONF_KEY = "elephantbird.class.for.DeprecatedInputFormatWrapper";
	public static final String VALUE_COPY_CONF_KEY = "elephantbird.class.for.ValueCopyClass";

  protected InputFormat<K, V> realInputFormat;
	protected DeprecatedInputFormatValueCopier<V> valueCopier = null;

  /**
   * Sets jobs input format to {@link DeprecatedInputFormatWrapper} and stores
   * supplied real {@link InputFormat} class name in job configuration.
   * This configuration is read on the remote tasks to instantiate actual
   * InputFormat correctly.
   */
  public static void setInputFormat(Class<?> realInputFormatClass, JobConf jobConf) {
    jobConf.setInputFormat(DeprecatedInputFormatWrapper.class);
    HadoopUtils.setClassConf(jobConf, CLASS_CONF_KEY, realInputFormatClass);
  }

  /**
   * For cases where we need to set hadoop1 input format in a hadoop2 Configuration object.
   */
  public static void setInputFormat(Class<?> realInputFormatClass, Configuration conf) {
    conf.setClass("mapred.input.format.class", DeprecatedInputFormatWrapper.class, org.apache.hadoop.mapred.InputFormat.class);
    HadoopUtils.setClassConf(conf, CLASS_CONF_KEY, realInputFormatClass);
  }

	public static void setInputFormat(Class<?> realInputFormatClass, JobConf jobConf,
			Class<? extends DeprecatedInputFormatValueCopier<?>> valueCopyClass) {
		jobConf.setInputFormat(DeprecatedInputFormatWrapper.class);
		HadoopUtils.setClassConf(jobConf, CLASS_CONF_KEY, realInputFormatClass);
		HadoopUtils.setClassConf(jobConf, VALUE_COPY_CONF_KEY, valueCopyClass);
	}

	/**
	 * Only used in very specific cases and does not absolve one from
	 * calling the static setInptuFormat methods to set up the hadoop job
	 * properly.
	 * @param inputFormat
	 */
	public void setInputFormatInstance(InputFormat<K, V> inputFormat) {
	  realInputFormat = inputFormat;
	}

  @SuppressWarnings("unchecked")
  private void initInputFormat(JobConf conf) {
    if (realInputFormat == null) {
			realInputFormat = ReflectionUtils.newInstance(conf.getClass(CLASS_CONF_KEY, null, InputFormat.class), conf);
			if (conf.get(VALUE_COPY_CONF_KEY) != null) {
				Class<? extends DeprecatedInputFormatValueCopier> copierClass = conf.getClass(VALUE_COPY_CONF_KEY,
						null, DeprecatedInputFormatValueCopier.class);
				if (null != copierClass) {
					valueCopier = ReflectionUtils.newInstance(copierClass, conf);
    }

			}
		}
  }

  public DeprecatedInputFormatWrapper() {
    // real inputFormat is initialized based on conf.
  }

  public DeprecatedInputFormatWrapper(InputFormat<K, V> realInputFormat) {
    this.realInputFormat = realInputFormat;
  }

  @Override
  public RecordReader<K, V> getRecordReader(InputSplit split, JobConf job,
                  Reporter reporter) throws IOException {
    initInputFormat(job);
		return new RecordReaderWrapper<K, V>(realInputFormat, split, job, reporter, valueCopier);
  }

  @Override
  public InputSplit[] getSplits(JobConf job, int numSplits) throws IOException {
    initInputFormat(job);

    try {
      List<org.apache.hadoop.mapreduce.InputSplit> splits =
        realInputFormat.getSplits(HadoopCompat.newJobContext(job, null));

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

      InputSplit[] resultSplits = new InputSplit[splits.size()];
      int i = 0;
      for (org.apache.hadoop.mapreduce.InputSplit split : splits) {
        if (split.getClass() == org.apache.hadoop.mapreduce.lib.input.FileSplit.class) {
          org.apache.hadoop.mapreduce.lib.input.FileSplit mapreduceFileSplit =
              ((org.apache.hadoop.mapreduce.lib.input.FileSplit)split);
          resultSplits[i++] = new FileSplit(
              mapreduceFileSplit.getPath(),
              mapreduceFileSplit.getStart(),
              mapreduceFileSplit.getLength(),
              mapreduceFileSplit.getLocations());
        } else {
          InputSplitWrapper wrapper = new InputSplitWrapper(split);
          wrapper.setConf(job);
          resultSplits[i++] = wrapper;
        }
      }

      return resultSplits;

    } catch (InterruptedException e) {
      throw new IOException(e);
    }
  }

  /**
   * A reporter that works with both mapred and mapreduce APIs.
   */
  public static class ReporterWrapper extends StatusReporter implements Reporter {
    private Reporter wrappedReporter;

    public ReporterWrapper(Reporter reporter) {
      wrappedReporter = reporter;
    }

    @Override
    public Counters.Counter getCounter(Enum<?> anEnum) {
      return wrappedReporter.getCounter(anEnum);
    }

    @Override
    public Counters.Counter getCounter(String s, String s1) {
      return wrappedReporter.getCounter(s, s1);
    }

    @Override
    public void incrCounter(Enum<?> anEnum, long l) {
      wrappedReporter.incrCounter(anEnum, l);
    }

    @Override
    public void incrCounter(String s, String s1, long l) {
      wrappedReporter.incrCounter(s, s1, l);
    }

    @Override
    public InputSplit getInputSplit() throws UnsupportedOperationException {
      return wrappedReporter.getInputSplit();
    }

    @Override
    public void progress() {
      wrappedReporter.progress();
    }

    // @Override
    public float getProgress() {
      throw new UnsupportedOperationException();
    }

    @Override
    public void setStatus(String s) {
      wrappedReporter.setStatus(s);
    }
  }

  private static class RecordReaderWrapper<K, V> implements RecordReader<K, V> {

    private org.apache.hadoop.mapreduce.RecordReader<K, V> realReader;
    private MapredInputFormatCompatible mifcReader = null;

    private long splitLen; // for getPos()

    // expect readReader return same Key & Value objects (common case)
    // this avoids extra serialization & deserialazion of these objects
    private K keyObj = null;
    private V valueObj = null;

    private boolean firstRecord = false;
    private boolean eof = false;

    private DeprecatedInputFormatValueCopier<V> valueCopier = null;

    public RecordReaderWrapper(InputFormat<K, V> newInputFormat, InputSplit oldSplit, JobConf oldJobConf,
				Reporter reporter, DeprecatedInputFormatValueCopier<V> valueCopier) throws IOException {

			this.valueCopier = valueCopier;
      splitLen = oldSplit.getLength();

      org.apache.hadoop.mapreduce.InputSplit split;
      if (oldSplit.getClass() == FileSplit.class) {
        split = new org.apache.hadoop.mapreduce.lib.input.FileSplit(
            ((FileSplit)oldSplit).getPath(),
            ((FileSplit)oldSplit).getStart(),
            ((FileSplit)oldSplit).getLength(),
            oldSplit.getLocations());
      } else {
        split = ((InputSplitWrapper)oldSplit).realSplit;
      }

      TaskAttemptID taskAttemptID = TaskAttemptID.forName(oldJobConf.get("mapred.task.id"));
      if (taskAttemptID == null) {
        taskAttemptID = new TaskAttemptID();
      }

      // create a MapContext to pass reporter to record reader (for counters)
      TaskAttemptContext taskContext = HadoopCompat
          .newMapContext(oldJobConf, taskAttemptID, null, null, null,
              new ReporterWrapper(reporter), null);
      try {
        realReader = newInputFormat.createRecordReader(split, taskContext);
        realReader.initialize(split, taskContext);

        if (realReader instanceof MapredInputFormatCompatible) {
          mifcReader = ((MapredInputFormatCompatible) realReader);
        }
      } catch (InterruptedException e) {
        throw new IOException(e);
      }
    }

    private void initKeyValueObjects() {
      // read once to gain access to key and value objects
      try {
        if (!firstRecord & !eof) {
          if (realReader.nextKeyValue()) {
            firstRecord = true;
            keyObj = realReader.getCurrentKey();
            valueObj = realReader.getCurrentValue();
          } else {
            eof = true;
          }
        }
      } catch (Exception e) {
        throw new RuntimeException("Could not read first record (and it was not an EOF)", e);
      }
    }

    @Override
    public void close() throws IOException {
      realReader.close();
    }

    @Override
    public K createKey() {
      initKeyValueObjects();
      return keyObj;
    }

    @Override
    public V createValue() {
      initKeyValueObjects();
      return valueObj;
    }

    @Override
    public long getPos() throws IOException {
      return (long) (splitLen * getProgress());
    }

    @Override
    public float getProgress() throws IOException {
      try {
        return realReader.getProgress();
      } catch (InterruptedException e) {
        throw new IOException(e);
      }
    }

    @Override
    public boolean next(K key, V value) throws IOException {
      if (eof) {
        return false;
      }

      if (firstRecord) { // key & value are already read.
        firstRecord = false;
        return true;
      }

      if (mifcReader != null) {
        mifcReader.setKeyValue(key, value);
      }

      try {
        if (realReader.nextKeyValue()) {

					if (key != realReader.getCurrentKey()) {

            if (mifcReader != null) {
              throw new IOException("The RecordReader returned a key and value that do not match "
                  + "the key and value sent to it. This means the RecordReader did not properly implement "
                  + "com.twitter.elephantbird.mapred.input.MapredInputFormatCompatible. "
                  + "Current reader class : " + realReader.getClass());

            } else {
              throw new IOException("DeprecatedInputFormatWrapper only "
                  + "supports RecordReaders that return the same key & value "
                  + "objects or implement com.twitter.elephantbird.mapred.input.MapredInputFormatCompatible. "
                  + "Current reader class : " + realReader.getClass());
            }

					}

					if (value != realReader.getCurrentValue()) {
						if (null != valueCopier)
							valueCopier.copyValue(value, realReader.getCurrentValue());
						else {
							throw new IOException("DeprecatedInputFormatWrapper - value is different "
									+ "and no value copier provided. "
									+ "Current reader class : " + realReader.getClass());
          }
					}
          return true;
        }
      } catch (InterruptedException e) {
        throw new IOException(e);
      }

      eof = true; // strictly not required, just for consistency
      return false;
    }
  }

  private static class InputSplitWrapper implements InputSplit, Configurable {

    org.apache.hadoop.mapreduce.InputSplit realSplit;
    private Configuration conf;

    @SuppressWarnings("unused") // MapReduce instantiates this.
    public InputSplitWrapper() {}

    public InputSplitWrapper(org.apache.hadoop.mapreduce.InputSplit realSplit) {
      this.realSplit = realSplit;
    }

    @Override
    public long getLength() throws IOException {
      try {
        return realSplit.getLength();
      } catch (InterruptedException e) {
        throw new IOException(e);
      }
    }

    @Override
    public String[] getLocations() throws IOException {
      try {
        return realSplit.getLocations();
      } catch (InterruptedException e) {
        throw new IOException(e);
      }
    }

    @Override
    public void readFields(DataInput in) throws IOException {
      realSplit = SplitUtil.deserializeInputSplit(conf, (DataInputStream) in);
    }

    @Override
    public void write(DataOutput out) throws IOException {
      SplitUtil.serializeInputSplit(conf, (DataOutputStream) out, realSplit);
    }

    @Override
    public void setConf(Configuration conf) {
      this.conf = conf;
    }

    @Override
    public Configuration getConf() {
      return conf;
    }
  }
}


package com.twitter.elephantbird.pig.load;

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

import com.google.common.base.Charsets;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.google.common.io.Files;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.pig.LoadFunc;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.PigSplit;
import org.apache.pig.data.Tuple;

import com.twitter.elephantbird.mapreduce.input.LuceneIndexInputFormat;
import com.twitter.elephantbird.mapreduce.input.LuceneIndexRecordReader;
import com.twitter.elephantbird.util.HdfsUtils;

/**
 * Base class for LoadFuncs that load data from lucene indexes.
 *

*  Wraps {@link LuceneIndexInputFormat}
 *

 * Subclasses must provide a {@link LuceneIndexInputFormat} and specify how to convert an MR
 * record to a tuple
 *

 * Constructor has two formats, one for loading queries from a file and one for supplying them
 * directly as pig string literals.
 * For example:
 *
 * x = load '/some/path' using MyLuceneIndexLoader('--queries', 'a query', 'another query');
 *
 * or
 *
 * x = load '/some/path' using MyLuceneIndexLoader('--file', 'path/to/local/file');
 *
 *
 * The file should have one query per line and be UTF-8 encoded
 * In both cases, the strings provided (as literals or in a file) are the serialized form
 * of the query used by {@link LuceneIndexInputFormat#setQueries(List, Configuration)}
 *
 * @param <T> type of records that will be converted to tuples
 * @author Alex Levenson
 */
public abstract class LuceneIndexLoader<T extends Writable> extends LoadFunc {

  private static final String USAGE_HELP = "LuceneIndexLoader's constructor usage:\n"
    + "LuceneIndexLoader('--queries', 'a query', 'another query')\nor\n"
    + "LuceneIndexLoader('--file', 'path/to/local/file')";

  private LuceneIndexRecordReader<T> reader;
  protected List<String> queries = null;
  protected File queryFile = null;

  /**
   * Convert a value from the InputFormat to a tuple
   *
   * @param key the id of the the query this record belongs to
   * @param value the value from the input format
   * @return a Tuple representing this key value pair (can have whatever schema you'd like)
   */
  protected abstract Tuple recordToTuple(int key, T value);

  /**
   * Provide an instance of the {@link LuceneIndexInputFormat} to be wrapped by this loader
   * @return an instance of the {@link LuceneIndexInputFormat} to be wrapped by this loader
   * @throws IOException
   */
  protected abstract LuceneIndexInputFormat<T> getLuceneIndexInputFormat() throws IOException;

  public LuceneIndexLoader(String[] args) {
    Preconditions.checkNotNull(args, USAGE_HELP);
    Preconditions.checkArgument(args.length >= 2, USAGE_HELP);
    Preconditions.checkNotNull(args[0], USAGE_HELP);

    if (args[0].equals("--queries")) {
      queries = Lists.newArrayList(Arrays.copyOfRange(args, 1, args.length));
    } else if (args[0].equals("--file")) {
      Preconditions.checkArgument(args.length == 2, USAGE_HELP);
      queryFile = new File(args[1]);
    } else {
      throw new IllegalArgumentException(USAGE_HELP);
    }
  }

  /**
   * Subclasses may use this constructor, but they will have to set queries or queryFile
   * themselves. Failure to do so will cause an IllegalArgumentException in setLocation
   */
  protected LuceneIndexLoader() { }

  /**
   * THIS INVOLVES AN UNCHECKED CAST
   * Pig gives us a raw type for the RecordReader unfortunately.
   * However, because {@link #getInputFormat()} is final and delegates
   * to {@link #getLuceneIndexInputFormat()} we can be reasonably sure
   * that this record reader is actually a LuceneIndexRecordReader<T>
   */
  @Override
  @SuppressWarnings("unchecked")
  public void prepareToRead(RecordReader reader, PigSplit split) {
    this.reader = (LuceneIndexRecordReader<T>) reader;
  }

  @Override
  public void setLocation(String location, Job job) throws IOException {
    Preconditions.checkArgument(queries != null || queryFile != null,
        "Either queires or queryFile must be set in the constructor!");

    Configuration conf = HadoopCompat.getConfiguration(job);
    // prevent pig from trying to combine splits, let LuceneIndexInputFormat do that
    conf.setBoolean("pig.noSplitCombination", true);

    // lazy load the queries into conf
    if (!LuceneIndexInputFormat.queriesSet(conf)) {
      if (queries != null) {
        LuceneIndexInputFormat.setQueries(queries, conf);
      } else {
        LuceneIndexInputFormat.setQueries(loadQueriesFromFile(conf), conf);
      }
    }

    // a little bit of logic to support comma separated locations and also
    // hdfs glob syntax locations
    String[] locationsWithGlobs = getPathStrings(location);
    List<Path> expandedPaths = HdfsUtils.expandGlobs(Arrays.asList(locationsWithGlobs), conf);
    LuceneIndexInputFormat.setInputPaths(expandedPaths, conf);
  }

  protected List<String> loadQueriesFromFile(Configuration conf) throws IOException {
    Preconditions.checkArgument(queryFile.exists(),
      "Query file: " + queryFile + " does not exist!");
    List<String> lines = Files.readLines(queryFile, Charsets.UTF_8);
    List<String> strippedLines = Lists.newArrayListWithCapacity(lines.size());
    for (String line : lines) {
      line = line.trim();
      if (!line.isEmpty()) {
        strippedLines.add(line);
      }
    }
    return strippedLines;
  }

  @Override
  public Tuple getNext() throws IOException {
    try {
      if (!reader.nextKeyValue()) {
        // all done
        return null;
      }

      int key;
      T value;

      try {
        key = reader.getCurrentKey().get();
        value = reader.getCurrentValue();
      } catch (ClassCastException e) {
        throw new IOException("Record reader did not return correct key or value type", e);
      }

      return recordToTuple(key, value);
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      throw new IOException(e);
    }
  }

  // enforce some more type safety than pig gives us by default
  @Override
  public final InputFormat getInputFormat() throws IOException {
    return getLuceneIndexInputFormat();
  }

}


package com.twitter.elephantbird.util;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.Iterator;
import java.util.List;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

import com.google.common.base.Charsets;
import com.google.common.base.Preconditions;
import com.google.common.base.Splitter;
import com.google.common.collect.Lists;
import com.google.common.io.Closeables;

import org.apache.commons.codec.binary.Base64;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.Counter;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.TaskInputOutputContext;
import org.json.simple.JSONArray;
import org.json.simple.JSONValue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Various Hadoop specific utilities.
 */
public class HadoopUtils {
  private static final Logger LOG = LoggerFactory.getLogger(HadoopUtils.class);
  private static final Splitter COMMA_SPLITTER = Splitter.on(',');

  private HadoopUtils() { }

  /**
   * MapReduce counters are available only with {@link TaskInputOutputContext},
   * but most interfaces use super classes, though the actual object is a
   * subclass (e.g. Mapper.Context).
   * <br>
   * This utility method checks the type and returns the appropriate counter.
   * In the rare (may be unexpected) case where ctx is not a
   * TaskInputOutputContext, a dummy counter is returned after printing
   * a warning.
   */
  public static Counter getCounter(JobContext ctx, String group, String counter) {
    if (ctx instanceof TaskInputOutputContext<?, ?, ?, ?>) {
      Counter c = HadoopCompat.getCounter((TaskInputOutputContext<?, ?, ?, ?>) ctx,
          group, counter);
      if (c != null) {
        return c;
      }
    }
    String name = group + ":" + counter;
    LOG.warn("Using a dummy counter for " + name + " because it does not already exist.");
    return HadoopCompat.newGenericCounter(name, name, 0);
  }

  /**
   * @deprecated use {@link #setClassConf(Configuration, String, Class)}
   */
  @Deprecated
  public static void setInputFormatClass(Configuration  conf,
                                         String         configKey,
                                         Class<?>       clazz) {
    setClassConf(conf, configKey, clazz);
  }

  /**
   * A helper to set configuration to class name.
   * Throws a RuntimeExcpetion if the
   * configuration is already set to a different class name.
   */
  public static void setClassConf(Configuration  conf,
                                  String         configKey,
                                  Class<?>       clazz) {
    String existingClass = conf.get(configKey);
    String className = clazz.getName();

    if (existingClass != null && !existingClass.equals(className)) {
        throw new RuntimeException(
            "Already registered a different thriftClass for "
            + configKey
            + ". old: " + existingClass
            + " new: " + className);
    } else {
      conf.set(configKey, className);
    }
  }

  /**
   * Writes an object into a configuration by converting it to a base64 encoded string.
   * This is done by serializing the object to bytes via an {@link ObjectOutputStream},
   * gzip compressing those bytes, and then base64 encoding the compressed bytes.
   *
   * NOTE: obj must implement {@link java.io.Serializable}
   *
   * @param key for the configuration
   * @param obj to write (must be Serializable)
   * @param conf to write to
   * @throws IOException
   */
  public static void writeObjectToConfAsBase64(String key, Object obj, Configuration conf) throws IOException {
    ByteArrayOutputStream baos = null;
    GZIPOutputStream gos = null;
    ObjectOutputStream oos = null;

    try {
      baos = new ByteArrayOutputStream();
      gos = new GZIPOutputStream(baos);
      oos = new ObjectOutputStream(gos);
      oos.writeObject(obj);
    } finally {
      Closeables.close(oos, false);
      Closeables.close(gos, false);
      Closeables.close(baos, false);
    }

    conf.set(key, new String(Base64.encodeBase64(baos.toByteArray()), Charsets.UTF_8));
  }

  /**
   * Reads an object (that was written using
   * {@link #writeObjectToConfAsBase64}) from a configuration.
   *
   * @param key for the configuration
   * @param conf to read from
   * @return the read object, or null if key is not present in conf
   * @throws IOException
   */
  @SuppressWarnings("unchecked")
  public static <T> T readObjectFromConfAsBase64(String key, Configuration conf) throws IOException {
    String b64 = conf.get(key);
    if (b64 == null) {
      return null;
    }

    byte[] bytes = Base64.decodeBase64(b64.getBytes(Charsets.UTF_8));

    ByteArrayInputStream bais = null;
    GZIPInputStream gis = null;
    ObjectInputStream ois = null;

    try {
      bais = new ByteArrayInputStream(bytes);
      gis = new GZIPInputStream(bais);
      ois = new ObjectInputStream(gis);
      return (T) ois.readObject();
    } catch (ClassNotFoundException e) {
      LOG.error("Could not read object from config with key " + key, e);
      throw new IOException(e);
    } catch (ClassCastException e) {
      LOG.error("Couldn't cast object read from config with key " + key, e);
      throw new IOException(e);
    } finally {
      Closeables.close(ois, false);
      Closeables.close(gis, false);
      Closeables.close(bais, false);
    }
  }

  /**
   * Writes a list of strings into a configuration by converting it to a json array
   *
   * @param key for the configuration
   * @param list to write
   * @param conf to write to
   */
  public static void writeStringListToConfAsJson(String key,
                                                 List<String> list,
                                                 Configuration conf) {
    Preconditions.checkNotNull(list);
    conf.set(key, JSONArray.toJSONString(list));
  }

  /**
   * Reads a list of strings stored as a json array from  a configuration
   *
   * @param key for the configuration
   * @param conf to read from
   * @return the read list of strings, or null if key is not present in conf
   */
  @SuppressWarnings("unchecked")
  public static List<String> readStringListFromConfAsJson(String key, Configuration conf) {
    String json = conf.get(key);
    if (json == null) {
      return null;
    }
    return Lists.<String>newArrayList(((JSONArray) JSONValue.parse(json)));
  }

  /**
   * Writes a list of strings into a configuration by base64 encoding them and separating
   * them with commas
   *
   * @param key for the configuration
   * @param list to write
   * @param conf to write to
   */
  public static void writeStringListToConfAsBase64(String key, List<String> list, Configuration
    conf) {
    Preconditions.checkNotNull(list);
    Iterator<String> iter = list.iterator();
    StringBuilder sb = new StringBuilder();
    while(iter.hasNext()) {
      byte[] bytes = Base64.encodeBase64(iter.next().getBytes(Charsets.UTF_8), false);
      sb.append(new String(bytes, Charsets.UTF_8));
      if (iter.hasNext()) {
        sb.append(',');
      }
    }
    conf.set(key, sb.toString());
  }

  /**
   * Reads a list of strings stored as comma separated base64
   *
   * @param key for the configuration
   * @param conf to read from
   * @return the read list of strings, or null if key is not present in conf
   */
  @SuppressWarnings("unchecked")
  public static List<String> readStringListFromConfAsBase64(String key, Configuration conf) {
    String b64List = conf.get(key);
    if (b64List == null) {
      return null;
    }

    List<String> strings = Lists.newArrayList();

    for (String b64 : COMMA_SPLITTER.split(b64List)) {
      byte[] bytes = Base64.decodeBase64(b64.getBytes(Charsets.UTF_8));
      strings.add(new String(bytes, Charsets.UTF_8));
    }

    return strings;
  }
}


package com.twitter.elephantbird.util;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import com.google.common.collect.Maps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/*
 * A parser for W3C-style log lines.  See LzoW3CLogInputFormat for more details about
 * the format itself.  Create the parser with an InputStream open to a list of
 * hash/field-list lines that give all the field-sets ever used for logging.  That is,
 * if you logged a set of fields for two days, then added a parameter and logged that set of
 * fields for a day, then added another parameter and removed a different one and logged
 * that set of fields for a while, the field definition file should have three lines
 * with the hashes and corresponding set of fields for each.
 */
public class W3CLogParser {
  protected static final Logger LOG = LoggerFactory.getLogger(W3CLogParser.class);

  private Map<String, List<String>> fieldDef_ = Maps.newHashMap();
  private final static String DELIMITER = "\\s+";

  // Initialize a W3CLogParser From a InputStream
  public W3CLogParser(InputStream is) throws IOException {
    BufferedReader buf = new BufferedReader(new InputStreamReader(is));
    String line;
    while ((line = buf.readLine()) != null) {
      List<String> fields = Arrays.asList(line.split(DELIMITER));
      fieldDef_.put(fields.get(0), fields.subList(1, fields.size()));
    }
  }

  // Parse a line using the field definition, and put results into a Map of
  // <FieldName, Value>
  public Map<String, String> parse(String line) throws IOException {
    Map<String, String> map = new HashMap<String, String>();

    // Get the version CRC and find the field definition
    List<String> fields = Arrays.asList(line.split(DELIMITER));
    List<String> fieldNames = fieldDef_.get(fields.get(0));
    if (fieldNames == null) {
      throw new IOException("cannot find matching field definition for CRC "
          + fields.get(0));
    }
    if (fieldNames.size() != fields.size()) {
      throw new IOException("W3C field definition and input line for CRC "
          + fields.get(0) + " does not match:\n" + line);
    }

    // Map values to field names
    for (int fieldNum = 1; fieldNum < fieldNames.size(); fieldNum++) {
      map.put(fieldNames.get(fieldNum), fields.get(fieldNum));
    }

    return map;
  }
}


package com.twitter.elephantbird.util;

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

import org.apache.hadoop.conf.Configuration;
import org.apache.thrift.TBase;
import org.apache.thrift.TEnum;
import org.apache.thrift.TException;
import org.apache.thrift.meta_data.EnumMetaData;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.protocol.TType;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;

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

public class ThriftUtils {

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

  public static void setClassConf(Configuration jobConf, Class<?> genericClass,
                                  Class<? extends TBase<?, ?>> thriftClass) {
    String configKey = CLASS_CONF_PREFIX + genericClass.getName();
    HadoopUtils.setClassConf(jobConf, configKey, thriftClass);
  }

  /**
   * Verify that clazz is a Thrift class. i.e. is a subclass of TBase
   */
  private static void verifyAncestry(Class<?> tClass) {
    if (!TBase.class.isAssignableFrom(tClass)) {
      Utils.ensureClassLoaderConsistency(TBase.class, tClass.getClassLoader());
      throw new ClassCastException(tClass.getName() + " is not a Thrift class");
    }
  }

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

    Class<?> tClass = null;
    try {
      tClass = jobConf.getClassByName(className);
    } catch (ClassNotFoundException e) {
      throw new RuntimeException(e);
    }

    verifyAncestry(tClass);

    return new TypeRef<M>(tClass){};
  }

   /**
    * returns TypeRef for thrift class. Verifies that the class is
    * a Thrift class (i.e extends TBase).
    */
  public static <M extends TBase<?, ?>> TypeRef<M> getTypeRef(Class<?> tClass) {
    verifyAncestry(tClass);
    return new TypeRef<M>(tClass){};
  }

  /**
   * returns TypeRef for a thrift class.
   */
  public static <M extends TBase<?, ?>> TypeRef<M> getTypeRef(String thriftClassName, ClassLoader classLoader) {
    try {
      Class<?> tClass = classLoader == null ?
          Class.forName(thriftClassName) :
          Class.forName(thriftClassName, true, classLoader);

      verifyAncestry(tClass);

      return new TypeRef<M>(tClass){};
    } catch (ClassNotFoundException e) {
      throw new RuntimeException(e);
    }
  }

  /**
   * returns TypeRef for a thrift class.
   */
  public static<M extends TBase<?, ?>> TypeRef<M> getTypeRef(String thriftClassName) {
    return getTypeRef(thriftClassName, null);
  }

   /**
   * Returns value of a fieldName in an object.
   */
  public static <M> M getFieldValue(Object containingObject, String fieldName, Class<M> fieldClass) {
    return getFieldValue(containingObject.getClass(), containingObject, fieldName, fieldClass);
  }

  /**
   * Returns value of a static field with given name in containingClass.
   */
  public static <M> M getFieldValue(Class<?> containingClass, String fieldName, Class<M> fieldClass) {
    return getFieldValue(containingClass, null, fieldName, fieldClass);
  }

  private static <M> M getFieldValue(Class<?> containingClass, Object obj, String fieldName, Class<M> fieldClass) {
    try {
      java.lang.reflect.Field field = containingClass.getDeclaredField(fieldName);
      return fieldClass.cast(field.get(obj));
    } catch (Exception e) {
      throw new RuntimeException("while trying to find " + fieldName + " in "
                                  +  containingClass.getName(), e);
    }
  }

  /**
   * Returns generic type for a field in a Thrift class. The type is the return
   * type for the accessor method for the field (e.g. isFieldName()
   * for a boolean type or getFieldName for other types). The return
   * type works for both structs and unions. Reflecting directly based on
   * fields does not work for unions.
   *
   * @return generic {@link Type} of the thrift field.
   */
  public static Type getFieldType(Class<?> containingClass, String fieldName) {

    String suffix = // uppercase first letter
        fieldName.substring(0, 1).toUpperCase() + fieldName.substring(1);

    // look for getFieldName() or isFieldName()

    for(String prefix : new String[]{"get", "is"}) {
      try {
        Method method = containingClass.getDeclaredMethod(prefix + suffix);
        return method.getGenericReturnType();
      } catch (NoSuchMethodException e) {
      }
    }

    // look for bean style accessors get_fieldName and is_fieldName

    for(String prefix : new String[]{"get_", "is_"}) {
      try {
        Method method = containingClass.getDeclaredMethod(prefix + fieldName);
        return method.getGenericReturnType();
      } catch (NoSuchMethodException e) {
      }
    }

    throw new RuntimeException("could not find type for " + fieldName +
                                 " in " + containingClass);
  }

  /**
   * Returns the value class of the given thrift field.
   *
   * @param field thrift field.
   * @return class of field value, or null in the case of thrift types {@link TType#STOP} or
   * {@link TType#VOID}.
   */
  public static Class<?> getFieldValueType(Field field) {
    switch (field.getType()) {
      case TType.BOOL:
        return Boolean.class;
      case TType.BYTE:
        return Byte.class;
      case TType.DOUBLE:
        return Double.class;
      case TType.ENUM:
        return field.getEnumClass();
      case TType.I16:
        return Short.class;
      case TType.I32:
        return Integer.class;
      case TType.I64:
        return Long.class;
      case TType.LIST:
        return List.class;
      case TType.MAP:
        return Map.class;
      case TType.SET:
        return Set.class;
      case TType.STOP:
        return null;
      case TType.STRING:
        return String.class;
      case TType.STRUCT:
        return field.gettStructDescriptor().getThriftClass();
      case TType.VOID:
        return null;
    }
    return null;
  }

  private static void writeSingleFieldNoTag(TProtocol proto,
                                            Field field,
                                            Object value) throws TException {
    switch(field.getType()) {

    case TType.BOOL:
      proto.writeBool((Boolean)value);              break;
    case TType.BYTE:
      proto.writeByte((Byte)value);                 break;
    case TType.I16:
      proto.writeI16((Short)value);                 break;
    case TType.I32:
      proto.writeI32((Integer)value);               break;
    case TType.ENUM:
      proto.writeI32(((TEnum)value).getValue());    break;
    case TType.I64:
      proto.writeI64((Long)value);                  break;
    case TType.DOUBLE:
      proto.writeDouble((Double)value);             break;
    case TType.STRING: {
      if (value instanceof String) {
        proto.writeString((String)value);
      } else {
        proto.writeBinary((ByteBuffer)value);
      }
    }                                               break;
    case TType.STRUCT:
      ((TBase<?, ?>)value).write(proto);            break;

    default:
      throw new IllegalArgumentException("Unexpected type : " + field.getType());
    }
  }

  /**
   * Serializes a single field of a thrift struct.
   *
   * @throws TException
   */
  public static void writeFieldNoTag(TProtocol proto,
                                     Field field,
                                     Object value) throws TException {
    if (value == null) {
      return;
    }

    Field innerField = null;

    switch (field.getType()) {

    case TType.LIST:
      innerField = field.getListElemField();    break;
    case TType.SET:
      innerField = field.getSetElemField();     break;
    case TType.MAP:
      innerField = field.getMapKeyField();      break;

    default:
      writeSingleFieldNoTag(proto, field, value);
      return;
    }

    // a map or a collection:

    if (field.getType() == TType.MAP) {

      Field valueField = field.getMapValueField();
      Map<?, ?> map = (Map<?, ?>)value;

      proto.writeByte(innerField.getType());
      proto.writeByte(valueField.getType());
      proto.writeI32(map.size());

      for(Entry<?, ?> entry : map.entrySet()) {
        writeSingleFieldNoTag(proto, innerField, entry.getKey());
        writeSingleFieldNoTag(proto, valueField, entry.getValue());
      }

    } else { // SET or LIST

      Collection<?> coll = (Collection<?>)value;

      proto.writeByte(innerField.getType());
      proto.writeI32(coll.size());

      for(Object v : coll) {
        writeSingleFieldNoTag(proto, innerField, v);
      }

    }
  }

  private static Object readSingleFieldNoTag(TProtocol  proto,
                                             Field      field)
                                             throws TException {
    switch(field.getType()) {

    case TType.BOOL:
      return proto.readBool();
    case TType.BYTE:
      return proto.readByte();
    case TType.I16:
      return proto.readI16();
    case TType.I32:
      return proto.readI32();
    case TType.ENUM:
      return field.getEnumValueOf(proto.readI32());
    case TType.I64:
      return proto.readI64();
    case TType.DOUBLE:
      return proto.readDouble();
    case TType.STRING:
      return field.isBuffer() ?  proto.readBinary() : proto.readString();
    case TType.STRUCT:
      TBase<?, ?> tObj = field.gettStructDescriptor().newThriftObject();
      tObj.read(proto);
      return tObj;

    default:
      throw new IllegalArgumentException("Unexpected type : " + field.getType());
    }

  }

  /**
   * Deserializes a thrift field that was serilized with
   * {@link #writeFieldNoTag(TProtocol, Field, Object)}.
   *
   * @throws TException in case of any Thrift errors.
   */
  public static Object readFieldNoTag(TProtocol   proto,
                                      Field       field)
                                      throws TException {

    Collection<Object> coll = null;
    Field innerField = null;

    switch (field.getType()) {

    case TType.LIST:
      innerField = field.getListElemField();
      coll = Lists.newArrayList();              break;
    case TType.SET:
      innerField = field.getSetElemField();
      coll = Sets.newHashSet();                 break;
    case TType.MAP:
      innerField = field.getMapKeyField();      break;

    default:
      return readSingleFieldNoTag(proto, field);
    }

    // collection or a map:

    if (field.getType() == TType.MAP) {

      proto.readByte();
      proto.readByte();
      int nEntries = proto.readI32();

      Map<Object, Object> map = Maps.newHashMap();
      Field valueField = field.getMapValueField();

      for (int i=0; i<nEntries; i++) {
        map.put(readFieldNoTag(proto, innerField),
                readFieldNoTag(proto, valueField));
      }
      return map;

    } else { // SET or LIST

      proto.readByte();
      int nEntries = proto.readI32();

      for(int i=0; i<nEntries; i++) {
        coll.add(readFieldNoTag(proto, innerField));
      }
      return coll;

    }
  }
}


package com.twitter.elephantbird.pig.util;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import com.google.protobuf.ByteString;
import com.google.protobuf.Descriptors;
import com.google.protobuf.Descriptors.Descriptor;
import com.google.protobuf.Descriptors.EnumValueDescriptor;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.Descriptors.DescriptorValidationException;
import com.google.protobuf.Message;
import com.google.protobuf.Message.Builder;
import com.google.protobuf.DescriptorProtos.DescriptorProto;
import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
import com.google.protobuf.DescriptorProtos.FieldDescriptorProto.Type;

import org.apache.pig.ResourceSchema;
import org.apache.pig.ResourceSchema.ResourceFieldSchema;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.DataByteArray;
import org.apache.pig.data.DataType;
import org.apache.pig.data.Tuple;
import org.apache.pig.impl.util.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.twitter.elephantbird.util.Protobufs;

/**
 * Converts a Pig Tuple into a Protobuf message. Tuple values should be ordered to match the natural
 * order of Protobuf field ordinal values. For example, say we define the following Protobuf
 * message:
 *
 *

 * message MyProtobufType {
 *   optional int32 f1 = 1;
 *   optional int32 f2 = 3;
 *   optional int32 f3 = 7;
 * }
 *

 *
 * Input Tuples are expected to contain field values in order {@code (f1, f2, f3)}. Tuples may
 * contain fewer values than Protobuf message fields (e.g. only {@code (f1, f2)} in the prior
 * example); Any remaining fields will be left unset.
 *
 * @author Vikram Oberoi
 */
public class PigToProtobuf {
  private static final Logger LOG = LoggerFactory.getLogger(PigToProtobuf.class);

  public PigToProtobuf() {}

  @SuppressWarnings("unchecked")
  public static <M extends Message> M tupleToMessage(Class<M> protoClass, Tuple tuple) {
    Builder builder = Protobufs.getMessageBuilder(protoClass);
    return (M) tupleToMessage(builder, tuple);
  }

  /**
   * Turn a Tuple into a Message with the given type.
   * @param builder a builder for the Message type the tuple will be converted to
   * @param tuple the tuple
   * @return a message representing the given tuple
   */
  public static Message tupleToMessage(Builder builder, Tuple tuple) {
    return tupleToMessage(builder, builder.getDescriptorForType().getFields(), tuple);
  }

  /**
   * @param builder
   * @param fieldDescriptors should be same as builder.getDescriptorForType.getFields().
   *        Avoids overhead of getFields() which creates an array each time.
   * @param tuple
   * @return
   */
  public static Message tupleToMessage(Builder builder, List<FieldDescriptor> fieldDescriptors, Tuple tuple) {
    if (tuple == null) {
      return  builder.build();
    }

    for (int i = 0; i < fieldDescriptors.size() && i < tuple.size(); i++) {
      Object tupleField = null;
      FieldDescriptor fieldDescriptor = fieldDescriptors.get(i);

      try {
        tupleField = tuple.get(i);
      } catch (ExecException e) {
        LOG.warn("Could not convert tuple field " + tupleField + " to field with descriptor " + fieldDescriptor);
        continue;
      }

      if (tupleField != null) {
        try {
          if (fieldDescriptor.isRepeated()) {
            // Repeated fields are set with Lists containing objects of the fields' Java type.
            builder.setField(fieldDescriptor,
                dataBagToRepeatedField(builder, fieldDescriptor, (DataBag) tupleField));
          } else {
            if (fieldDescriptor.getType() == FieldDescriptor.Type.MESSAGE) {
              Builder nestedMessageBuilder = builder.newBuilderForField(fieldDescriptor);
              builder.setField(fieldDescriptor,
                  tupleToMessage(nestedMessageBuilder, (Tuple) tupleField));
            } else {
              builder.setField(fieldDescriptor,
                  tupleFieldToSingleField(fieldDescriptor, tupleField));
            }
          }
        } catch (Exception e) {
          String value = String.valueOf(tupleField);
          final int max_length = 100;
          if (max_length < value.length()) {
            value = value.substring(0, max_length - 3) + "...";
          }
          String type = tupleField == null ? "unknown" : tupleField.getClass().getName();
          throw new RuntimeException(String.format(
              "Failed to set field '%s' using tuple value '%s' of type '%s' at index %d",
              fieldDescriptor.getName(), value, type, i), e);
        }
      }
    }

    return builder.build();
  }

  /**
   * For a given ResourceSchema, generate a protobufs Descriptor with analagous field names
   * and types.
   *
   * @param schema Pig schema.
   * @return Protobufs Descriptor
   * @throws Descriptors.DescriptorValidationException
   */
  public static Descriptor schemaToProtoDescriptor(ResourceSchema schema)
      throws DescriptorValidationException {
      return schemaToProtoDescriptor(schema, null);
  }

  /**
   * For a given ResourceSchema, generate a protobufs Descriptor with analogous field names
   * and types.
   *
   * @param schema Pig schema.
   * @param extraFields optionally pass a List of extra fields (Pairs of name:type) to be included.
   * @return Protobufs Descriptor
   * @throws Descriptors.DescriptorValidationException
   */
  public static Descriptor schemaToProtoDescriptor(ResourceSchema schema, List<Pair<String, Type>> extraFields)
      throws DescriptorValidationException {

    // init protobufs
    DescriptorProto.Builder desBuilder = DescriptorProto.newBuilder();

    int count = 0;
    for (ResourceFieldSchema fieldSchema : schema.getFields()) {
      // Pig types
      int position = ++count;
      String fieldName = fieldSchema.getName();
      byte dataTypeId = fieldSchema.getType();

      // determine and add protobuf types
      Type protoType = pigTypeToProtoType(dataTypeId);
      LOG.info("Mapping Pig field " + fieldName + " of type " + dataTypeId + " to protobuf type: " + protoType);

      addField(desBuilder, fieldName, position, protoType);
    }

    if (count == 0) {
      throw new IllegalArgumentException("ResourceSchema does not have any fields");
    }

    // If extra fields are needed, let's add them
    if (extraFields != null) {
      for (Pair<String, Type> extraField : extraFields) {
        addField(desBuilder, extraField.first, ++count, extraField.second);
      }
    }

    desBuilder.setName("PigToProtobufDynamicBuilder");

    return Protobufs.makeMessageDescriptor(desBuilder.build());
  }

  /**
   * Converts a DataBag into a List of objects with the type in the given FieldDescriptor. DataBags
   * don't map cleanly to repeated protobuf types, so each Tuple has to be unwrapped (by taking the
   * first element if the type is primitive or by converting the Tuple to a Message if the type is
   * MESSAGE), and the contents have to be appended to a List.
   * @param containingMessageBuilder a Message builder for the Message that contains this repeated field
   * @param fieldDescriptor a FieldDescriptor for this repeated field
   * @param bag the DataBag being serialized
   * @return a protobuf-friendly List of fieldDescriptor-type objects
   */
  private static List<Object> dataBagToRepeatedField(Builder containingMessageBuilder, FieldDescriptor fieldDescriptor, DataBag bag) {
    ArrayList<Object> bagContents = new ArrayList<Object>((int)bag.size());
    Iterator<Tuple> bagIter = bag.iterator();

    while (bagIter.hasNext()) {
      Tuple tuple = bagIter.next();
      if (fieldDescriptor.getType() == FieldDescriptor.Type.MESSAGE) {
        Builder nestedMessageBuilder = containingMessageBuilder.newBuilderForField(fieldDescriptor);
        bagContents.add(tupleToMessage((Builder)nestedMessageBuilder, tuple));
      } else {
        try {
          bagContents.add(tupleFieldToSingleField(fieldDescriptor, tuple.get(0)));
        } catch (ExecException e) {
          LOG.warn("Could not add a value for repeated field with descriptor " + fieldDescriptor);
        }
      }
    }

    return bagContents;
  }

  /**
   * Converts a tupleField string to its corresponding protobuf enum type if necessary, otherwise
   * returns the tupleField as is.
   * @param fieldDescriptor the FieldDescriptor for the given tuple field
   * @param tupleField the tupleField being converted to a protobuf field
   * @return the protobuf type for the given tupleField. This will be the tupleField itself unless it's an enum, in which case this will return the enum type for the field.
   */
  private static Object tupleFieldToSingleField(FieldDescriptor fieldDescriptor, Object tupleField) {
    // type convertion should match with ProtobufToPig.getPigScriptDataType
    switch (fieldDescriptor.getType()) {
    case ENUM:
      return toEnumValueDescriptor(fieldDescriptor, (String) tupleField);
    case BOOL:
      return Boolean.valueOf((Integer)tupleField != 0);
    case BYTES:
      return ByteString.copyFrom(((DataByteArray)tupleField).get());
    default:
      return tupleField;
    }
  }

  private static EnumValueDescriptor toEnumValueDescriptor(FieldDescriptor fieldDescriptor,
      String name) {
    EnumValueDescriptor out = fieldDescriptor.getEnumType().findValueByName(name);
    if (out == null) {
      throw new IllegalArgumentException(
          String.format("Failed to convert string '%s'" +
              " to enum value of type '%s'", name,
              fieldDescriptor.getEnumType().getFullName()));
    }
    return out;
  }

  /**
   * Add a field to a protobuf builder
   */
  private static void addField(DescriptorProto.Builder builder, String name, int fieldId, Type type) {
    FieldDescriptorProto.Builder fdBuilder = FieldDescriptorProto.newBuilder()
      .setName(name)
      .setNumber(fieldId)
      .setType(type);
    builder.addField(fdBuilder.build());
  }

  /**
   * For a given Pig type, return the protobufs type that maps to it.
   */
  private static Type pigTypeToProtoType(byte pigTypeId) {

    switch(pigTypeId) {
        case DataType.BOOLEAN:
          return Type.TYPE_BOOL;
        case DataType.INTEGER:
          return Type.TYPE_INT32;
        case DataType.LONG:
          return Type.TYPE_INT64;
        case DataType.FLOAT:
          return Type.TYPE_FLOAT;
        case DataType.DOUBLE:
          return Type.TYPE_DOUBLE;
        case DataType.CHARARRAY:
          return Type.TYPE_STRING;
        case DataType.BYTEARRAY:
          return Type.TYPE_BYTES;
        default:
          throw new IllegalArgumentException("Unsupported Pig type passed (" + pigTypeId +
              ") where a simple type is expected while converting Pig to a dynamic Protobuf");
    }
  }
}


package com.twitter.elephantbird.pig.load;

import com.google.common.collect.Maps;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.GnuParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Options;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.mapreduce.lib.input.TextInputFormat;
import org.apache.pig.PigException;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.NonSpillableDataBag;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
import org.apache.pig.impl.PigContext;
import org.json.simple.JSONArray;
import org.json.simple.JSONObject;
import org.json.simple.parser.JSONParser;
import org.json.simple.parser.ParseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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

/**
 * Decodes each line as JSON passes the resulting map of values
 * to Pig as a single-element tuple.
 *
 *
This Loader supports loading of nested JSON structures, but this feature
 * is disabled by default. There are two ways to enable it:

 *
setting the -nestedLoad option in the
 * {@link JsonLoader#JsonLoader(String)} constructor
 *
setting the elephantbird.jsonloader.nestedLoad
 * property to true. This can be done with the following pig command:
 *
grunt> set elephantbird.jsonloader.nestedLoad 'true'

 *

 */
public class JsonLoader extends LzoBaseLoadFunc {
  private static final Logger LOG = LoggerFactory.getLogger(JsonLoader.class);
  private static final TupleFactory tupleFactory = TupleFactory.getInstance();

  public static final String NESTED_LOAD_KEY = "elephantbird.jsonloader.nestedLoad";

  private final static Options validOptions_ = new Options();
  private final static CommandLineParser parser_ = new GnuParser();
  private final CommandLine configuredOptions_;

  private final JSONParser jsonParser = new JSONParser();

  private enum JsonLoaderCounters {
    LinesRead,
    LinesJsonDecoded,
    LinesParseError,
    LinesParseErrorBadNumber
  }

  private String inputFormatClassName;
  private boolean isNestedLoadEnabled = false;

  private static void populateValidOptions() {
    validOptions_.addOption("nestedLoad", false, "Enables loading of " +
        "nested JSON structures");
    validOptions_.addOption("inputFormat", true, "The input format class name" +
        " used by this loader instance");
  }

  public JsonLoader() {
    // defaults to no options
    this("");
  }

  /**
   * Constructor. Construct a JsonLoader LoadFunc to load.
   * @param inputFormatClassName the inputFormat class name
   *
   * @param optString Loader options. Available options:

   *
-nestedLoad==(true|false) Enables loading of nested JSON
   * structures. When enabled, JSON objects are loaded as nested Maps
   * and JSON arrays are loaded as Bags.
   *
-inputFormat=className The input format class name used
   * by this loader instance.
   *

   */
  public JsonLoader(String optString) {
    populateValidOptions();
    String[] optsArr = optString.split(" ");

    try {
      configuredOptions_ = parser_.parse(validOptions_, optsArr);
    } catch (org.apache.commons.cli.ParseException e) {
        HelpFormatter formatter = new HelpFormatter();
        formatter.printHelp( "[-nestedLoad] [-inputFormat]", validOptions_ );
        throw new RuntimeException(e);
    }
    isNestedLoadEnabled = configuredOptions_.hasOption("nestedLoad");

    if (configuredOptions_.getOptionValue("inputFormat") != null) {
      this.inputFormatClassName = configuredOptions_.getOptionValue("inputFormat");
    } else {
      this.inputFormatClassName = TextInputFormat.class.getName();
    }
  }

  public void setInputFormatClassName(String inputFormatClassName) {
    this.inputFormatClassName = inputFormatClassName;
  }

  /**
   * Return every non-null line as a single-element tuple to Pig.
   */
  @Override
  public Tuple getNext() throws IOException {
    if (reader == null) {
      return null;
    }
    // nested load can be enabled through a pig property
    if ("true".equals(jobConf.get(NESTED_LOAD_KEY)))
      isNestedLoadEnabled = true;
    try {
      while (reader.nextKeyValue()) {
        Text value = (Text) reader.getCurrentValue();
        incrCounter(JsonLoaderCounters.LinesRead, 1L);
        Tuple t = parseStringToTuple(value.toString());
        if (t != null) {
          incrCounter(JsonLoaderCounters.LinesJsonDecoded, 1L);
          return t;
        }
      }
      return null;

    } catch (InterruptedException e) {
      int errCode = 6018;
      String errMsg = "Error while reading input";
      throw new ExecException(errMsg, errCode,
          PigException.REMOTE_ENVIRONMENT, e);
    }
  }

  @Override
  public InputFormat getInputFormat() throws IOException {
    try {
      return (FileInputFormat) PigContext.resolveClassName(inputFormatClassName).newInstance();
    } catch (InstantiationException e) {
      throw new IOException("Failed creating input format " + inputFormatClassName, e);
    } catch (IllegalAccessException e) {
      throw new IOException("Failed creating input format " + inputFormatClassName, e);
    }
  }

  protected Tuple parseStringToTuple(String line) {
    try {
      JSONObject jsonObj = (JSONObject) jsonParser.parse(line);
      if (jsonObj != null) {
        return tupleFactory.newTuple(walkJson(jsonObj));
      } else {
        // JSONParser#parse(String) may return a null reference, e.g. when
        // the input parameter is the string "null".  A single line with
        // "null" is not valid JSON though.
        LOG.warn("Could not json-decode string: " + line);
        incrCounter(JsonLoaderCounters.LinesParseError, 1L);
        return null;
      }
    } catch (ParseException e) {
      LOG.warn("Could not json-decode string: " + line, e);
      incrCounter(JsonLoaderCounters.LinesParseError, 1L);
      return null;
    } catch (NumberFormatException e) {
      LOG.warn("Very big number exceeds the scale of long: " + line, e);
      incrCounter(JsonLoaderCounters.LinesParseErrorBadNumber, 1L);
      return null;
    } catch (ClassCastException e) {
      LOG.warn("Could not convert to Json Object: " + line, e);
      incrCounter(JsonLoaderCounters.LinesParseError, 1L);
      return null;
    }
  }

  private Object wrap(Object value) {

    if (isNestedLoadEnabled && value instanceof JSONObject) {
      return walkJson((JSONObject) value);
    }  else if (isNestedLoadEnabled && value instanceof JSONArray) {

      JSONArray a = (JSONArray) value;
      DataBag mapValue = new NonSpillableDataBag(a.size());
      for (int i=0; i<a.size(); i++) {
        Tuple t = tupleFactory.newTuple(wrap(a.get(i)));
        mapValue.add(t);
      }
      return mapValue;

    } else {
      return value != null ? value.toString() : null;
    }
  }

  private Map<String,Object> walkJson(JSONObject jsonObj) {
    Map<String,Object> v = Maps.newHashMap();
    for (Object key: jsonObj.keySet()) {
      v.put(key.toString(), wrap(jsonObj.get(key)));
    }
    return v;
  }

}


package com.twitter.elephantbird.pig.util;

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

import com.google.common.base.Charsets;

import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.DataByteArray;
import org.apache.pig.data.Tuple;
import org.apache.thrift.TBase;
import org.apache.thrift.TEnum;
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.ThriftUtils;
import com.twitter.elephantbird.util.TypeRef;

/**
 * Converts a Pig Tuple into a Thrift struct. Tuple values should be ordered to match the natural
 * order of Thrift field ordinal values. For example, say we define the following Thrift struct:
 *
 * <pre>
 * struct MyThriftType {
 *   1: i32 f1
 *   3: i32 f2
 *   7: i32 f3
 * }
 * </pre>
 *
 * Input Tuples are expected to contain field values in order {@code (f1, f2, f3)}. Tuples may
 * contain fewer values than Thrift struct fields (e.g. only {@code (f1, f2)} in the prior example);
 * Any remaining fields will be left unset.
 */
public class PigToThrift<T extends TBase<?, ?>> {
  public static final Logger LOG = LoggerFactory.getLogger(PigToThrift.class);

  private TStructDescriptor structDesc;

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

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

  public PigToThrift(Class<T> tClass) {
    structDesc = TStructDescriptor.getInstance(tClass);
    // may be TODO : compare the schemas to catch errors early.
  }

  @SuppressWarnings("unchecked")
  public T getThriftObject(Tuple tuple) {
    return (T)toThrift(structDesc, tuple);
  }

  /**
   * Construct a Thrift object from the tuple.
   */
  @SuppressWarnings("unchecked")
  private static TBase<?, ?> toThrift(TStructDescriptor tDesc, Tuple tuple) {
    int size = tDesc.getFields().size();
    int tupleSize = tuple.size();
    @SuppressWarnings("rawtypes")
    TBase tObj = newTInstance(tDesc.getThriftClass());
    for(int i = 0; i<size && i<tupleSize; i++) {
      Object pObj;
      try {
        pObj = tuple.get(i);
      } catch (ExecException e) {
        throw new RuntimeException(e);
      }
      if (pObj != null) {
        Field field = tDesc.getFieldAt(i);
        try {
          tObj.setFieldValue(field.getFieldIdEnum(), toThriftValue(field, pObj));
        } catch (Exception e) {
          String value = String.valueOf(tObj);
          final int max_length = 100;
          if (max_length < value.length()) {
            value = value.substring(0, max_length - 3) + "...";
          }
          String type = tObj == null ? "unknown" : tObj.getClass().getName();
          throw new RuntimeException(String.format(
              "Failed to set field '%s' using tuple value '%s' of type '%s' at index %d",
              field.getName(), value, type, i), e);
        }
      }
      // if tDesc is a union, at least one field needs to be non-null.
      // user is responsible for ensuring that.
    }
    return tObj;
  }

  /**
   * For a given Pig value, return a Thrift object of the same type as the Thrift field passed. The
   * thrift field is expected to be compatible with the value passed. If it is not, a warning will
   * be logged and a null value will be returned.
   *
   * @param thriftField the Thrift field used to determine the type of the response object
   * @param pigValue the value to convert to Thrift
   * @return a Thrift object
   */
  @SuppressWarnings("unchecked")
  public static Object toThriftValue(Field thriftField, Object pigValue) {
    try {
      switch (thriftField.getType()) {
      case TType.BOOL:
        return Boolean.valueOf(((Integer)pigValue) != 0);
      case TType.BYTE :
        return ((Integer)pigValue).byteValue();
      case TType.I16 :
        return Short.valueOf(((Integer)pigValue).shortValue());
      case TType.STRING:
        return toStringType(pigValue);
      case TType.STRUCT:
        return toThrift(thriftField.gettStructDescriptor(), (Tuple)pigValue);
      case TType.MAP:
        return toThriftMap(thriftField, (Map<String, Object>)pigValue);
      case TType.SET:
        return toThriftSet(thriftField.getSetElemField(), (DataBag) pigValue);
      case TType.LIST:
        return toThriftList(thriftField.getListElemField(), (DataBag)pigValue);
      case TType.ENUM:
        return toThriftEnum(thriftField, (String) pigValue);
      default:
        // standard types : I32, I64, DOUBLE, etc.
        return pigValue;
      }
    } catch (Exception e) {
      // mostly a schema mismatch.
      LOG.warn(String.format(
          "Failed to set field '%s' of type '%s' with value '%s' of type '%s'",
          thriftField.getName(), ThriftUtils.getFieldValueType(thriftField).getName(),
          pigValue, pigValue.getClass().getName()), e);
    }
    return null;
  }

  /* TType.STRING could be either a DataByteArray or a String */
  private static Object toStringType(Object value) {
    if (value instanceof String) {
      return value;
    } else if (value instanceof DataByteArray) {
      byte[] buf = ((DataByteArray)value).get();
      // mostly there is no need to copy.
      return ByteBuffer.wrap(Arrays.copyOf(buf, buf.length));
    }
    return null;
  }

  private static Map<Object, Object> toThriftMap(Field field, Map<String, Object> map) {
    Field keyField   = field.getMapKeyField();
    Field valueField = field.getMapValueField();
    HashMap<Object, Object> out = new HashMap<Object, Object>(map.size());
    for(Entry<String, Object> e : map.entrySet()) {
      String s = e.getKey();
      Object key;
      switch (keyField.getType()) {
        case TType.STRING: key = s; break;
        case TType.BOOL: key = Boolean.parseBoolean(s); break;
        case TType.BYTE: key = Byte.parseByte(s); break;
        case TType.I16: key = Short.parseShort(s); break;
        case TType.I32: key = Integer.parseInt(s); break;
        case TType.I64: key = Long.parseLong(s); break;
        case TType.DOUBLE: key = Double.parseDouble(s); break;
        case TType.ENUM: key = toThriftEnum(keyField, s); break;
        default:
          // LIST, MAP, SET, STOP, STRUCT, VOID types are unsupported
          throw new RuntimeException(String.format(
              "Conversion from string map key to type '%s' is unsupported",
              ThriftUtils.getFieldValueType(keyField).getName()));
      }
      if (keyField.isBuffer()) {
        key = ByteBuffer.wrap(s.getBytes(Charsets.UTF_8));
      }
      out.put(key, toThriftValue(valueField, e.getValue()));
    }
    return out;
  }

  private static Set<Object> toThriftSet(Field elemField, DataBag bag) {
    Set<Object> set = new HashSet<Object>((int)bag.size());
    fillThriftCollection(set, elemField, bag);
    return set;
  }

  private static List<Object> toThriftList(Field elemField, DataBag bag) {
    List<Object> list = new ArrayList<Object>((int)bag.size());
    fillThriftCollection(list, elemField, bag);
    return list;
  }

  private static TEnum toThriftEnum(Field elemField, String name) {
    TEnum out = elemField.getEnumValueOf(name);
    if (out == null) {
      throw new IllegalArgumentException(
          String.format("Failed to convert string '%s'" +
              " to enum value of type '%s'", name,
              ThriftUtils.getFieldValueType(elemField).getName()));
    }
    return out;
  }

  private static void fillThriftCollection(Collection<Object> tColl, Field elemField, DataBag bag) {
    for (Tuple tuple : bag) {
      if (!elemField.isStruct()) {
        // this tuple is a just wrapper for another object.
        try {
          tColl.add(toThriftValue(elemField, tuple.get(0)));
        } catch (ExecException e) {
          throw new RuntimeException(e);
        }
      } else { // tuple for a struct.
        tColl.add(toThriftValue(elemField, tuple));
      }
    }
  }

  /** return an instance assuming tClass is a Thrift class */
  private static TBase<?, ?> newTInstance(Class<?> tClass) {
    try {
      return (TBase<?, ?>) tClass.newInstance();
    } catch (Exception e) { // not expected.
      throw new RuntimeException(e);
    }
  }
}


package com.twitter.elephantbird.mapreduce.input.combine;

import com.twitter.elephantbird.mapred.input.DeprecatedInputFormatWrapper;
import com.twitter.elephantbird.util.HadoopCompat;
import com.twitter.elephantbird.util.HadoopUtils;
import com.twitter.elephantbird.util.SplitUtil;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.mapreduce.*;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.util.ReflectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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

/**
 * This class allows for combining the InputSplit of an underlying {@link InputFormat}
 * in a way which functions properly with Elephant Bird's other input formats. It works
 * seamlessly with {@link com.twitter.elephantbird.mapred.input.DeprecatedInputFormatWrapper}
 * which means that any of Elephant Bird's Input Formats can be used in Cascading
 * as a CombineFileInputFormat.
 *
 * @author Jonathan Coveney
 */
public class DelegateCombineFileInputFormat<K, V> extends FileInputFormat<K, V> {
  private static final Logger LOG = LoggerFactory.getLogger(DelegateCombineFileInputFormat.class);

  public static final String USE_COMBINED_INPUT_FORMAT = "elephantbird.use.combine.input.format";
  public static final String COMBINED_INPUT_FORMAT_DELEGATE = "elephantbird.combine.input.format.delegate";

  public static void setUseCombinedInputFormat(Configuration conf) {
    conf.setBoolean(USE_COMBINED_INPUT_FORMAT, true);
  }

  // This sets configures the delegate, though it does not configure DelegateCombineFileInputFormat.
  public static void setCombinedInputFormatDelegate(Configuration conf, Class<? extends InputFormat> clazz) {
    HadoopUtils.setClassConf(conf, COMBINED_INPUT_FORMAT_DELEGATE, clazz);
  }

  private InputFormat<K, V> delegate;
  private long maxSplitSize;
  private long minSplitSizeNode;
  private long minSplitSizeRack;

  public static void setDelegateInputFormat(JobConf jobConf, Class<? extends InputFormat> inputFormat) {
    DeprecatedInputFormatWrapper.setInputFormat(DelegateCombineFileInputFormat.class, jobConf);
    setCombinedInputFormatDelegate(jobConf, inputFormat);
  }

  public static void setDelegateInputFormat(Configuration conf, Class<? extends InputFormat> inputFormat) {
    DeprecatedInputFormatWrapper.setInputFormat(DelegateCombineFileInputFormat.class, conf);
    setCombinedInputFormatDelegate(conf, inputFormat);
  }

  public static void setDelegateInputFormatHadoop2(Configuration conf, Class<? extends InputFormat> inputFormat) {
    conf.setClass("mapreduce.inputformat.class", DelegateCombineFileInputFormat.class, InputFormat.class);
    setCombinedInputFormatDelegate(conf, inputFormat);
  }

  private void initInputFormat(Configuration conf) throws IOException {
    if (delegate == null) {
      Class<? extends InputFormat> delegateClass =
        conf.getClass(COMBINED_INPUT_FORMAT_DELEGATE, null, InputFormat.class);
      if (delegateClass == null) {
        throw new IOException("No delegate class was set on key: " + COMBINED_INPUT_FORMAT_DELEGATE);
      }
      delegate = ReflectionUtils.newInstance(delegateClass, conf);
    }
  }

  public DelegateCombineFileInputFormat() {
    // Will instantiate the delegate via reflection
  }

  public DelegateCombineFileInputFormat(InputFormat<K, V> delegate) {
    this.delegate = delegate;
  }

  private boolean shouldCombine(Configuration conf) {
    return conf.getBoolean(USE_COMBINED_INPUT_FORMAT, false);
  }

  @Override
  public RecordReader createRecordReader(
          InputSplit inputSplit, TaskAttemptContext taskAttemptContext)
      throws IOException, InterruptedException {
    Configuration conf = HadoopCompat.getConfiguration(taskAttemptContext);
    initInputFormat(conf);
    if (shouldCombine(conf)) {
      return new CompositeRecordReader(delegate);
    } else {
      return delegate.createRecordReader(inputSplit, taskAttemptContext);
    }
  }

  @Override
  public List<InputSplit> getSplits(JobContext job) throws IOException {
    Configuration conf = HadoopCompat.getConfiguration(job);
    initInputFormat(conf);
    try {
      if (shouldCombine(conf)) {
          List<InputSplit> inputSplits = delegate.getSplits(job);
          List<InputSplit> combinedInputSplits = new ArrayList<InputSplit>();
          for (CompositeInputSplit split : SplitUtil.getCombinedCompositeSplits(inputSplits, conf)) {
            split.setConf(conf);
            combinedInputSplits.add(split);
          }
          return combinedInputSplits;
      } else {
        return delegate.getSplits(job);
      }
    } catch (InterruptedException e) {
      LOG.error("Thread interrupted", e);
      Thread.currentThread().interrupt();
      throw new IOException(e);
    }
  }
}


package com.twitter.elephantbird.mapreduce.input;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;

import com.twitter.elephantbird.util.HadoopCompat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.fs.PathFilter;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.JobContext;

import com.twitter.elephantbird.mapreduce.output.LuceneIndexOutputFormat;
import com.twitter.elephantbird.util.HadoopUtils;
import com.twitter.elephantbird.util.HdfsUtils;

/**
 * Base class for input formats that read lucene indexes stored in HDFS directories.
 * Given a list of indexes and queries, runs each query over each index. Implements split
 * combining (combines multiple indexes into one split) based on
 * the total size of the index directory and the configured max combined split size.
 *

 * Emits key, value records where key is the query that resulted in value
 * (key is actually the position in the list of queries, not the query string itself)
 *

 * Subclasses must provide:
 *

 *
a {@link LuceneIndexRecordReader} which describes how to convert a String into a
 *      Query and how to convert a Document into a value of type T

 *

 * Subclasses may provide:d
 *

 *
 a {@link PathFilter} for identifying HDFS directories that contain Lucene indexes

 *

 *
 * @param <T> - the type that your lucene Documents will be converted to
 * @author Alex Levenson
 */
public abstract class LuceneIndexInputFormat<T extends Writable>
    extends InputFormat<IntWritable, T> {

  public static final String QUERIES_KEY = LuceneIndexInputFormat.class.getCanonicalName()
    + ".queries";

  public static final String INPUT_PATHS_KEY = LuceneIndexInputFormat.class.getCanonicalName()
    + ".inputpaths";

  public static final String MAX_NUM_INDEXES_PER_SPLIT_KEY =
    LuceneIndexInputFormat.class.getCanonicalName() + ".max_num_indexes_per_split";

  // Empirically it seems that 200 is a reasonable number of small
  // indexes for one mapper to process in a few minutes, and thousands of
  // of small indexes can take one mapper more than 30 minutes to process
  // due to overhead for each index.
  private static final long DEFAULT_MAX_NUM_INDEXES_PER_SPLIT = 200;

  public static final String MAX_COMBINED_INDEX_SIZE_PER_SPLIT_KEY =
    LuceneIndexInputFormat.class.getCanonicalName() + ".max_combined_index_size_per_split";

  // default to 10GB
  // back of the envelope reasoning:
  // Assume 1 mapper should process 1 GB, and each index will return 1/10th its size in records
  private static final long DEFAULT_MAX_COMBINED_INDEX_SIZE_PER_SPLIT = 10*1024*1024*1024L;

  private static final String[] EMPTY_NODE_ARRAY = new String[0];

  private Path[] inputPaths = null;
  private PathFilter indexDirPathFilter = null;
  private long maxCombinedIndexSizePerSplit;
  private long maxNumIndexesPerSplit;

  /**
   * Subclasses may provide a {@link PathFilter} for identifying HDFS
   * directories that contain lucene indexes. When directories are being
   * searched recursively for index directories, this path filter will be used
   * to determine if a directory is a lucene index.
   *

   * The default is to treat any directory whose name begins with "-index" as a lucene index,
   * which matches what {@link LuceneIndexOutputFormat} generates.
   *
   * @param conf job conf
   * @return a path filter that accepts directories with lucene indexes in them
   * @throws IOException
   */
  public PathFilter getIndexDirPathFilter(Configuration conf) throws IOException {
    return LuceneIndexOutputFormat.newIndexDirFilter(conf);
  }

  @VisibleForTesting
  void loadConfig(Configuration conf) throws IOException {
    inputPaths = getInputPaths(conf);

    indexDirPathFilter = Preconditions.checkNotNull(getIndexDirPathFilter(conf),
      "You must provide a non-null PathFilter");

    maxCombinedIndexSizePerSplit = Preconditions.checkNotNull(
        getMaxCombinedIndexSizePerSplit(conf),
        MAX_COMBINED_INDEX_SIZE_PER_SPLIT_KEY + " cannot be null");

    maxNumIndexesPerSplit = Preconditions.checkNotNull(getMaxNumIndexesPerSplit(conf),
        MAX_NUM_INDEXES_PER_SPLIT_KEY + " cannot be null");
  }

  /**
   * Creates splits with multiple indexes per split
   * (if they are smaller than maxCombinedIndexSizePerSplit).
   * It is possible for a split to be larger than maxCombinedIndexSizePerSplit,
   * if it consists of a single index that is
   * larger than maxCombinedIndexSizePerSplit.
   *

   * All inputPaths will be searched for indexes recursively
   *

   * The bin-packing problem of combining splits is solved naively:
   *

   *
Sort all indexes by size

   *
Begin packing indexes into splits until adding the next split would cause the split to
   *       exceed maxCombinedIndexSizePerSplit

   *
Begin packing subsequent indexes into the next split, and so on

   *

   */
  @Override
  public List<InputSplit> getSplits(JobContext job) throws IOException, InterruptedException {

    // load settings from job conf
    loadConfig(HadoopCompat.getConfiguration(job));

    // find all the index dirs and create a split for each
    PriorityQueue<LuceneIndexInputSplit> splits = findSplits(HadoopCompat.getConfiguration(job));

    // combine the splits based on maxCombineSplitSize
    List<InputSplit> combinedSplits = combineSplits(splits, maxCombinedIndexSizePerSplit,
      maxNumIndexesPerSplit);

    return combinedSplits;
  }

  /**
   * Finds and creates all the index splits based on the input paths set in conf
   * @param conf job conf
   * @return a priority queue of the splits, default is sorted by size
   * @throws IOException
   */
  protected PriorityQueue<LuceneIndexInputSplit> findSplits(Configuration conf) throws IOException {
    PriorityQueue<LuceneIndexInputSplit> splits = new PriorityQueue<LuceneIndexInputSplit>();
    List<Path> indexDirs = Lists.newLinkedList();

    // find all indexes nested under all the input paths
    // (which happen to be directories themselves)
    for (Path path : inputPaths) {
      HdfsUtils.collectPaths(path, path.getFileSystem(conf), indexDirPathFilter, indexDirs);
    }

    // compute the size of each index
    // and create a single split per index
    for (Path indexDir : indexDirs) {
      long size = HdfsUtils.getDirectorySize(indexDir, indexDir.getFileSystem(conf));
      splits.add(new LuceneIndexInputSplit(Lists.newLinkedList(Arrays.asList(indexDir)), size));
    }
    return splits;
  }

  protected List<InputSplit> combineSplits(PriorityQueue<LuceneIndexInputSplit> splits,
                                           long maxCombinedIndexSizePerSplit,
                                           long maxNumIndexesPerSplit) {

    // now take the one-split-per-index splits and combine them into multi-index-per-split splits
    List<InputSplit> combinedSplits = Lists.newLinkedList();
    LuceneIndexInputSplit currentSplit = splits.poll();
    while (currentSplit != null) {
      while (currentSplit.getLength() < maxCombinedIndexSizePerSplit) {
        LuceneIndexInputSplit nextSplit = splits.peek();
        if (nextSplit == null) {
          break;
        }
        if (currentSplit.getLength() + nextSplit.getLength() > maxCombinedIndexSizePerSplit) {
          break;
        }
        if (currentSplit.getIndexDirs().size() >= maxNumIndexesPerSplit) {
          break;
        }
        currentSplit.combine(nextSplit);
        splits.poll();
      }
      combinedSplits.add(currentSplit);
      currentSplit = splits.poll();
    }
    return combinedSplits;
  }

  /**
   * Set the queries to run over the indexes
   * These can be Strings suitable for parsing with a QueryParser, or they can be
   * a custom serialized Query object. They have to be Strings so that they can be
   * written to the job conf. They will be deserialized / parsed by the abstract
   * method {@link LuceneIndexRecordReader#deserializeQuery(String)}
   *
   * @param queries queries to run over the indexes
   * @param conf job conf
   * @throws IOException
   */
  public static void setQueries(List<String> queries, Configuration conf) throws IOException {
    Preconditions.checkNotNull(queries);
    Preconditions.checkArgument(!queries.isEmpty());
    HadoopUtils.writeStringListToConfAsBase64(QUERIES_KEY, queries, conf);
  }

  /**
   * Get the queries to run over the indexes
   * These are the queries in the same form as they were passed to {@link #setQueries} above
   *
   * @param conf job conf
   * @return queries as passed to {@link #setQueries}
   * @throws IOException
   */
  public static List<String> getQueries(Configuration conf) throws IOException {
    return Preconditions.checkNotNull(HadoopUtils.readStringListFromConfAsBase64(QUERIES_KEY, conf),
      "You must call LuceneIndexInputFormat.setQueries()");
  }

  /**
   * Check whether queries have been set yet for this job, useful for lazy loading
   * queries into the config
   *
   * @param conf job conf
   * @return whether the queries have been set yet
   */
  public static boolean queriesSet(Configuration conf) {
    return conf.get(QUERIES_KEY) != null;
  }

  /**
   * Sets the input paths for for this input format.
   * All paths will be searched for indexes recursively
   *
   * @param paths the input paths
   * @param conf the job conf
   * @throws IOException
   */
  public static void setInputPaths(List<Path> paths, Configuration conf) throws IOException {
    Preconditions.checkNotNull(paths);
    Preconditions.checkArgument(!paths.isEmpty());
    String[] pathStrs = new String[paths.size()];
    int i = 0;
    for (Path p : paths) {
      FileSystem fs = p.getFileSystem(conf);
      pathStrs[i++] = fs.makeQualified(p).toString();
    }
    conf.setStrings(INPUT_PATHS_KEY, pathStrs);
  }

  /**
   * Gets the input paths for this input format
   *
   * @param conf the job conf
   */
  public static Path[] getInputPaths(Configuration conf) {
    String[] pathStrs = Preconditions.checkNotNull(conf.getStrings(INPUT_PATHS_KEY),
        "You must call LuceneIndexInputFormat.setInputPaths()");
    Path[] paths = new Path[pathStrs.length];
    for (int i = 0; i < pathStrs.length; i++) {
      paths[i] = new Path(pathStrs[i]);
    }
    return paths;
  }

  /**
   * Set the max combined size of indexes to be processed by one split.
   *
   * If an index is larger than size than it will be put in its own split, but all splits
   * containing multiple indexes will have a combined size <= size.
   *
   * @param size the max combined size of indexes to be processed by one split in bytes
   * @param conf job conf
   */
  public static void setMaxCombinedIndexSizePerSplitBytes(long size, Configuration conf) {
    conf.setLong(MAX_COMBINED_INDEX_SIZE_PER_SPLIT_KEY, size);
  }

  /**
   * Get the max size of a combined split in bytes
   * @param conf job conf
   * @return the max size of a combined split in bytes
   */
  public static long getMaxCombinedIndexSizePerSplit(Configuration conf) {
    return conf.getLong(MAX_COMBINED_INDEX_SIZE_PER_SPLIT_KEY,
      DEFAULT_MAX_COMBINED_INDEX_SIZE_PER_SPLIT);
  }

  /**
   * Set the max number of indexes to process for a single split.
   *
   * If a combined split still has room for more indexes (as determined by
   * {@link #getMaxCombinedIndexSizePerSplit}) then more indexes will be added to it
   * UNLESS that would cause the split to have more than num indexes combined into it.
   *
   * This helps preventing one split from getting thousands of small indexes which can make it
   * significantly slower than the others.
   *
   * @param num max number of splits per combined split
   * @param conf job conf
   */
  public static void setMaxNumIndexesPerSplit(long num, Configuration conf) {
    conf.setLong(MAX_NUM_INDEXES_PER_SPLIT_KEY, num);
  }

  /**
   * Get the max number of indexes per split
   *
   * @param conf job conf
   * @return the max number of indexes per split
   */
  public static long getMaxNumIndexesPerSplit(Configuration conf) {
    return conf.getLong(MAX_NUM_INDEXES_PER_SPLIT_KEY, DEFAULT_MAX_NUM_INDEXES_PER_SPLIT);
  }

  /**
   * A split that represents multiple index (directories).
   * Has to be {@link Writable} to work with pig
   */
  public static class LuceneIndexInputSplit extends InputSplit
                                            implements Writable, Comparable<LuceneIndexInputSplit> {
    private List<Path> indexDirs;
    private Long length;

    /**
     * Required for instantiation by reflection
     */
    public LuceneIndexInputSplit() { }

    /**
     * Constructor for this class.
     * @param indexDirs a {@link java.util.List} of directories
     * containing existing Lucene indexes.
     * @param length the size of the split in bytes
     */
    public LuceneIndexInputSplit(List<Path> indexDirs, long length) {
      this.indexDirs = indexDirs;
      this.length = length;
    }

    /**
     * Merge other into this split
     * Will have no effect on other
     * @param other the split to combine
     */
    public void combine(LuceneIndexInputSplit other) {
      indexDirs.addAll(other.getIndexDirs());
      length += other.getLength();
    }

    /**
     * Get the size of this split in bytes
     * @return the size of this split in bytes
     */
    @Override
    public long getLength() {
      return length;
    }

    /**
     * Because an index consists of multiple (multi-block) files there's not much to be gained from
     * finding nodes where there is locality
     * @return an empty String[]
     */
    @Override
    public String[] getLocations() throws IOException, InterruptedException {
      return EMPTY_NODE_ARRAY;
    }

    /**
     * @return the list of indexes in this split (which are directories)
     */
    public List<Path> getIndexDirs() {
      return indexDirs;
    }

    @Override
    public void write(DataOutput out) throws IOException {
      out.writeLong(length);
      out.writeInt(indexDirs.size());
      for(Path p : indexDirs) {
        Text.writeString(out, p.toString());
      }
    }

    @Override
    public void readFields(DataInput in) throws IOException {
      this.length = in.readLong();
      int numDirs = in.readInt();
      this.indexDirs = Lists.newLinkedList();
      for (int i = 0; i < numDirs; i++) {
        String path = Text.readString(in);
        this.indexDirs.add(new Path(path));
      }
    }

    /**
     * sorts by length (size in bytes)
     */
    @Override
    public int compareTo(LuceneIndexInputSplit other) {
      return length.compareTo(other.getLength());
    }

    /**
     * Prints the directories and the combined length/size of the split in bytes.
     */
    @Override
    public String toString() {
      return "LuceneIndexInputSplit<indexDirs:" + indexDirs.toString() + " length:" + length + ">";
    }
  }
}


package com.twitter.elephantbird.mapreduce.input;

import java.io.IOException;
import java.util.BitSet;
import java.util.Iterator;

import com.google.common.collect.AbstractIterator;

import org.apache.hadoop.io.Writable;
import org.apache.hadoop.util.Progressable;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Scorer;

/**
 * Collects all documents that match queries
 *
 * @author Alex Levenson
 */
public abstract class LuceneIndexCollectAllRecordReader<T extends Writable>
    extends LuceneIndexRecordReader<T> {

  // TODO: This bit set can get big (up to about Integer.MAX_VALUE bits),
  // TODO: ideally it should get thrown away when its no longer in use.
  // TODO: However, it's difficult to know when that is, and the jvm will be torn down after
  // TODO: each map task anyway.
  private static final BitSet docIds = new BitSet();

  /**
   * Convert a {@link Document} to a value to be emitted by this record reader
   *
   * @param doc document to convert
   * @return a value to be emitted from this record reader
   */
  protected abstract T docToValue(Document doc) throws IOException;

  /**
   * Applies {@link #docToValue(Document)} to every document
   * found by executing query over searcher
   *
   * @param searcher the index searcher to query
   * @param query the query to run
   * @return a list of values to be emitted as records (one by one) by this record reader
   * @throws IOException
   */
  @Override
  protected Iterator<T> search(final IndexSearcher searcher, final Query query) throws IOException {
    // grow the bit set if needed
    docIds.set(searcher.getIndexReader().maxDoc());
    // clear it
    docIds.clear();
    searcher.search(query, new Collector() {
      private int docBase;

      @Override
      public void setScorer(Scorer scorer) {
      }

      @Override
      public boolean acceptsDocsOutOfOrder() {
        return true;
      }

      @Override
      public void collect(int doc) {
        docIds.set(doc + docBase);
      }

      public void setNextReader(AtomicReaderContext context) {
        this.docBase = context.docBase;
      }
    });

    return new AbstractIterator<T>() {
      private int doc = docIds.nextSetBit(0);

      @Override
      protected T computeNext() {
        ((Progressable)context).progress(); // casting to avoid Hadoop 2 incompatibility
        if (doc < 0) {
          return endOfData();
        }
        try {
          T ret = docToValue(searcher.doc(doc));
          doc = docIds.nextSetBit(doc + 1);
          return  ret;
        } catch (IOException e) {
          throw new RuntimeException(e);
        }
      }
    };
  }
}