com.twitter.elephantbird.pig.store.LzoProtobufBlockPigStorage Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of elephant-bird-pig Show documentation

Pig utilities.

There is a newer version: 4.17

package com.twitter.elephantbird.pig.store;

import java.io.IOException;

import org.apache.hadoop.mapreduce.OutputFormat;
import org.apache.pig.data.Tuple;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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


/**
 * Serializes Pig Tuples into Block encoded protocol buffers.
 * The fields in the pig tuple must correspond exactly to the fields in the protobuf, as
 * no name-matching is performed (that's a TODO).

 *
 *
 * @param  Protocol Buffer Message class being serialized
 */
public class LzoProtobufBlockPigStorage extends BaseStoreFunc {
  private static final Logger LOG = LoggerFactory.getLogger(LzoProtobufBlockPigStorage.class);

  private TypeRef typeRef_;
  private ProtobufWritable writable;

  public LzoProtobufBlockPigStorage() {}

  public LzoProtobufBlockPigStorage(String protoClassName) {
    TypeRef typeRef = Protobufs.getTypeRef(protoClassName);
    setTypeRef(typeRef);
  }

  protected void setTypeRef(TypeRef typeRef) {
    typeRef_ = typeRef;
    writable = ProtobufWritable.newInstance(typeRef_.getRawClass());
  }

  @SuppressWarnings("unchecked")
  @Override
  public void putNext(Tuple f) throws IOException {
    if (f == null) {
      return;
    }
    Builder builder = Protobufs.getMessageBuilder(typeRef_.getRawClass());
    try {
      writable.set((M) PigToProtobuf.tupleToMessage(builder, f));
      writer.write(null, writable);
    } catch (InterruptedException e) {
      // TODO Auto-generated catch block
      e.printStackTrace();
    }
  }

  @Override
  public OutputFormat> getOutputFormat() throws IOException {
    if (typeRef_ == null) {
      LOG.error("Protobuf class must be specified before an OutputFormat can be created. Do not use the no-argument constructor.");
      throw new IllegalArgumentException("Protobuf class must be specified before an OutputFormat can be created. Do not use the no-argument constructor.");
    }
    return new LzoProtobufBlockOutputFormat(typeRef_);
  }

}