Fix bug vertices has joint primary keys load error

Fix #17 Change-Id: I74ff2445602ff8afb76b58c3e9da5ea68eebeb27
apache · Nov 23, 2018 · b9cdb9c · b9cdb9c
1 parent a0b27a1
commit b9cdb9c
Show file tree

Hide file tree

Showing 7 changed files with 62 additions and 9 deletions.
diff --git a/pom.xml b/pom.xml
@@ -6,7 +6,7 @@
 
     <groupId>com.baidu.hugegraph</groupId>
     <artifactId>hugegraph-loader</artifactId>
-    <version>0.7.0</version>
+    <version>0.8.0</version>
 
     <properties>
         <release.name>hugegraph-loader</release.name>

diff --git 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<String> fieldNames,
                                  Map<String, Object> keyValues) {
         List<String> primaryKeys = vertexLabel.primaryKeys();
-        List<Object> 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
@@ -151,8 +151,8 @@ protected EdgeLabel getEdgeLabel(String name) {
     }
 
     protected String spliceVertexId(VertexLabel vertexLabel,
-                                    List<Object> 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
@@ -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<String, Object> keyValues) {
         } else {
             assert isPrimaryKey(this.vertexLabel.idStrategy());
             List<String> primaryKeys = this.vertexLabel.primaryKeys();
-            List<Object> primaryValues = new ArrayList<>(primaryKeys.size());
+            Object[] primaryValues = new Object[primaryKeys.size()];
             for (Map.Entry<String, Object> entry : keyValues.entrySet()) {
                 String fieldName = entry.getKey();
                 Object fieldValue = entry.getValue();
@@ -95,7 +94,7 @@ private void assignIdIfNeed(Vertex vertex, Map<String, Object> 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
@@ -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<Vertex> 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
@@ -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
@@ -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"
+      }
+    }
+  ]
+}