neo4j数据库导入到arangodb

参考文章： 图数据库 Neo4j 之 Java Api 的使用

java向Neo4j添加节点及其关系

简介

主要功能：读取neo4j数据库文件，将数据导入到arangodb中

pom.xml文件

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <properties>
        <project-name>graphdb-data-util</project-name>
    </properties>

    <groupId>org.fiend</groupId>
    <artifactId>graphdb-data-util</artifactId>
    <version>1.0-SNAPSHOT</version>
    <name>${project-name}</name>

    <dependencies>
        <!-- Neo4j dependency -->
        <dependency>
            <groupId>org.neo4j</groupId>
            <artifactId>neo4j</artifactId>
            <version>3.5.0</version>
        </dependency>

        <!-- Arangodb dependency -->
        <!-- syn -->
        <dependency>
            <groupId>com.arangodb</groupId>
            <artifactId>arangodb-java-driver</artifactId>
            <version>5.0.4</version>
        </dependency>
        <!-- async -->
        <dependency>
            <groupId>com.arangodb</groupId>
            <artifactId>arangodb-java-driver-async</artifactId>
            <version>5.0.4</version>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <!-- maven打source包 -->
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-source-plugin</artifactId>
                <executions>
                    <execution>
                        <id>attach-sources</id>
                        <goals>
                            <goal>jar</goal>
                        </goals>
                    </execution>
                </executions>
            </plugin>
        </plugins>
    </build>
</project>

java文件

package org.fiend.graphdb;

import com.arangodb.ArangoCollection;
import com.arangodb.ArangoCursor;
import com.arangodb.ArangoDB;
import com.arangodb.ArangoDatabase;
import com.arangodb.entity.*;
import com.arangodb.model.CollectionCreateOptions;
import com.arangodb.util.MapBuilder;
import org.neo4j.graphdb.GraphDatabaseService;
import org.neo4j.graphdb.Node;
import org.neo4j.graphdb.Relationship;
import org.neo4j.graphdb.Transaction;
import org.neo4j.graphdb.factory.GraphDatabaseFactory;

import java.io.File;
import java.util.*;

public class CopyFromNeo4jToArangoDB {
    private static GraphDatabaseService neo4jdb = null;
    
    private static ArangoDatabase arangodb1 = null;
    private static ArangoDatabase arangodb2 = null;
    
    private static final String argdbName2 = "langpf_graph";

    public static void main(String[] args) {
        GraphDatabaseFactory dbFactory = new GraphDatabaseFactory();
        // neo4jdb= dbFactory.newEmbeddedDatabase(new File("/data/arangodb/graph8.db"));
        neo4jdb = dbFactory.newEmbeddedDatabase(new File("/data/neo4j/graph23.db"));

        // ArangoDB builder = new ArangoDB.Builder().host("127.0.0.1", 8529).user("admin").password("123456").build();
        // ArangoDB builder = new ArangoDB.Builder().host("10.101.15.205", 8531)
        //         .host("10.101.15.206", 8531).host("10.101.15.207", 8531).build();
        ArangoDB builder = new ArangoDB.Builder().host("10.101.15.254", 8531).user("langpf").password("123abc").build();

        // clearDB(builder, db_name1);
        clearDB(builder, argdbName2);

        // arangodb1 = builder.db(db_name1);
        arangodb2 = builder.db(argdbName2);

        // long time1 = System.currentTimeMillis()/1000;
        // copy1();
        long time2 = System.currentTimeMillis() / 1000;
        copy2();
        long time3 = System.currentTimeMillis() / 1000;

        // System.out.println(db_name1+"==="+(time2-time1)+"s");
        System.out.println(argdbName2 + "===" + (time3 - time2) + "s");

        System.out.println("Finish.");
        neo4jdb.shutdown();
        builder.shutdown();
    }

    /**
     * arangodb -- clear
     */
    private static void clearDB(ArangoDB builder, String dbName) {
        Set<String> dbs = new HashSet<>(builder.getDatabases());
        
        if(dbs.contains(dbName)) {
            builder.db(dbName).drop();
        }
        
        builder.createDatabase(dbName);
    }

