TestRCFileThriftStorage.java

package com.twitter.elephantbird.pig.load;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Iterator;

import com.twitter.elephantbird.pig.util.PigTestUtil;
import com.twitter.elephantbird.util.HadoopCompat;
import com.twitter.elephantbird.util.CoreTestUtil;

import org.apache.commons.codec.binary.Base64;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileUtil;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.BytesWritable;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.OutputFormat;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.TaskAttemptID;
import org.apache.pig.PigServer;
import org.apache.pig.data.DataByteArray;
import org.apache.pig.data.Tuple;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import com.google.common.collect.ImmutableMap;
import com.twitter.elephantbird.mapreduce.io.ThriftConverter;
import com.twitter.elephantbird.mapreduce.output.RCFileThriftOutputFormat;
import com.twitter.elephantbird.pig.piggybank.ThriftBytesToTuple;
import com.twitter.elephantbird.pig.store.RCFileThriftPigStorage;
import com.twitter.elephantbird.pig.util.ThriftToPig;
import com.twitter.elephantbird.thrift.test.TestName;
import com.twitter.elephantbird.thrift.test.TestPerson;
import com.twitter.elephantbird.thrift.test.TestPersonExtended;
import com.twitter.elephantbird.thrift.test.TestPhoneType;
import com.twitter.elephantbird.util.Codecs;
import com.twitter.elephantbird.util.ThriftUtils;

/**
 * Test RCFile loader and storage with Thrift objects
 */
public class TestRCFileThriftStorage {

  private PigServer pigServer;
  private final String testDir =
                    CoreTestUtil.getTestDataDir(TestRCFileThriftStorage.class);
  private final File inputDir = new File(testDir, "in");
  private final File rcfile_in = new File(testDir, "rcfile_in");

  private ThriftToPig<TestPersonExtended> thriftToPig = ThriftToPig.newInstance(TestPersonExtended.class);
  private ThriftConverter<TestPersonExtended> thriftConverter = ThriftConverter.newInstance(TestPersonExtended.class);

  private final TestPersonExtended[] records = new TestPersonExtended[]{
                                                    makePerson(0),
                                                    makePerson(1),
                                                    makePerson(2) };

  private static final Base64 base64 = Codecs.createStandardBase64();

  public static class B64ToTuple extends ThriftBytesToTuple<TestPersonExtended> {
    public B64ToTuple(String className) {
      super(className);
    }

    @Override
    public Tuple exec(Tuple input) throws IOException {
      byte[] bytes = ((DataByteArray)input.get(0)).get();
      input.set(0, new DataByteArray(base64.decode(bytes)));
      return super.exec(input);
    }
  }

  @Before
  public void setUp() throws Exception {

    FileUtil.fullyDelete(new File(testDir));

    pigServer = PigTestUtil.makePigServer();

    pigServer.getPigContext().getProperties().setProperty(
        "mapred.output.compress", "true"); //default codec

    inputDir.mkdirs();

    // create an text file with b64 encoded thrift objects.

    FileOutputStream out = new FileOutputStream(new File(inputDir, "persons_b64.txt"));
    for (TestPersonExtended rec : records) {
      out.write(base64.encode(thriftConverter.toBytes(rec)));
      out.write('\n');
    }
    out.close();
  }

  @Test
  public void testRCFileSThrifttorage() throws Exception {
    /* Create a directory with two files:
     *  - one created with TestPersonExtended objects using RCFileThriftPigStorage
     *  - one created with TestPerson using serialized TestPersonExtended objects
     *         to test handling of unknown fields.
     *
     *  Then load both files using RCFileThriftPigLoader.
     */

    // write to rcFile using RCFileThriftPigStorage
    for(String line : String.format(
            "DEFINE b64ToTuple %s('%s');\n" +
            "A = load '%s' as (line);\n" +
            "A = foreach A generate b64ToTuple(line) as t;\n" +
            "A = foreach A generate FLATTEN(t);\n" +
            "STORE A into '%s' using %s('%s');\n"

            , B64ToTuple.class.getName()
            , TestPersonExtended.class.getName()
            , inputDir.toURI().toString()
            , rcfile_in.toURI().toString()
            , RCFileThriftPigStorage.class.getName()
            , TestPersonExtended.class.getName()

            ).split("\n")) {

      pigServer.registerQuery(line + "\n");
    }
    // the RCFile created above has 5 columns : 4 fields in extended Person
    // and one for unknown fields (this column is empty in this case).

    // store another file with unknowns, by writing the person object with
    // TestPerson rather than with TestPersionExtended, but using
    // serialized TestPersionExtended objects.

    RecordWriter<Writable, Writable> thriftWriter =
      createThriftWriter(TestPerson.class, new File(rcfile_in, "persons_with_unknows.rc"));
    for(TestPersonExtended person : records) {
      // write the bytes from TestPersonExtened
      thriftWriter.write(null, new BytesWritable(thriftConverter.toBytes(person)));
    }
    thriftWriter.close(null);
    // this RCFile has 3 columns : 2 fields in TestPerson and one for unknown
    // fields. The unknowns-columns contains 2 fields from TestPersonExtended
    // that are not understood by TestPerson.

    // load using RCFileThriftPigLoader
    pigServer.registerQuery(String.format(
        "A = load '%s' using %s('%s');\n"
        , rcfile_in.toURI().toString()
        , RCFileThriftPigLoader.class.getName()
        , TestPersonExtended.class.getName()));

    // verify the result:
    Iterator<Tuple> rows = pigServer.openIterator("A");
    for (int i=0; i<2; i++) {
      for(TestPersonExtended person : records) {
        String expected = personToString(person);
        Assert.assertEquals(expected, rows.next().toString());
      }
    }

    // clean up on successful run
    FileUtil.fullyDelete(new File(testDir));
  }

  // return a Person thrift object
  private TestPersonExtended makePerson(int index) {
    return new TestPersonExtended(
              new TestName("bob " + index, "jenkins"),
              ImmutableMap.of(TestPhoneType.HOME,
                              "408-555-5555" + "ex" + index),
              "bob_" + index + "@examle.com",
              new TestName("alice " + index, "smith"));
  }

  @SuppressWarnings("unchecked")
  private static RecordWriter<Writable, Writable>
  createThriftWriter(Class<?> thriftClass, final File file)
                    throws IOException, InterruptedException {

    OutputFormat outputFormat = (
      new RCFileThriftOutputFormat(ThriftUtils.getTypeRef(thriftClass.getName())) {
        @Override
        public Path getDefaultWorkFile(TaskAttemptContext context,
            String extension) throws IOException {
          return new Path(file.toURI().toString());
        }
    });

    Configuration conf = new Configuration();
    // TODO: figure out why Gzip or BZip2 compression fails on OSX
    //conf.setBoolean("mapred.output.compress", true);
    //conf.set("mapred.output.compression.codec", "org.apache.hadoop.io.compress.BZip2Codec");


    return outputFormat.getRecordWriter(
        HadoopCompat.newTaskAttemptContext(conf, new TaskAttemptID()));
  }

  private String personToString(TestPersonExtended person) {
    return thriftToPig.getPigTuple(person).toString();
  }

}