RCFileProtobufPigStorage.java

package com.twitter.elephantbird.pig.store;

import java.io.IOException;

import org.apache.hadoop.mapreduce.OutputFormat;
import org.apache.pig.data.Tuple;

import com.google.protobuf.Message;
import com.google.protobuf.Message.Builder;
import com.twitter.elephantbird.mapreduce.io.ProtobufWritable;
import com.twitter.elephantbird.mapreduce.output.RCFileProtobufOutputFormat;
import com.twitter.elephantbird.pig.util.PigToProtobuf;
import com.twitter.elephantbird.util.Protobufs;
import com.twitter.elephantbird.util.TypeRef;

/**
 * StoreFunc for storing Protobuf messages in RCFiles. <p>
 *
 * @see {@link RCFileProtobufOutputFormat}
 */
public class RCFileProtobufPigStorage extends BaseStoreFunc {
  // add stats?

  private TypeRef<? extends Message> typeRef;
  private Builder msgBuilder;
  private ProtobufWritable<Message> writable;

  public RCFileProtobufPigStorage(String protoClassName) {
    typeRef = Protobufs.getTypeRef(protoClassName);
    msgBuilder = Protobufs.getMessageBuilder(typeRef.getRawClass());
    writable = ProtobufWritable.newInstance(Message.class);
  }

  @Override @SuppressWarnings("unchecked")
  public OutputFormat getOutputFormat() throws IOException {
    return new RCFileProtobufOutputFormat(typeRef);
  }

  public void putNext(Tuple t) throws IOException {
    Message msg = PigToProtobuf.tupleToMessage(msgBuilder.clone(), t);
    writable.set(msg);
    writeRecord(null, writable);
  }
}