RubixRecordWriter.java

/* (c) 2014 LinkedIn Corp. All rights reserved.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use
 * this file except in compliance with the License. You may obtain a copy of the
 * License at  http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software distributed
 * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied.
 */

package com.linkedin.cubert.io.rubix;

import java.io.ByteArrayOutputStream;
import java.io.DataOutput;
import java.io.DataOutputStream;
import java.io.IOException;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.CompressionOutputStream;
import org.apache.hadoop.io.serializer.SerializationFactory;
import org.apache.hadoop.io.serializer.Serializer;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.pig.data.Tuple;
import org.codehaus.jackson.JsonNode;
import org.codehaus.jackson.map.ObjectMapper;
import org.codehaus.jackson.node.ObjectNode;

import com.linkedin.cubert.block.BlockSchema;
import com.linkedin.cubert.io.BlockSerializationType;
import com.linkedin.cubert.io.CompactSerializer;
import com.linkedin.cubert.plan.physical.CubertStrings;
import com.linkedin.cubert.utils.print;

/**
 * Writes data in rubix binary file format.
 * 
 * The rubix file format stores all values together in the beginning of the file, followed
 * by all keys together at the end of the file. This RecordWriter caches all keys in
 * memory, and dumps them to file when all record values are written. It is, therefore,
 * expected that the cumulative size of keys does not exceed the available memory of the
 * process.
 * <p>
 * The data is organized in the following manner:
 * <p>
 * {@literal File :=  [Value Section] [Key Section] <int = Key Section size>}
 * <p>
 * {@literal Value Section := [Value] [Value] [Value] ...}
 * <p>
 * {@literal Key Section := <string = key class> <string = value class> <BlockSchema of the file>
 * <parition keys> <sort keys> [Key <long=offset of first value in value section> <number
 * of records for this key>]}
 * <p>
 * Each non-null key is considered as an InputSplit. As an example, say, the client
 * emitted the following key-values:
 * <ul>
 * <li>key1, value1</li>
 * 
 * <li>null, value2</li>
 * 
 * <li>null, value3</li>
 * 
 * <li>key2, value4</li>
 * </ul>
 * <p>
 * The file will be stored as:
 * <p>
 * value1 value2 value3 value4 key1 offset_value1 blockid1 output_record_count1 key2
 * offset_value4 blockid2 output_record_count2
 * <p>
 * There are two input splits, one for each non-null key. The the first split, the mapper
 * will receive the following pairs of key-values: (key1, value1), (null, value2), (null,
 * value3). In the second split, the mapper will receive (key2, value4).
 * 
 * @author Maneesh Varshney
 * 
 * @author spyne
 * 
 */
public class RubixRecordWriter<K, V> extends RecordWriter<K, V>
{
    /* Main output stream (connected to output file) */
    private final FSDataOutputStream out;

    /* Serializer for value objects */
    private final Serializer<V> valueSerializer;
    private final Serializer<K> keySerializer;

    /* Used to compress values. Dumps to out stream */
    private final CompressionOutputStream compressedStream;

    private final ByteArrayOutputStream keySectionStream = new ByteArrayOutputStream();
    private final DataOutput keySectionOut = new DataOutputStream(keySectionStream);

    private JsonNode metadataJson;

    private int blockCount = 0;
    private long outputRecordCount = -1;

    @SuppressWarnings({ "unchecked", "rawtypes" })
    public RubixRecordWriter(Configuration conf,
                             FSDataOutputStream out,
                             Class keyClass,
                             Class valueClass,
                             CompressionCodec codec) throws IOException
    {
        this.out = out;

        final SerializationFactory serializationFactory = new SerializationFactory(conf);
        keySerializer = serializationFactory.getSerializer(keyClass);

        ObjectMapper mapper = new ObjectMapper();
        metadataJson =
                mapper.readValue(conf.get(CubertStrings.JSON_METADATA), JsonNode.class);
        ((ObjectNode) metadataJson).put("keyClass", keyClass.getCanonicalName());
        ((ObjectNode) metadataJson).put("valueClass", valueClass.getCanonicalName());
        BlockSchema schema = new BlockSchema(metadataJson.get("schema"));

        if (conf.getBoolean(CubertStrings.USE_COMPACT_SERIALIZATION, false)
                && schema.isFlatSchema())
        {
            valueSerializer = new CompactSerializer<V>(schema);
            ((ObjectNode) metadataJson).put("serializationType",
                                            BlockSerializationType.COMPACT.toString());
        }
        else
        {
            valueSerializer = serializationFactory.getSerializer(valueClass);
            ((ObjectNode) metadataJson).put("serializationType",
                                            BlockSerializationType.DEFAULT.toString());
        }

        keySerializer.open(keySectionStream);

        if (codec == null)
        {
            valueSerializer.open(out);
            compressedStream = null;
        }
        else
        {
            compressedStream = codec.createOutputStream(out);
            valueSerializer.open(compressedStream);
        }

    }

    @Override
    public void write(K key, V value) throws IOException,
            InterruptedException
    {
        if (key != null)
        {
            handleNewKey(key);
        }

        if (value != null)
        {
            outputRecordCount++;
            valueSerializer.serialize(value);
        }
    }

    private void handleNewKey(K key) throws IOException
    {
        /* Increment block count when encountering new key */
        blockCount++;

        if (compressedStream != null)
        {
            compressedStream.finish();
        }

        // write the count for the previous block now;
        if (outputRecordCount >= 0)
        {
            keySectionOut.writeLong(outputRecordCount);
        }
        outputRecordCount = 0;

        @SuppressWarnings("unchecked")
        K partitionKey = (K) ((Tuple) key).get(0);
        long blockId = (Long) ((Tuple) key).get(1);

        keySerializer.serialize(partitionKey);
        keySectionOut.writeLong(out.getPos());
        keySectionOut.writeLong(blockId);

        if (compressedStream != null)
        {
            compressedStream.resetState();
        }
    }

    @Override
    public void close(TaskAttemptContext context) throws IOException,
            InterruptedException
    {
        if (compressedStream != null)
        {
            compressedStream.finish();
        }

        ((ObjectNode) metadataJson).put("numberOfBlocks", blockCount);

        // write out the number of records for the last block
        keySectionOut.writeLong(outputRecordCount);
        byte[] trailerBytes = keySectionStream.toByteArray();

        long trailerstartOffset = out.getPos();

        out.writeUTF(metadataJson.toString());
        out.writeInt(trailerBytes.length);
        out.write(trailerBytes);
        long trailerSize = out.getPos() - trailerstartOffset;
        out.writeLong(trailerstartOffset);
        print.f("Trailer size withe metadata + keys is %d, only keys is %d",
                trailerSize,
                trailerBytes.length);

        out.close();
    }
}