VertexToFaunusBinary.java

package com.thinkaurelius.faunus.formats.rexster.util;

import com.thinkaurelius.faunus.mapreduce.util.CounterMap;
import com.thinkaurelius.titan.diskstorage.StaticBuffer;
import com.thinkaurelius.titan.diskstorage.util.ByteBufferUtil;
import com.thinkaurelius.titan.graphdb.database.serialize.kryo.KryoSerializer;
import com.tinkerpop.blueprints.Direction;
import com.tinkerpop.blueprints.Edge;
import com.tinkerpop.blueprints.Element;
import com.tinkerpop.blueprints.Vertex;
import org.apache.hadoop.io.WritableUtils;

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

/**
 * @author Marko A. Rodriguez (http://markorodriguez.com)
 */
public class VertexToFaunusBinary {

    private static final ElementIdHandler DEFAULT_ELEMENT_ID_HANDLER = new DefaultElementIdHandler();
    private final ElementIdHandler elementIdHandler;
    protected static final KryoSerializer serialize = new KryoSerializer();

    public VertexToFaunusBinary() {
        this(DEFAULT_ELEMENT_ID_HANDLER);
    }

    public VertexToFaunusBinary(final ElementIdHandler elementIdHandler) {
        this.elementIdHandler = elementIdHandler;
    }

    public static void write(final Vertex vertex, final DataOutput out) throws IOException {
        new VertexToFaunusBinary().writeVertex(vertex, out);
    }

    public static void write(final Vertex vertex, final DataOutput out,
                             final ElementIdHandler elementIdHandler) throws IOException {
        new VertexToFaunusBinary(elementIdHandler).writeVertex(vertex, out);
    }

    public void writeVertex(final Vertex vertex, final DataOutput out) throws IOException {
        WritableUtils.writeVLong(out, elementIdHandler.convertIdentifier(vertex));
        out.writeBoolean(false);
        WritableUtils.writeVLong(out, 0);
        writeProperties(vertex, out);
        writeEdges(vertex, Direction.IN, out);
        writeEdges(vertex, Direction.OUT, out);

    }

    private void writeEdges(final Vertex vertex, final Direction direction, final DataOutput out) throws IOException {
        final CounterMap<String> map = new CounterMap<String>();
        for (final Edge edge : vertex.getEdges(direction)) {
            map.incr(edge.getLabel(), 1);
        }
        WritableUtils.writeVInt(out, map.size());
        for (final Map.Entry<String, Long> entry : map.entrySet()) {
            out.writeUTF(entry.getKey());
            WritableUtils.writeVInt(out, entry.getValue().intValue());
            for (final Edge edge : vertex.getEdges(direction, entry.getKey())) {
                WritableUtils.writeVLong(out, elementIdHandler.convertIdentifier(edge));
                out.writeBoolean(false);
                WritableUtils.writeVLong(out, 0);
                writeProperties(edge, out);
                WritableUtils.writeVLong(out, elementIdHandler.convertIdentifier(edge.getVertex(direction.opposite())));
            }
        }
    }

    private static void writeProperties(final Element element, final DataOutput out) throws IOException {
        WritableUtils.writeVInt(out, element.getPropertyKeys().size());
        if (element.getPropertyKeys().size() > 0) {
            final com.thinkaurelius.titan.graphdb.database.serialize.DataOutput o = serialize.getDataOutput(128, true);
            for (final String key : element.getPropertyKeys()) {
                o.writeObject(key, String.class);
                o.writeClassAndObject(element.getProperty(key));
            }
            final StaticBuffer buffer = o.getStaticBuffer();
            WritableUtils.writeVInt(out, buffer.length());
            out.write(ByteBufferUtil.getArray(buffer.asByteBuffer()));
        }
    }
}