diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/io/HugeGraphSONModule.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/io/HugeGraphSONModule.java index 1a982c4a7c..8e214adb03 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/io/HugeGraphSONModule.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/io/HugeGraphSONModule.java @@ -24,10 +24,15 @@ import java.util.Date; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; +import org.apache.tinkerpop.gremlin.process.traversal.Path; +import org.apache.tinkerpop.gremlin.process.traversal.step.util.Tree; +import org.apache.tinkerpop.gremlin.structure.Element; import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONIo; +import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONTokens; import org.apache.tinkerpop.gremlin.structure.io.graphson.TinkerPopJacksonModule; import org.apache.tinkerpop.shaded.jackson.core.JsonGenerator; import org.apache.tinkerpop.shaded.jackson.core.JsonParser; @@ -70,10 +75,10 @@ public class HugeGraphSONModule extends TinkerPopJacksonModule { private static final long serialVersionUID = 6480426922914059122L; - public static boolean OPTIMIZE_SERIALIZE = true; - private static final String TYPE_NAMESPACE = "hugegraph"; + private static boolean OPTIMIZE_SERIALIZE = true; + @SuppressWarnings("rawtypes") private static final Map TYPE_DEFINITIONS; @@ -198,6 +203,9 @@ public static void registerGraphSerializers(SimpleModule module) { */ module.addSerializer(HugeVertex.class, new HugeVertexSerializer()); module.addSerializer(HugeEdge.class, new HugeEdgeSerializer()); + + module.addSerializer(Path.class, new PathSerializer()); + module.addSerializer(Tree.class, new TreeSerializer()); } @SuppressWarnings("rawtypes") @@ -485,6 +493,49 @@ public void serializeWithType(HugeEdge value, JsonGenerator generator, } } + private static class PathSerializer extends StdSerializer { + + public PathSerializer() { + super(Path.class); + } + + @Override + public void serialize(Path path, JsonGenerator jsonGenerator, + SerializerProvider provider) throws IOException { + jsonGenerator.writeStartObject(); + jsonGenerator.writeObjectField(GraphSONTokens.LABELS, + path.labels()); + jsonGenerator.writeObjectField(GraphSONTokens.OBJECTS, + path.objects()); + jsonGenerator.writeEndObject(); + } + } + + @SuppressWarnings("rawtypes") // Tree + private static class TreeSerializer extends StdSerializer { + + public TreeSerializer() { + super(Tree.class); + } + + @Override + public void serialize(Tree tree, JsonGenerator jsonGenerator, + SerializerProvider provider) throws IOException { + jsonGenerator.writeStartArray(); + @SuppressWarnings("unchecked") + Set> set = tree.entrySet(); + for (Map.Entry entry : set) { + jsonGenerator.writeStartObject(); + jsonGenerator.writeObjectField(GraphSONTokens.KEY, + entry.getKey()); + jsonGenerator.writeObjectField(GraphSONTokens.VALUE, + entry.getValue()); + jsonGenerator.writeEndObject(); + } + jsonGenerator.writeEndArray(); + } + } + private static class ShardSerializer extends StdSerializer { public ShardSerializer() { @@ -493,8 +544,7 @@ public ShardSerializer() { @Override public void serialize(Shard shard, JsonGenerator jsonGenerator, - SerializerProvider provider) - throws IOException { + SerializerProvider provider) throws IOException { jsonGenerator.writeStartObject(); jsonGenerator.writeStringField("start", shard.start()); jsonGenerator.writeStringField("end", shard.end()); @@ -511,8 +561,7 @@ public FileSerializer() { @Override public void serialize(File file, JsonGenerator jsonGenerator, - SerializerProvider provider) - throws IOException { + SerializerProvider provider) throws IOException { jsonGenerator.writeStartObject(); jsonGenerator.writeStringField("file", file.getName()); jsonGenerator.writeEndObject(); @@ -527,8 +576,7 @@ public BlobSerializer() { @Override public void serialize(Blob blob, JsonGenerator jsonGenerator, - SerializerProvider provider) - throws IOException { + SerializerProvider provider) throws IOException { jsonGenerator.writeBinary(blob.bytes()); } } diff --git a/hugegraph-test/src/main/java/com/baidu/hugegraph/core/TaskCoreTest.java b/hugegraph-test/src/main/java/com/baidu/hugegraph/core/TaskCoreTest.java index ed4faf2ece..4f8a8d9a0c 100644 --- a/hugegraph-test/src/main/java/com/baidu/hugegraph/core/TaskCoreTest.java +++ b/hugegraph-test/src/main/java/com/baidu/hugegraph/core/TaskCoreTest.java @@ -255,9 +255,9 @@ public void testGremlinJobWithScript() throws TimeoutException { + "schema.propertyKey('lang').asText().ifNotExist().create();" + "schema.propertyKey('date').asDate().ifNotExist().create();" + "schema.propertyKey('price').asInt().ifNotExist().create();" - + "person1=schema.vertexLabel('person1').properties('name','age').ifNotExist().create();" - + "person2=schema.vertexLabel('person2').properties('name','age').ifNotExist().create();" - + "knows=schema.edgeLabel('knows').sourceLabel('person1').targetLabel('person2').properties('date').ifNotExist().create();" + + "schema.vertexLabel('person1').properties('name','age').ifNotExist().create();" + + "schema.vertexLabel('person2').properties('name','age').ifNotExist().create();" + + "schema.edgeLabel('knows').sourceLabel('person1').targetLabel('person2').properties('date').ifNotExist().create();" + "for(int i = 0; i < 1000; i++) {" + " p1=graph.addVertex(T.label,'person1','name','p1-'+i,'age',29);" + " p2=graph.addVertex(T.label,'person2','name','p2-'+i,'age',27);" @@ -302,6 +302,89 @@ public void testGremlinJobWithScript() throws TimeoutException { Assert.assertEquals("[1000]", task.result()); } + @Test + public void testGremlinJobWithSerializedResults() throws TimeoutException { + HugeGraph graph = graph(); + TaskScheduler scheduler = graph.taskScheduler(); + + String script = "schema=graph.schema();" + + "schema.propertyKey('name').asText().ifNotExist().create();" + + "schema.vertexLabel('char').useCustomizeNumberId().properties('name').ifNotExist().create();" + + "schema.edgeLabel('next').sourceLabel('char').targetLabel('char').properties('name').ifNotExist().create();" + + "g.addV('char').property(id,1).property('name','A').as('a')" + + " .addV('char').property(id,2).property('name','B').as('b')" + + " .addV('char').property(id,3).property('name','C').as('c')" + + " .addV('char').property(id,4).property('name','D').as('d')" + + " .addV('char').property(id,5).property('name','E').as('e')" + + " .addV('char').property(id,6).property('name','F').as('f')" + + " .addE('next').from('a').to('b').property('name','ab')" + + " .addE('next').from('b').to('c').property('name','bc')" + + " .addE('next').from('b').to('d').property('name','bd')" + + " .addE('next').from('c').to('d').property('name','cd')" + + " .addE('next').from('c').to('e').property('name','ce')" + + " .addE('next').from('d').to('e').property('name','de')" + + " .addE('next').from('e').to('f').property('name','ef')" + + " .addE('next').from('f').to('d').property('name','fd')" + + " .iterate();" + + "g.tx().commit(); g.E().count();"; + + HugeTask task = runGremlinJob(script); + task = scheduler.waitUntilTaskCompleted(task.id(), 10); + Assert.assertEquals("test-gremlin-job", task.name()); + Assert.assertEquals("gremlin", task.type()); + Assert.assertEquals(TaskStatus.SUCCESS, task.status()); + Assert.assertEquals("[8]", task.result()); + + Id edgeLabelId = graph.schema().getEdgeLabel("next").id(); + + script = "g.V(1).outE().inV().path()"; + task = runGremlinJob(script); + task = scheduler.waitUntilTaskCompleted(task.id(), 10); + Assert.assertEquals(TaskStatus.SUCCESS, task.status()); + String expected = String.format("[{\"labels\":[[],[],[]],\"objects\":[" + + "{\"id\":1,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"A\"}}," + + "{\"id\":\"L1>%s>>L2\",\"label\":\"next\",\"type\":\"edge\",\"outV\":1,\"outVLabel\":\"char\",\"inV\":2,\"inVLabel\":\"char\",\"properties\":{\"name\":\"ab\"}}," + + "{\"id\":2,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"B\"}}" + + "]}]", edgeLabelId); + Assert.assertEquals(expected, task.result()); + + script = "g.V(1).out().out().path()"; + task = runGremlinJob(script); + task = scheduler.waitUntilTaskCompleted(task.id(), 10); + Assert.assertEquals(TaskStatus.SUCCESS, task.status()); + expected = "[{\"labels\":[[],[],[]],\"objects\":[" + + "{\"id\":1,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"A\"}}," + + "{\"id\":2,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"B\"}}," + + "{\"id\":3,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"C\"}}]}," + + "{\"labels\":[[],[],[]],\"objects\":[" + + "{\"id\":1,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"A\"}}," + + "{\"id\":2,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"B\"}}," + + "{\"id\":4,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"D\"}}]}]"; + Assert.assertEquals(expected, task.result()); + + script = "g.V(1).outE().inV().tree()"; + task = runGremlinJob(script); + task = scheduler.waitUntilTaskCompleted(task.id(), 10); + Assert.assertEquals(TaskStatus.SUCCESS, task.status()); + expected = String.format("[[{\"key\":{\"id\":1,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"A\"}}," + + "\"value\":[" + + "{\"key\":{\"id\":\"L1>%s>>L2\",\"label\":\"next\",\"type\":\"edge\",\"outV\":1,\"outVLabel\":\"char\",\"inV\":2,\"inVLabel\":\"char\",\"properties\":{\"name\":\"ab\"}}," + + "\"value\":[{\"key\":{\"id\":2,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"B\"}},\"value\":[]}]}]}]]", + edgeLabelId); + Assert.assertEquals(expected, task.result()); + + script = "g.V(1).out().out().tree()"; + task = runGremlinJob(script); + task = scheduler.waitUntilTaskCompleted(task.id(), 10); + Assert.assertEquals(TaskStatus.SUCCESS, task.status()); + expected = "[[{\"key\":{\"id\":1,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"A\"}}," + + "\"value\":[{\"key\":{\"id\":2,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"B\"}}," + + "\"value\":[" + + "{\"key\":{\"id\":3,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"C\"}},\"value\":[]}," + + "{\"key\":{\"id\":4,\"label\":\"char\",\"type\":\"vertex\",\"properties\":{\"name\":\"D\"}},\"value\":[]}]}]}]]"; + Assert.assertEquals(expected, task.result()); + } + @Test public void testGremlinJobWithFailure() throws TimeoutException { HugeGraph graph = graph(); diff --git a/hugegraph-test/src/main/java/com/baidu/hugegraph/tinkerpop/TestGraph.java b/hugegraph-test/src/main/java/com/baidu/hugegraph/tinkerpop/TestGraph.java index 20c2cb547c..22cf61454b 100644 --- a/hugegraph-test/src/main/java/com/baidu/hugegraph/tinkerpop/TestGraph.java +++ b/hugegraph-test/src/main/java/com/baidu/hugegraph/tinkerpop/TestGraph.java @@ -44,6 +44,7 @@ import com.baidu.hugegraph.schema.PropertyKey; import com.baidu.hugegraph.schema.SchemaManager; import com.baidu.hugegraph.task.TaskScheduler; +import com.baidu.hugegraph.testutil.Whitebox; import com.baidu.hugegraph.type.define.IdStrategy; import com.baidu.hugegraph.type.define.NodeRole; import com.google.common.collect.ImmutableSet; @@ -251,7 +252,8 @@ public GraphComputer compute() throws IllegalArgumentException { @SuppressWarnings({ "unchecked", "rawtypes" }) @Override public I io(final Io.Builder builder) { - HugeGraphSONModule.OPTIMIZE_SERIALIZE = false; + Whitebox.setInternalState(HugeGraphSONModule.class, + "OPTIMIZE_SERIALIZE", false); return (I) builder.graph(this).onMapper(mapper -> mapper.addRegistry(HugeGraphIoRegistry.instance()) ).create();