From b9cdb9c703d3100488b680f3b1cef3a1dd0e9472 Mon Sep 17 00:00:00 2001 From: liningrui Date: Fri, 16 Nov 2018 16:19:14 +0800 Subject: [PATCH] Fix bug vertices has joint primary keys load error Fix #17 Change-Id: I74ff2445602ff8afb76b58c3e9da5ea68eebeb27 --- pom.xml | 2 +- .../hugegraph/loader/parser/EdgeParser.java | 5 ++-- .../loader/parser/ElementParser.java | 4 +-- .../hugegraph/loader/parser/VertexParser.java | 5 ++-- .../loader/test/functional/LoaderTest.java | 27 +++++++++++++++++++ src/test/resources/schema_joint_pk.groovy | 10 +++++++ src/test/resources/struct_joint_pk.json | 18 +++++++++++++ 7 files changed, 62 insertions(+), 9 deletions(-) create mode 100644 src/test/resources/schema_joint_pk.groovy create mode 100644 src/test/resources/struct_joint_pk.json diff --git a/pom.xml b/pom.xml index b9cfa0632..a9fc466b0 100644 --- a/pom.xml +++ b/pom.xml @@ -6,7 +6,7 @@ com.baidu.hugegraph hugegraph-loader - 0.7.0 + 0.8.0 hugegraph-loader diff --git a/src/main/java/com/baidu/hugegraph/loader/parser/EdgeParser.java b/src/main/java/com/baidu/hugegraph/loader/parser/EdgeParser.java index f83965d07..67220c53a 100644 --- a/src/main/java/com/baidu/hugegraph/loader/parser/EdgeParser.java +++ b/src/main/java/com/baidu/hugegraph/loader/parser/EdgeParser.java @@ -19,7 +19,6 @@ package com.baidu.hugegraph.loader.parser; -import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -82,7 +81,7 @@ private Object buildVertexId(VertexLabel vertexLabel, List fieldNames, Map keyValues) { List primaryKeys = vertexLabel.primaryKeys(); - List primaryValues = new ArrayList<>(primaryKeys.size()); + Object[] primaryValues = new Object[primaryKeys.size()]; for (String fieldName : fieldNames) { if (!keyValues.containsKey(fieldName)) { continue; @@ -110,7 +109,7 @@ private Object buildVertexId(VertexLabel vertexLabel, // The id strategy of source/target label must be PRIMARY_KEY if (primaryKeys.contains(key)) { int index = primaryKeys.indexOf(key); - primaryValues.add(index, value); + primaryValues[index] = value; } } } diff --git a/src/main/java/com/baidu/hugegraph/loader/parser/ElementParser.java b/src/main/java/com/baidu/hugegraph/loader/parser/ElementParser.java index a39ea5e38..b28b52ec0 100644 --- a/src/main/java/com/baidu/hugegraph/loader/parser/ElementParser.java +++ b/src/main/java/com/baidu/hugegraph/loader/parser/ElementParser.java @@ -151,8 +151,8 @@ protected EdgeLabel getEdgeLabel(String name) { } protected String spliceVertexId(VertexLabel vertexLabel, - List primaryValues) { - E.checkArgument(vertexLabel.primaryKeys().size() == primaryValues.size(), + Object[] primaryValues) { + E.checkArgument(vertexLabel.primaryKeys().size() == primaryValues.length, "Missing some primary key columns, expect %s, " + "but only got %s for vertex label '%s'", vertexLabel.primaryKeys(), primaryValues, vertexLabel); diff --git a/src/main/java/com/baidu/hugegraph/loader/parser/VertexParser.java b/src/main/java/com/baidu/hugegraph/loader/parser/VertexParser.java index bd2fbc8b6..ba6127ef4 100644 --- a/src/main/java/com/baidu/hugegraph/loader/parser/VertexParser.java +++ b/src/main/java/com/baidu/hugegraph/loader/parser/VertexParser.java @@ -19,7 +19,6 @@ package com.baidu.hugegraph.loader.parser; -import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -85,7 +84,7 @@ private void assignIdIfNeed(Vertex vertex, Map keyValues) { } else { assert isPrimaryKey(this.vertexLabel.idStrategy()); List primaryKeys = this.vertexLabel.primaryKeys(); - List primaryValues = new ArrayList<>(primaryKeys.size()); + Object[] primaryValues = new Object[primaryKeys.size()]; for (Map.Entry entry : keyValues.entrySet()) { String fieldName = entry.getKey(); Object fieldValue = entry.getValue(); @@ -95,7 +94,7 @@ private void assignIdIfNeed(Vertex vertex, Map keyValues) { if (primaryKeys.contains(key)) { int index = primaryKeys.indexOf(key); - primaryValues.add(index, value); + primaryValues[index] = value; } } String id = this.spliceVertexId(this.vertexLabel, primaryValues); diff --git a/src/test/java/com/baidu/hugegraph/loader/test/functional/LoaderTest.java b/src/test/java/com/baidu/hugegraph/loader/test/functional/LoaderTest.java index e7a42bdcf..be0abbf12 100644 --- a/src/test/java/com/baidu/hugegraph/loader/test/functional/LoaderTest.java +++ b/src/test/java/com/baidu/hugegraph/loader/test/functional/LoaderTest.java @@ -509,6 +509,33 @@ public void testLoadVerticesWithCustomizedNumberId() { FileUtil.delete(path("vertex_person_number_id.csv")); } + @Test + public void testLoadVerticesWithJointPrimaryKeys() { + String line = FileUtil.newCSVLine("marko", 29, "Beijing"); + FileUtil.append(path("vertex_person.csv"), line); + + String[] args = new String[]{"-f", path("struct_joint_pk.json"), + "-s", path("schema_joint_pk.groovy"), + "-g", "hugegraph", + "--test-mode", "true"}; + try { + HugeGraphLoader.main(args); + } catch (Exception e) { + Assert.fail("Should not throw exception, but throw " + e); + } + + List vertices = client.graph().listVertices(); + + Assert.assertEquals(1, vertices.size()); + Vertex vertex = vertices.get(0); + + Assert.assertTrue(vertex.id().toString().contains("marko!Beijing")); + Assert.assertEquals("person", vertex.label()); + Assert.assertEquals("marko", vertex.property("name")); + Assert.assertEquals(29, vertex.property("age")); + Assert.assertEquals("Beijing", vertex.property("city")); + } + private static String path(String fileName) { return Paths.get(PATH_PREFIX, fileName).toString(); } diff --git a/src/test/resources/schema_joint_pk.groovy b/src/test/resources/schema_joint_pk.groovy new file mode 100644 index 000000000..5e0616e0a --- /dev/null +++ b/src/test/resources/schema_joint_pk.groovy @@ -0,0 +1,10 @@ +// Define schema +schema.propertyKey("name").asText().ifNotExist().create(); +schema.propertyKey("age").asInt().ifNotExist().create(); +schema.propertyKey("city").asText().ifNotExist().create(); +schema.propertyKey("weight").asDouble().ifNotExist().create(); +schema.propertyKey("date").asText().ifNotExist().create(); + +schema.vertexLabel("person").properties("name", "age", "city").primaryKeys("name", "city").ifNotExist().create(); + +schema.edgeLabel("knows").sourceLabel("person").targetLabel("person").properties("date", "weight").ifNotExist().create(); diff --git a/src/test/resources/struct_joint_pk.json b/src/test/resources/struct_joint_pk.json new file mode 100644 index 000000000..7bb682636 --- /dev/null +++ b/src/test/resources/struct_joint_pk.json @@ -0,0 +1,18 @@ +{ + "vertices": [ + { + "label": "person", + "input": { + "type": "file", + "path": "src/test/resources/vertex_person.csv", + "format": "CSV", + "charset": "UTF-8" + }, + "mapping": { + "name": "name", + "age": "age", + "city": "city" + } + } + ] +}