package com.tinkerpop.gremlin.structure.io.kryo;
import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.tinkerpop.gremlin.process.T;
import com.tinkerpop.gremlin.structure.Direction;
import com.tinkerpop.gremlin.structure.Edge;
import com.tinkerpop.gremlin.structure.Graph;
import com.tinkerpop.gremlin.structure.Vertex;
import com.tinkerpop.gremlin.structure.VertexProperty;
import com.tinkerpop.gremlin.structure.io.GraphReader;
import com.tinkerpop.gremlin.structure.util.batch.BatchGraph;
import com.tinkerpop.gremlin.structure.util.detached.DetachedEdge;
import com.tinkerpop.gremlin.structure.util.detached.DetachedVertex;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
/**
* The {@link GraphReader} for the Gremlin Structure serialization format based on Kryo. The format is meant to be
* non-lossy in terms of Gremlin Structure to Gremlin Structure migrations (assuming both structure implementations
* support the same graph features).
* <br/>
* This implementation is not thread-safe.
*
* @author Stephen Mallette (http://stephen.genoprime.com)
*/
public class KryoReader implements GraphReader {
private final Kryo kryo;
private final GremlinKryo.HeaderReader headerReader;
private final long batchSize;
private final String vertexIdKey;
private final String edgeIdKey;
private final File tempFile;
final AtomicLong counter = new AtomicLong(0);
private KryoReader(final File tempFile, final long batchSize,
final String vertexIdKey, final String edgeIdKey,
final GremlinKryo gremlinKryo) {
this.kryo = gremlinKryo.createKryo();
this.headerReader = gremlinKryo.getHeaderReader();
this.vertexIdKey = vertexIdKey;
this.edgeIdKey = edgeIdKey;
this.tempFile = tempFile;
this.batchSize = batchSize;
}
@Override
public Iterator<Vertex> readVertices(final InputStream inputStream, final Direction direction,
final Function<DetachedVertex, Vertex> vertexMaker,
final Function<DetachedEdge, Edge> edgeMaker) throws IOException {
final Input input = new Input(inputStream);
return new Iterator<Vertex>() {
@Override
public boolean hasNext() {
return !input.eof();
}
@Override
public Vertex next() {
try {
final Vertex v = readVertex(direction, vertexMaker, edgeMaker, input);
// read the vertex terminator
kryo.readClassAndObject(input);
return v;
} catch (Exception ex) {
throw new RuntimeException(ex);
}
}
};
}
@Override
public Edge readEdge(final InputStream inputStream, final Function<DetachedEdge, Edge> edgeMaker) throws IOException {
final Input input = new Input(inputStream);
this.headerReader.read(kryo, input);
final Object o = kryo.readClassAndObject(input);
return edgeMaker.apply((DetachedEdge) o);
}
@Override
public Vertex readVertex(final InputStream inputStream, final Function<DetachedVertex, Vertex> vertexMaker) throws IOException {
return readVertex(inputStream, null, vertexMaker, null);
}
@Override
public Vertex readVertex(final InputStream inputStream, final Direction direction, Function<DetachedVertex, Vertex> vertexMaker, final Function<DetachedEdge, Edge> edgeMaker) throws IOException {
final Input input = new Input(inputStream);
return readVertex(direction, vertexMaker, edgeMaker, input);
}
@Override
public void readGraph(final InputStream inputStream, final Graph graphToWriteTo) throws IOException {
this.counter.set(0);
final Input input = new Input(inputStream);
this.headerReader.read(kryo, input);
final BatchGraph graph;
try {
// will throw an exception if not constructed properly
graph = BatchGraph.build(graphToWriteTo)
.vertexIdKey(vertexIdKey)
.edgeIdKey(edgeIdKey)
.bufferSize(batchSize).create();
} catch (Exception ex) {
throw new IOException("Could not instantiate BatchGraph wrapper", ex);
}
try (final Output output = new Output(new FileOutputStream(tempFile))) {
final boolean supportedMemory = input.readBoolean();
if (supportedMemory) {
// if the graph that serialized the data supported sideEffects then the sideEffects needs to be read
// to advance the reader forward. if the graph being read into doesn't support the sideEffects
// then we just setting the data to sideEffects.
final Map<String, Object> memMap = (Map<String, Object>) kryo.readObject(input, HashMap.class);
if (graphToWriteTo.features().graph().variables().supportsVariables()) {
final Graph.Variables variables = graphToWriteTo.variables();
memMap.forEach(variables::set);
}
}
final boolean hasSomeVertices = input.readBoolean();
if (hasSomeVertices) {
while (!input.eof()) {
final List<Object> vertexArgs = new ArrayList<>();
final DetachedVertex current = (DetachedVertex) kryo.readClassAndObject(input);
appendToArgList(vertexArgs, T.id, current.id());
appendToArgList(vertexArgs, T.label, current.label());
final Vertex v = graph.addVertex(vertexArgs.toArray());
current.iterators().propertyIterator().forEachRemaining(p -> createVertexProperty(graphToWriteTo, v, p, false));
current.iterators().hiddenPropertyIterator().forEachRemaining(p -> createVertexProperty(graphToWriteTo, v, p, true));
// the gio file should have been written with a direction specified
final boolean hasDirectionSpecified = input.readBoolean();
final Direction directionInStream = kryo.readObject(input, Direction.class);
final Direction directionOfEdgeBatch = kryo.readObject(input, Direction.class);
// graph serialization requires that a direction be specified in the stream and that the
// direction of the edges be OUT
if (!hasDirectionSpecified || directionInStream != Direction.OUT || directionOfEdgeBatch != Direction.OUT)
throw new IllegalStateException(String.format("Stream must specify edge direction and that direction must be %s", Direction.OUT));
// if there are edges then read them to end and write to temp, otherwise read what should be
// the vertex terminator
if (!input.readBoolean())
kryo.readClassAndObject(input);
else
readToEndOfEdgesAndWriteToTemp(input, output);
}
}
} catch (Exception ex) {
throw new IOException(ex);
}
// done writing to temp
// start reading in the edges now from the temp file
try (final Input edgeInput = new Input(new FileInputStream(tempFile))) {
readFromTempEdges(edgeInput, graph);
graph.tx().commit();
} catch (Exception ex) {
ex.printStackTrace();
throw new IOException(ex);
} finally {
deleteTempFileSilently();
}
}
private static void createVertexProperty(final Graph graphToWriteTo, final Vertex v, final VertexProperty<Object> p, final boolean hidden) {
final List<Object> propertyArgs = new ArrayList<>();
if (graphToWriteTo.features().vertex().properties().supportsUserSuppliedIds())
appendToArgList(propertyArgs, T.id, p.id());
p.iterators().propertyIterator().forEachRemaining(it -> appendToArgList(propertyArgs, it.key(), it.value()));
p.iterators().hiddenPropertyIterator().forEachRemaining(it -> appendToArgList(propertyArgs, Graph.Key.hide(it.key()), it.value()));
v.property(hidden ? Graph.Key.hide(p.key()) : p.key(), p.value(), propertyArgs.toArray());
}
private static void appendToArgList(final List<Object> propertyArgs, final Object key, final Object val) {
propertyArgs.add(key);
propertyArgs.add(val);
}
private Vertex readVertex(final Direction directionRequested, final Function<DetachedVertex, Vertex> vertexMaker,
final Function<DetachedEdge, Edge> edgeMaker, final Input input) throws IOException {
if (null != directionRequested && null == edgeMaker)
throw new IllegalArgumentException("If a directionRequested is specified then an edgeAdder function should also be specified");
this.headerReader.read(kryo, input);
final DetachedVertex detachedVertex = (DetachedVertex) kryo.readClassAndObject(input);
final Vertex v = vertexMaker.apply(detachedVertex);
final boolean streamContainsEdgesInSomeDirection = input.readBoolean();
if (!streamContainsEdgesInSomeDirection && directionRequested != null)
throw new IllegalStateException(String.format("The direction %s was requested but no attempt was made to serialize edges into this stream", directionRequested));
// if there are edges in the stream and the direction is not present then the rest of the stream is
// simply ignored
if (directionRequested != null) {
final Direction directionsInStream = kryo.readObject(input, Direction.class);
if (directionsInStream != Direction.BOTH && directionsInStream != directionRequested)
throw new IllegalStateException(String.format("Stream contains %s edges, but requesting %s", directionsInStream, directionRequested));
final Direction firstDirection = kryo.readObject(input, Direction.class);
if (firstDirection == Direction.OUT && (directionRequested == Direction.BOTH || directionRequested == Direction.OUT))
readEdges(input, edgeMaker);
else {
// requested direction in, but BOTH must be serialized so skip this. the illegalstateexception
// prior to this IF should have caught a problem where IN is not supported at all
if (firstDirection == Direction.OUT && directionRequested == Direction.IN)
skipEdges(input);
}
if (directionRequested == Direction.BOTH || directionRequested == Direction.IN) {
// if the first direction was OUT then it was either read or skipped. in that case, the marker
// of the stream is currently ready to read the IN direction. otherwise it's in the perfect place
// to start reading edges
if (firstDirection == Direction.OUT)
kryo.readObject(input, Direction.class);
readEdges(input, edgeMaker);
}
}
return v;
}
private void readEdges(final Input input, final Function<DetachedEdge, Edge> edgeMaker) {
if (input.readBoolean()) {
Object next = kryo.readClassAndObject(input);
while (!next.equals(EdgeTerminator.INSTANCE)) {
final DetachedEdge detachedEdge = (DetachedEdge) next;
edgeMaker.apply(detachedEdge);
next = kryo.readClassAndObject(input);
}
}
}
private void skipEdges(final Input input) {
if (input.readBoolean()) {
Object next = kryo.readClassAndObject(input);
while (!next.equals(EdgeTerminator.INSTANCE)) {
// next edge to skip or the terminator
next = kryo.readClassAndObject(input);
}
}
}
/**
* Reads through the all the edges for a vertex and writes the edges to a temp file which will be read later.
*/
private void readToEndOfEdgesAndWriteToTemp(final Input input, final Output output) throws IOException {
Object next = kryo.readClassAndObject(input);
while (!next.equals(EdgeTerminator.INSTANCE)) {
kryo.writeClassAndObject(output, next);
// next edge or terminator
next = kryo.readClassAndObject(input);
}
// this should be the vertex terminator
kryo.readClassAndObject(input);
kryo.writeClassAndObject(output, EdgeTerminator.INSTANCE);
kryo.writeClassAndObject(output, VertexTerminator.INSTANCE);
}
/**
* Read the edges from the temp file and load them to the graph.
*/
private void readFromTempEdges(final Input input, final Graph graphToWriteTo) {
while (!input.eof()) {
// in this case the outId is the id assigned by the graph
Object next = kryo.readClassAndObject(input);
while (!next.equals(EdgeTerminator.INSTANCE)) {
final List<Object> edgeArgs = new ArrayList<>();
final DetachedEdge detachedEdge = (DetachedEdge) next;
final Vertex vOut = graphToWriteTo.v(detachedEdge.iterators().vertexIterator(Direction.OUT).next().id());
final Vertex inV = graphToWriteTo.v(detachedEdge.iterators().vertexIterator(Direction.IN).next().id());
detachedEdge.iterators().propertyIterator().forEachRemaining(p -> edgeArgs.addAll(Arrays.asList(p.key(), p.value())));
detachedEdge.iterators().hiddenPropertyIterator().forEachRemaining(p -> edgeArgs.addAll(Arrays.asList(Graph.Key.hide(p.key()), p.value())));
appendToArgList(edgeArgs, T.id, detachedEdge.id());
vOut.addEdge(detachedEdge.label(), inV, edgeArgs.toArray());
next = kryo.readClassAndObject(input);
}
// vertex terminator
kryo.readClassAndObject(input);
}
}
@SuppressWarnings("ResultOfMethodCallIgnored")
private void deleteTempFileSilently() {
try {
tempFile.delete();
} catch (Exception ignored) {
}
}
public static Builder build() {
return new Builder();
}
public static class Builder {
private File tempFile;
private long batchSize = BatchGraph.DEFAULT_BUFFER_SIZE;
private String vertexIdKey = T.id.getAccessor();
private String edgeIdKey = T.id.getAccessor();
/**
* Always use the most recent kryo version by default
*/
private GremlinKryo gremlinKryo = GremlinKryo.build().create();
private Builder() {
this.tempFile = new File(UUID.randomUUID() + ".tmp");
}
/**
* Set the size between commits when reading into the {@link Graph} instance. This value defaults to
* {@link BatchGraph#DEFAULT_BUFFER_SIZE}.
*/
public Builder batchSize(final long batchSize) {
this.batchSize = batchSize;
return this;
}
/**
* Supply a custom {@link GremlinKryo} instance to use as the serializer for the {@code KryoWriter}.
*/
public Builder custom(final GremlinKryo gremlinKryo) {
this.gremlinKryo = gremlinKryo;
return this;
}
/**
* The name of the key to supply to
* {@link com.tinkerpop.gremlin.structure.util.batch.BatchGraph.Builder#vertexIdKey} when reading data into
* the {@link Graph}.
*/
public Builder vertexIdKey(final String vertexIdKey) {
this.vertexIdKey = vertexIdKey;
return this;
}
/**
* The name of the key to supply to
* {@link com.tinkerpop.gremlin.structure.util.batch.BatchGraph.Builder#edgeIdKey} when reading data into
* the {@link Graph}.
*/
public Builder edgeIdKey(final String edgeIdKey) {
this.edgeIdKey = edgeIdKey;
return this;
}
/**
* The reader requires a working directory to write temp files to. If this value is not set, it will write
* the temp file to the local directory.
*/
public Builder workingDirectory(final String workingDirectory) {
final File f = new File(workingDirectory);
if (!f.exists() || !f.isDirectory())
throw new IllegalArgumentException("The workingDirectory is not a directory or does not exist");
tempFile = new File(workingDirectory + File.separator + UUID.randomUUID() + ".tmp");
return this;
}
public KryoReader create() {
return new KryoReader(tempFile, batchSize, this.vertexIdKey, this.edgeIdKey, this.gremlinKryo);
}
}
}