This is an automated email from the ASF dual-hosted git repository.
jin pushed a commit to branch master
in repository
https://gitbox.apache.org/repos/asf/incubator-hugegraph-toolchain.git
The following commit(s) were added to refs/heads/master by this push:
new 9ee7abb feat: introduce HugeGraphSparkLoader (#282)
9ee7abb is described below
commit 9ee7abb2ee397db636e4c1a1f258f0c7bcd1b5c1
Author: Simon Cheung <[email protected]>
AuthorDate: Fri May 27 11:27:08 2022 +0800
feat: introduce HugeGraphSparkLoader (#282)
---
.../assembly/static/bin/hugegraph-spark-loader.sh | 64 ++++++
hugegraph-loader/pom.xml | 11 +
.../hugegraph/loader/executor/LoadContext.java | 3 +-
.../hugegraph/loader/executor/LoadOptions.java | 3 +-
.../hugegraph/loader/mapping/ElementMapping.java | 10 +-
.../hugegraph/loader/mapping/InputStruct.java | 3 +-
.../hugegraph/loader/source/AbstractSource.java | 3 +-
.../hugegraph/loader/source/file/FileFilter.java | 3 +-
.../hugegraph/loader/source/file/SkippedLine.java | 3 +-
.../loader/spark/HugeGraphSparkLoader.java | 245 +++++++++++++++++++++
.../loader/test/unit/MappingConverterTest.java | 4 +-
11 files changed, 343 insertions(+), 9 deletions(-)
diff --git a/hugegraph-loader/assembly/static/bin/hugegraph-spark-loader.sh
b/hugegraph-loader/assembly/static/bin/hugegraph-spark-loader.sh
new file mode 100755
index 0000000..a2e6a06
--- /dev/null
+++ b/hugegraph-loader/assembly/static/bin/hugegraph-spark-loader.sh
@@ -0,0 +1,64 @@
+#!/bin/bash
+
+PARAMS=""
+while (( "$#" )); do
+ case "$1" in
+ -m|--master)
+ MASTER=$2
+ shift 2
+ ;;
+
+ -n|--name)
+ APP_NAME=$2
+ shift 2
+ ;;
+
+ -e|--deploy-mode)
+ DEPLOY_MODE=$2
+ shift 2
+ ;;
+
+ -c|--conf)
+ SPARK_CONFIG=${SPARK_CONFIG}" --conf "$2
+ shift 2
+ ;;
+
+ --) # end argument parsing
+ shift
+ break
+ ;;
+
+ *) # preserve positional arguments
+ PARAMS="$PARAMS $1"
+ shift
+ ;;
+
+ esac
+done
+
+if [ -z ${MASTER} ] || [ -z ${DEPLOY_MODE} ]; then
+ echo "Error: The following options are required:
+ [-e | --deploy-mode], [-m | --master]"
+ usage
+ exit 0
+fi
+
+BIN_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+APP_DIR=$(dirname ${BIN_DIR})
+LIB_DIR=${APP_DIR}/lib
+
+assemblyJarName=$(find ${LIB_DIR} -name hugegraph-loader*.jar)
+
+DEFAULT_APP_NAME="hugegraph-spark-loader"
+APP_NAME=${APP_NAME:-$DEFAULT_APP_NAME}
+
+CMD="${SPARK_HOME}/bin/spark-submit
+ --name ${APP_NAME} \
+ --master ${MASTER} \
+ --deploy-mode ${DEPLOY_MODE} \
+ --class com.baidu.hugegraph.loader.spark.HugeGraphSparkLoader \
+ ${SPARK_CONFIG}
+ --jars $(echo ${LIB_DIR}/*.jar | tr ' ' ',') ${assemblyJarName} ${PARAMS}"
+
+echo ${CMD}
+exec ${CMD}
diff --git a/hugegraph-loader/pom.xml b/hugegraph-loader/pom.xml
index 4ffdd52..afcb9d6 100644
--- a/hugegraph-loader/pom.xml
+++ b/hugegraph-loader/pom.xml
@@ -36,6 +36,17 @@
</properties>
<dependencies>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-core_2.12</artifactId>
+ <version>3.1.2</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_2.12</artifactId>
+ <version>3.1.2</version>
+ </dependency>
+
<dependency>
<groupId>com.baidu.hugegraph</groupId>
<artifactId>hugegraph-client</artifactId>
diff --git
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/executor/LoadContext.java
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/executor/LoadContext.java
index 24a8df6..117124d 100644
---
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/executor/LoadContext.java
+++
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/executor/LoadContext.java
@@ -20,6 +20,7 @@
package com.baidu.hugegraph.loader.executor;
import java.io.IOException;
+import java.io.Serializable;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
@@ -38,7 +39,7 @@ import com.baidu.hugegraph.loader.util.HugeClientHolder;
import com.baidu.hugegraph.structure.constant.GraphMode;
import com.baidu.hugegraph.util.Log;
-public final class LoadContext {
+public final class LoadContext implements Serializable {
private static final Logger LOG = Log.logger(LoadContext.class);
diff --git
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/executor/LoadOptions.java
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/executor/LoadOptions.java
index c35b5a6..c5f1bfd 100644
---
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/executor/LoadOptions.java
+++
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/executor/LoadOptions.java
@@ -20,6 +20,7 @@
package com.baidu.hugegraph.loader.executor;
import java.io.File;
+import java.io.Serializable;
import java.util.Set;
import org.apache.commons.lang3.StringUtils;
@@ -35,7 +36,7 @@ import com.beust.jcommander.Parameter;
import com.beust.jcommander.ParameterException;
import com.google.common.collect.ImmutableSet;
-public class LoadOptions {
+public class LoadOptions implements Serializable {
private static final Logger LOG = Log.logger(LoadOptions.class);
diff --git
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/mapping/ElementMapping.java
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/mapping/ElementMapping.java
index fbdeac2..48aaec6 100644
---
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/mapping/ElementMapping.java
+++
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/mapping/ElementMapping.java
@@ -19,6 +19,7 @@
package com.baidu.hugegraph.loader.mapping;
+import java.io.Serializable;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
@@ -37,7 +38,7 @@ import com.fasterxml.jackson.annotation.JsonPropertyOrder;
import com.google.common.collect.ImmutableSet;
@JsonPropertyOrder({"label", "skip"})
-public abstract class ElementMapping implements Checkable {
+public abstract class ElementMapping implements Checkable, Serializable {
@JsonProperty("label")
private String label;
@@ -55,6 +56,8 @@ public abstract class ElementMapping implements Checkable {
private Set<Object> nullValues;
@JsonProperty("update_strategies")
private Map<String, UpdateStrategy> updateStrategies;
+ @JsonProperty("batch_size")
+ private long batchSize;
public ElementMapping() {
this.skip = false;
@@ -64,6 +67,7 @@ public abstract class ElementMapping implements Checkable {
this.ignoredFields = new HashSet<>();
this.nullValues = ImmutableSet.of(Constants.EMPTY_STR);
this.updateStrategies = new HashMap<>();
+ this.batchSize = 1000;
}
public abstract ElemType type();
@@ -173,6 +177,10 @@ public abstract class ElementMapping implements Checkable {
return mappingValue;
}
+ public long batchSize() {
+ return this.batchSize;
+ }
+
public Set<String> selectedFields() {
return this.selectedFields;
}
diff --git
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/mapping/InputStruct.java
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/mapping/InputStruct.java
index 0b97974..850015c 100644
---
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/mapping/InputStruct.java
+++
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/mapping/InputStruct.java
@@ -19,6 +19,7 @@
package com.baidu.hugegraph.loader.mapping;
+import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;
@@ -33,7 +34,7 @@ import com.fasterxml.jackson.annotation.JsonPropertyOrder;
import com.google.common.collect.ImmutableList;
@JsonPropertyOrder({"id", "skip", "input", "vertices", "edges"})
-public class InputStruct implements Checkable {
+public class InputStruct implements Checkable, Serializable {
public static final InputStruct EMPTY = new InputStruct(ImmutableList.of(),
ImmutableList.of());
diff --git
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/AbstractSource.java
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/AbstractSource.java
index c264f83..518831a 100644
---
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/AbstractSource.java
+++
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/AbstractSource.java
@@ -19,6 +19,7 @@
package com.baidu.hugegraph.loader.source;
+import java.io.Serializable;
import java.nio.charset.Charset;
import java.util.Arrays;
import java.util.List;
@@ -29,7 +30,7 @@ import com.baidu.hugegraph.util.CollectionUtil;
import com.baidu.hugegraph.util.E;
import com.fasterxml.jackson.annotation.JsonProperty;
-public abstract class AbstractSource implements InputSource {
+public abstract class AbstractSource implements InputSource, Serializable {
@JsonProperty("header")
private String[] header;
diff --git
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/file/FileFilter.java
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/file/FileFilter.java
index a190f7c..5e9acc2 100644
---
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/file/FileFilter.java
+++
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/file/FileFilter.java
@@ -19,6 +19,7 @@
package com.baidu.hugegraph.loader.source.file;
+import java.io.Serializable;
import java.util.Set;
import org.apache.commons.io.FilenameUtils;
@@ -26,7 +27,7 @@ import org.apache.commons.io.FilenameUtils;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.collect.ImmutableSet;
-public class FileFilter {
+public class FileFilter implements Serializable {
private static final String ALL_EXTENSION = "*";
diff --git
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/file/SkippedLine.java
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/file/SkippedLine.java
index 24a62da..ad16ab8 100644
---
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/file/SkippedLine.java
+++
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/source/file/SkippedLine.java
@@ -19,13 +19,14 @@
package com.baidu.hugegraph.loader.source.file;
+import java.io.Serializable;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import com.baidu.hugegraph.loader.constant.Constants;
import com.fasterxml.jackson.annotation.JsonProperty;
-public class SkippedLine {
+public class SkippedLine implements Serializable {
@JsonProperty("regex")
private String regex;
diff --git
a/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/spark/HugeGraphSparkLoader.java
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/spark/HugeGraphSparkLoader.java
new file mode 100644
index 0000000..5d552e2
--- /dev/null
+++
b/hugegraph-loader/src/main/java/com/baidu/hugegraph/loader/spark/HugeGraphSparkLoader.java
@@ -0,0 +1,245 @@
+/*
+ * 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.spark;
+
+import com.baidu.hugegraph.driver.GraphManager;
+import com.baidu.hugegraph.loader.builder.EdgeBuilder;
+import com.baidu.hugegraph.loader.builder.ElementBuilder;
+import com.baidu.hugegraph.loader.builder.VertexBuilder;
+import com.baidu.hugegraph.loader.executor.LoadContext;
+import com.baidu.hugegraph.loader.executor.LoadOptions;
+import com.baidu.hugegraph.loader.mapping.EdgeMapping;
+import com.baidu.hugegraph.loader.mapping.VertexMapping;
+import com.baidu.hugegraph.loader.mapping.ElementMapping;
+import com.baidu.hugegraph.loader.mapping.InputStruct;
+import com.baidu.hugegraph.loader.mapping.LoadMapping;
+import com.baidu.hugegraph.loader.source.InputSource;
+import com.baidu.hugegraph.loader.source.file.FileFilter;
+import com.baidu.hugegraph.loader.source.file.FileFormat;
+import com.baidu.hugegraph.loader.source.file.FileSource;
+import com.baidu.hugegraph.loader.source.file.SkippedLine;
+import com.baidu.hugegraph.loader.source.file.Compression;
+import com.baidu.hugegraph.loader.util.Printer;
+import com.baidu.hugegraph.structure.GraphElement;
+import com.baidu.hugegraph.structure.graph.UpdateStrategy;
+import com.baidu.hugegraph.structure.graph.Vertex;
+import com.baidu.hugegraph.structure.graph.Edge;
+import com.baidu.hugegraph.structure.graph.BatchEdgeRequest;
+import com.baidu.hugegraph.structure.graph.BatchVertexRequest;
+import com.baidu.hugegraph.util.Log;
+
+import org.apache.spark.sql.DataFrameReader;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSession;
+import org.slf4j.Logger;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+public class HugeGraphSparkLoader implements Serializable {
+
+ public static final Logger LOG = Log.logger(HugeGraphSparkLoader.class);
+
+ private final LoadOptions loadOptions;
+ private final Map<ElementBuilder, List<GraphElement>> builders;
+
+ public static void main(String[] args) {
+ HugeGraphSparkLoader loader;
+ try {
+ loader = new HugeGraphSparkLoader(args);
+ } catch (Throwable e) {
+ Printer.printError("Failed to start loading", e);
+ return;
+ }
+ loader.load();
+ }
+
+ public HugeGraphSparkLoader(String[] args) {
+ this.loadOptions = LoadOptions.parseOptions(args);
+ this.builders = new HashMap<>();
+ }
+
+ public void load() {
+ LoadMapping mapping = LoadMapping.of(this.loadOptions.file);
+ List<InputStruct> structs = mapping.structs();
+
+ SparkSession session = SparkSession.builder().getOrCreate();
+ for (InputStruct struct : structs) {
+ Dataset<Row> ds = read(session, struct);
+ ds.foreachPartition((Iterator<Row> p) -> {
+ LoadContext context = initPartition(this.loadOptions, struct);
+ p.forEachRemaining((Row row) -> {
+ loadRow(struct, row, p, context);
+ });
+ context.close();
+ });
+ }
+ session.close();
+ session.stop();
+ }
+
+ private LoadContext initPartition(
+ LoadOptions loadOptions, InputStruct struct) {
+ LoadContext context = new LoadContext(loadOptions);
+ for (VertexMapping vertexMapping : struct.vertices()) {
+ this.builders.put(
+ new VertexBuilder(context, struct, vertexMapping),
+ new ArrayList<>());
+ }
+ for (EdgeMapping edgeMapping : struct.edges()) {
+ this.builders.put(new EdgeBuilder(context, struct, edgeMapping),
+ new ArrayList<>());
+ }
+ return context;
+ }
+
+ private void loadRow(InputStruct struct, Row row, Iterator<Row> p,
+ LoadContext context) {
+ for (Map.Entry<ElementBuilder, List<GraphElement>> builderMap :
+ this.builders.entrySet()) {
+ ElementMapping elementMapping = builderMap.getKey().mapping();
+ // Parse
+ if (elementMapping.skip()) {
+ continue;
+ }
+ parse(row, builderMap, struct);
+
+ // Insert
+ List<GraphElement> graphElements =
+ builderMap.getValue();
+ if (graphElements.size() > elementMapping.batchSize() ||
+ (!p.hasNext() && graphElements.size() > 0)) {
+ sink(builderMap, context.client().graph(),
+ this.loadOptions.checkVertex);
+ }
+ }
+ }
+
+ private Dataset<Row> read(SparkSession ss, InputStruct struct) {
+ InputSource input = struct.input();
+ String charset = input.charset();
+ FileSource fileSource = input.asFileSource();
+
+ String[] header = fileSource.header();
+ String delimiter = fileSource.delimiter();
+ String path = fileSource.path();
+ FileFilter filter = fileSource.filter();
+ FileFormat format = fileSource.format();
+ String dateFormat = fileSource.dateFormat();
+ String timeZone = fileSource.timeZone();
+ SkippedLine skippedLine = fileSource.skippedLine();
+ Compression compression = fileSource.compression();
+ int batchSize = fileSource.batchSize();
+
+ DataFrameReader reader = ss.read();
+ Dataset<Row> ds;
+ switch (input.type()) {
+ case FILE:
+ case HDFS:
+ switch (format) {
+ case TEXT:
+ ds = reader.text(path);
+ break;
+ case JSON:
+ ds = reader.json(path);
+ break;
+ case CSV:
+ ds = reader.csv(path);
+ break;
+ default:
+ throw new IllegalStateException(
+ "Unexpected format value: " + format);
+ }
+ break;
+ case JDBC:
+ // TODO: implement jdbc
+ default:
+ throw new AssertionError(String.format(
+ "Unsupported input source '%s'", input.type()));
+ }
+ return ds;
+ }
+
+ private void parse(Row row,
+ Map.Entry<ElementBuilder, List<GraphElement>>
builderMap,
+ InputStruct struct) {
+ ElementBuilder builder = builderMap.getKey();
+ List<GraphElement> graphElements = builderMap.getValue();
+ if ("".equals(row.mkString())) {
+ return;
+ }
+ List<GraphElement> elements;
+ switch (struct.input().type()) {
+ case FILE:
+ case HDFS:
+ FileSource fileSource = struct.input().asFileSource();
+ elements = builder.build(fileSource.header(),
+ row.mkString()
+ .split(fileSource.delimiter()));
+ break;
+ case JDBC:
+ //TODO: implement jdbc
+ default:
+ throw new AssertionError(String.format(
+ "Unsupported input source '%s'",
+ struct.input().type()));
+ }
+ graphElements.addAll(elements);
+ }
+
+ private void sink(Map.Entry<ElementBuilder, List<GraphElement>> builderMap,
+ GraphManager g, boolean isCheckVertex) {
+ ElementBuilder builder = builderMap.getKey();
+ ElementMapping elementMapping = builder.mapping();
+ List<GraphElement> graphElements = builderMap.getValue();
+ boolean isVertex = builder.mapping().type().isVertex();
+ Map<String, UpdateStrategy> updateStrategyMap =
elementMapping.updateStrategies();
+ if (updateStrategyMap.isEmpty()) {
+ if (isVertex) {
+ g.addVertices((List<Vertex>) (Object) graphElements);
+ } else {
+ g.addEdges((List<Edge>) (Object) graphElements);
+ }
+ } else {
+ // CreateIfNotExist dose not support false now
+ if (isVertex) {
+ BatchVertexRequest.Builder req =
+ new BatchVertexRequest.Builder();
+ req.vertices((List<Vertex>) (Object) graphElements)
+ .updatingStrategies(updateStrategyMap)
+ .createIfNotExist(true);
+ g.updateVertices(req.build());
+ } else {
+ BatchEdgeRequest.Builder req = new BatchEdgeRequest.Builder();
+ req.edges((List<Edge>) (Object) graphElements)
+ .updatingStrategies(updateStrategyMap)
+ .checkVertex(isCheckVertex)
+ .createIfNotExist(true);
+ g.updateEdges(req.build());
+ }
+ }
+ graphElements.clear();
+ }
+}
diff --git
a/hugegraph-loader/src/test/java/com/baidu/hugegraph/loader/test/unit/MappingConverterTest.java
b/hugegraph-loader/src/test/java/com/baidu/hugegraph/loader/test/unit/MappingConverterTest.java
index e09dfaf..459b3b5 100644
---
a/hugegraph-loader/src/test/java/com/baidu/hugegraph/loader/test/unit/MappingConverterTest.java
+++
b/hugegraph-loader/src/test/java/com/baidu/hugegraph/loader/test/unit/MappingConverterTest.java
@@ -99,7 +99,7 @@ public class MappingConverterTest {
"\"value_mapping\":{},\"selected\":[]," +
"\"ignored\":[\"Occupation\",\"Zip-code\",\"Gender\"," +
"\"Age\"],\"null_values\":[\"\"]," +
- "\"update_strategies\":{}}],\"edges\":[]},{\"id\":\"2\"," +
+
"\"update_strategies\":{},\"batch_size\":1000}],\"edges\":[]},{\"id\":\"2\"," +
"\"skip\":false,\"input\":{\"type\":\"FILE\"," +
"\"path\":\"ratings.dat\"," +
"\"file_filter\":{\"extensions\":[\"*\"]}," +
@@ -116,7 +116,7 @@ public class MappingConverterTest {
"\"field_mapping\":{\"UserID\":\"id\",\"MovieID\":\"id\"," +
"\"Rating\":\"rate\"},\"value_mapping\":{},\"selected\":[]," +
"\"ignored\":[\"Timestamp\"],\"null_values\":[\"\"]," +
- "\"update_strategies\":{}}]}]}";
+ "\"update_strategies\":{},\"batch_size\":1000}]}]}";
Assert.assertEquals(expectV2Json, actualV2Json);
FileUtils.forceDelete(inputFile);