    /* 第二种导出方案:
        graph关联多个collection, 类型为document的collection名称为label, 仅包含属性name;
        类型为edge的collection名称为type, 无自定义属性。
        共11个document collection, 19个edge collection。
    */
    private static void copy2() {
        Map<Long, String> id_map = new HashMap<>();
        Map<String, Set<String>> edge_label_map = new HashMap<>();
        
        int nodeNum = 0;
        int edgeNum = 0;
        
        try (Transaction tx = neo4jdb.beginTx()) {
            for (Node node : neo4jdb.getAllNodes()) {
                addNode2(node, id_map);

                nodeNum++;
                if (nodeNum % 10000 == 0) {
                    System.out.println("finish node===" + nodeNum);
                }
            }

            for (Relationship edge : neo4jdb.getAllRelationships()) {
                Node start_node = edge.getStartNode();
                Node end_node = edge.getEndNode();

                String start_label = start_node.getLabels().iterator().next().name();
                String end_label = end_node.getLabels().iterator().next().name();

                String from_id = id_map.get(start_node.getId());
                String to_id = id_map.get(end_node.getId());
                String type = start_label + end_label;
                addEdge2(from_id, to_id, type);

                Set<String> edge_set = edge_label_map.containsKey(start_label) ? edge_label_map.get(start_label) : new HashSet<>();
                edge_set.add(end_label);
                edge_label_map.put(start_label, edge_set);

                edgeNum++;
                if (edgeNum % 1000 == 0) {
                    System.out.println("finish edge===" + edgeNum);
                }
            }
            
            tx.success();
        }

        createGraph2(edge_label_map);

        System.out.println("node===" + nodeNum);
        System.out.println("edge===" + edgeNum);
    }

    private static void addNode2(Node node, Map<Long, String> id_map) {
        String label = node.getLabels().iterator().next().name();
        String name  = node.getProperty("name").toString();

        if (!hasCollection(arangodb2, label)) {
            arangodb2.createCollection(label);
        }
        ArangoCollection node_collection = arangodb2.collection(label);

        String query = "for n in @@coll filter n.name==@name return n";
        Map<String, Object> bindVars = new MapBuilder().put("@coll", label).put("name", name).get();
        ArangoCursor<BaseDocument> cursor = arangodb2.query(query, bindVars, null, BaseDocument.class);

        if (cursor.hasNext()) {
            // System.out.println(label+"\t"+name);
            id_map.put(node.getId(), cursor.next().getId());

            return;
        }

        BaseDocument doc = new BaseDocument();
        doc.addAttribute("name", name);
        DocumentCreateEntity<BaseDocument> dce = node_collection.insertDocument(doc);

        id_map.put(node.getId(), dce.getId());
    }

    private static void addEdge2(String from_id, String to_id, String type) {
        if (null == from_id || null == to_id) {
            return;
        }

        if (!hasCollection(arangodb2, type)) {
            CollectionCreateOptions cco = new CollectionCreateOptions();
            cco.type(CollectionType.EDGES);
            arangodb2.createCollection(type, cco);
        }
        ArangoCollection edge_collection = arangodb2.collection(type);

        String query = "for e in @@coll filter e._from==@from and e._to==@to return e";
        Map<String, Object> bindVars = new MapBuilder().put("@coll", type).put("from", from_id).put("to", to_id).get();
        ArangoCursor<BaseEdgeDocument> cursor = arangodb2.query(query, bindVars, null, BaseEdgeDocument.class);

        if (cursor.hasNext()) {
            // System.out.println(from_id+"\t"+to_id+"\t"+type);
            return;
        }

        BaseEdgeDocument edge = new BaseEdgeDocument();
        edge.setFrom(from_id);
        edge.setTo(to_id);
        edge_collection.insertDocument(edge);
    }

    private static void createGraph2(Map<String, Set<String>> map) {
        List<EdgeDefinition> ed_list = new ArrayList<>();

        for (String sl : map.keySet()) {
            Set<String> set = map.get(sl);
            for (String el : set) {

                EdgeDefinition ed = new EdgeDefinition();
                ed.collection(sl + el);
                ed.from(sl);
                ed.to(el);
                ed_list.add(ed);
            }
        }

        arangodb2.createGraph("graph", ed_list);
    }

