Skip to content

Commit

Permalink
reform loader to adapt hugegraph-computer input
Browse files Browse the repository at this point in the history
  • Loading branch information
coderzc committed Dec 20, 2021
1 parent 5e30894 commit 9e9c794
Show file tree
Hide file tree
Showing 9 changed files with 94 additions and 9 deletions.
2 changes: 1 addition & 1 deletion pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

<groupId>com.baidu.hugegraph</groupId>
<artifactId>hugegraph-loader</artifactId>
<version>0.11.2</version>
<version>0.11.3</version>
<packaging>jar</packaging>

<name>hugegraph-loader</name>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ private EdgeKVPairs newEdgeKVPairs() {
}

@Override
protected SchemaLabel schemaLabel() {
public SchemaLabel schemaLabel() {
return this.edgeLabel;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ public ElementBuilder(LoadContext context, InputStruct struct) {

public abstract List<GE> build(String[] names, Object[] values);

protected abstract SchemaLabel schemaLabel();
public abstract SchemaLabel schemaLabel();

protected abstract Collection<String> nonNullableKeys();

Expand Down Expand Up @@ -549,9 +549,8 @@ public List<Vertex> buildVertices(boolean withProperty) {
}
addProperties(vertex, this.properties);
checkNonNullableKeys(vertex);
} else {
vertex.id(id);
}
vertex.id(id);
return ImmutableList.of(vertex);
}
}
Expand Down Expand Up @@ -646,9 +645,8 @@ public List<Vertex> buildVertices(boolean withProperty) {
addProperty(vertex, this.pkName, pkValue, false);
addProperties(vertex, this.properties);
checkNonNullableKeys(vertex);
} else {
vertex.id(id);
}
vertex.id(id);
vertices.add(vertex);
}
return vertices;
Expand Down
27 changes: 27 additions & 0 deletions src/main/java/com/baidu/hugegraph/loader/builder/SchemaCache.java
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
package com.baidu.hugegraph.loader.builder;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import com.baidu.hugegraph.driver.HugeClient;
Expand All @@ -28,7 +29,11 @@
import com.baidu.hugegraph.structure.schema.EdgeLabel;
import com.baidu.hugegraph.structure.schema.PropertyKey;
import com.baidu.hugegraph.structure.schema.VertexLabel;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;

@JsonIgnoreProperties(ignoreUnknown = true)
public final class SchemaCache {

private final HugeClient client;
Expand All @@ -43,6 +48,28 @@ public SchemaCache(HugeClient client) {
this.edgeLabels = new HashMap<>();
}

@JsonCreator
public SchemaCache(@JsonProperty(value = "propertykeys")
List<PropertyKey> propertyKeyList,
@JsonProperty("vertexlabels")
List<VertexLabel> vertexLabelList,
@JsonProperty("edgelabels")
List<EdgeLabel> edgeLabelList) {
this.client = null;
this.propertyKeys = new HashMap<>();
this.vertexLabels = new HashMap<>();
this.edgeLabels = new HashMap<>();
propertyKeyList.forEach(pk -> {
this.propertyKeys.put(pk.name(), pk);
});
vertexLabelList.forEach(vl -> {
this.vertexLabels.put(vl.name(), vl);
});
edgeLabelList.forEach(el -> {
this.edgeLabels.put(el.name(), el);
});
}

public void updateAll() {
this.propertyKeys.clear();
client.schema().getPropertyKeys().forEach(pk -> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ public List<Vertex> build(String[] names, Object[] values) {
}

@Override
protected SchemaLabel schemaLabel() {
public SchemaLabel schemaLabel() {
return this.vertexLabel;
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
/*
* Copyright 2017 HugeGraph Authors
*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with this
* work for additional information regarding copyright ownership. The ASF
* licenses this file to You under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*/

package com.baidu.hugegraph.loader.executor;

import com.baidu.hugegraph.loader.builder.SchemaCache;

public class ComputerLoadOptions extends LoadOptions {

private final SchemaCache schemaCache;

public ComputerLoadOptions(SchemaCache schemaCache) {
super();
this.schemaCache = schemaCache;
}

public SchemaCache schemaCache() {
return this.schemaCache;
}
}
14 changes: 14 additions & 0 deletions src/main/java/com/baidu/hugegraph/loader/executor/LoadContext.java
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,20 @@ public LoadContext(LoadOptions options) {
this.schemaCache = new SchemaCache(this.client);
}

public LoadContext(ComputerLoadOptions options) {
this.timestamp = DateUtil.now("yyyyMMdd-HHmmss");
this.closed = false;
this.stopped = false;
this.noError = true;
this.options = options;
this.summary = new LoadSummary();
this.oldProgress = LoadProgress.parse(options);
this.newProgress = new LoadProgress();
this.loggers = new ConcurrentHashMap<>();
this.client = null;
this.schemaCache = options.schemaCache();
}

public String timestamp() {
return this.timestamp;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@
import com.beust.jcommander.ParameterException;
import com.google.common.collect.ImmutableSet;

public final class LoadOptions {
public class LoadOptions {

private static final Logger LOG = Log.logger(LoadOptions.class);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@
import com.baidu.hugegraph.loader.metrics.LoadMetrics;
import com.baidu.hugegraph.loader.reader.line.Line;
import com.baidu.hugegraph.structure.GraphElement;
import com.baidu.hugegraph.structure.graph.Vertex;
import com.baidu.hugegraph.structure.schema.VertexLabel;
import com.baidu.hugegraph.util.E;
import com.baidu.hugegraph.util.Log;

Expand Down Expand Up @@ -103,7 +105,15 @@ private ParseTask buildTask(ElementBuilder builder, List<Line> lines) {
batches.add(records);
records = new ArrayList<>(batchSize);
}
final boolean needRemoveId =
builder instanceof VertexBuilder &&
((VertexLabel) builder.schemaLabel())
.idStrategy().isPrimaryKey();

for (GraphElement element : elements) {
if (needRemoveId) {
((Vertex) element).id(null);
}
records.add(new Record(line.rawLine(), element));
count++;
}
Expand Down

0 comments on commit 9e9c794

Please sign in to comment.