    /* 第一种导出方案: 
        graph关联两个collection, 分别为edge(类型：edge)和node(类型：document),
        其中node中属性包含label和name, edge中属性包含type。
    */
    private static void copy1() {
        Map<Long, String> id_map = new HashMap<>();

        int node_num = 0;
        int edge_num = 0;

        try (Transaction tx = neo4jdb.beginTx()) {
            Iterator<Node> nodeIter = neo4jdb.getAllNodes().iterator();
            while (nodeIter.hasNext()) {
                Node node = nodeIter.next();
                addNode1(node, id_map);

                node_num++;
                if (node_num % 10000 == 0) {
                    System.out.println("finish node===" + node_num);
                }
            }

            Iterator<Relationship> edgeIter = neo4jdb.getAllRelationships().iterator();
            while (edgeIter.hasNext()) {
                Relationship edge = edgeIter.next();
                String from_id = id_map.get(edge.getStartNodeId());
                String to_id = id_map.get(edge.getEndNodeId());
                String type = edge.getType().name();
                addEdge1(from_id, to_id, type);

                edge_num++;
                if (edge_num % 10000 == 0) {
                    System.out.println("finish edge===" + edge_num);
                }
            }

            tx.success();
        }

        createGraph1();

        System.out.println("node===" + node_num);
        System.out.println("edge===" + edge_num);
    }

    private static void addNode1(Node node, Map<Long, String> id_map) {
        if (!hasCollection(arangodb1, "node")) {
            CollectionCreateOptions cco = new CollectionCreateOptions();
            cco.numberOfShards(6);
            arangodb1.createCollection("node", cco);
        }
        ArangoCollection node_collection = arangodb1.collection("node");

        String label = node.getLabels().iterator().next().name();
        String name = node.getProperty("name").toString();

        String query = "for n in node filter n.label==@label and n.name==@name return n";
        Map<String, Object> bindVars = new MapBuilder().put("name", name).put("label", label).get();
        ArangoCursor<BaseDocument> cursor = arangodb1.query(query, bindVars, null, BaseDocument.class);

        if (cursor.hasNext()) {
            // System.out.println(label+"\t"+name);
            id_map.put(node.getId(), cursor.next().getId());
            return;
        }

        BaseDocument doc = new BaseDocument();
        doc.addAttribute("label", label);
        doc.addAttribute("name", name);
        DocumentCreateEntity<BaseDocument> dce = node_collection.insertDocument(doc);

        id_map.put(node.getId(), dce.getId());
    }

    private static void addEdge1(String from_id, String to_id, String type) {
        if (null == from_id || null == to_id) {
            return;
        }

        if (!hasCollection(arangodb1, "edge")) {
            CollectionCreateOptions cco = new CollectionCreateOptions();
            cco.type(CollectionType.EDGES);
            cco.numberOfShards(6);
            arangodb1.createCollection("edge", cco);
        }
        ArangoCollection edge_collection = arangodb1.collection("edge");

        String query = "for e in edge filter e._from==@from and e._to==@to return e";
        Map<String, Object> bindVars = new MapBuilder().put("from", from_id).put("to", to_id).get();
        ArangoCursor<BaseEdgeDocument> cursor = arangodb1.query(query, bindVars, null, BaseEdgeDocument.class);

        if (cursor.hasNext()) {
            // System.out.println(from_id+"\t"+to_id+"\t"+type);
            return;
        }

        BaseEdgeDocument edge = new BaseEdgeDocument();
        edge.setFrom(from_id);
        edge.setTo(to_id);
        edge.addAttribute("type", type);
        edge_collection.insertDocument(edge);
    }

    private static void createGraph1() {
        List<EdgeDefinition> ed_list = new ArrayList<>();

        EdgeDefinition ed = new EdgeDefinition();
        ed.collection("edge");
        ed.from("node");
        ed.to("node");
        ed_list.add(ed);

        arangodb1.createGraph("graph", ed_list);
    }

    private static boolean hasCollection(ArangoDatabase arangodb, String name) {
        final List<Byte> list = new ArrayList<>();

        arangodb.getCollections().forEach(coll -> {
            if (coll.getName().equals(name)) {
                list.add((byte) 0);
            }
        });

        return list.size() == 1;
    }
}