DgraphSample.java

package com.sample.dgraph;

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

import io.dgraph.DgraphClient;
import io.dgraph.DgraphGrpc;
import io.dgraph.DgraphProto.Operation;
import io.dgraph.Transaction;
import io.dgraph.DgraphGrpc.DgraphStub;
import io.dgraph.DgraphProto.Assigned;
import io.dgraph.DgraphProto.Mutation;
import io.dgraph.DgraphProto.Response;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;

import org.json.JSONArray;
import org.json.JSONObject;
import org.junit.Test;

import com.google.protobuf.ByteString;

public class DgraphSample {

    private static String hostname;
    private static int portNo;

    protected static ManagedChannel channel;
    protected static DgraphClient dgraphClient;

    public void initializeDgraph() {
        ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost",18271).usePlaintext(true).build();
        DgraphStub stub = DgraphGrpc.newStub(channel);
        DgraphClient dgraphClient = new DgraphClient(stub);
        dgraphClient.alter(Operation.newBuilder().setDropAll(true).build());
        Operation op = Operation.newBuilder().setSchema("name: string @index(term, trigram) . ").build();
        dgraphClient.alter(op);
    }

    public DgraphSample() {
        hostname = "localhost";
        portNo = 18271;

        if (channel == null || channel.isShutdown() || dgraphClient == null) {
            channel =
                    ManagedChannelBuilder.forAddress(hostname, portNo).usePlaintext(true).build();
            DgraphStub stub = DgraphGrpc.newStub(channel);
            dgraphClient = new DgraphClient(stub);
        }
    }

    public void createNodeListInBatch(List<Node> nodeList) {
        Transaction txn = dgraphClient.newTransaction();
        try {
            int i = 0;
            for (Node node : nodeList) {
                if (i % 3000 == 0) {
                    txn.commit(); //txn commit because it gets destroy after some mutation
                    txn.close();
                    txn = dgraphClient.newTransaction();
                }
                i++;

                JSONObject nodeJsonObj = new JSONObject();
                nodeJsonObj.put("name", node.getName());

                Mutation mu = Mutation.newBuilder().setSetJson(ByteString.copyFromUtf8(nodeJsonObj.toString())).build();
                txn.mutate(mu);
            }
            txn.commit();

        } catch(Exception e) {
            e.printStackTrace();
        } finally {
            txn.discard();
        }
    }

    public String createNode(Node node) {
        Transaction txn = dgraphClient.newTransaction();
        try {
            JSONObject nodeJsonObj = new JSONObject();

            nodeJsonObj.put("name", node.getName());

            Mutation mu = Mutation.newBuilder().setSetJson(ByteString.copyFromUtf8(nodeJsonObj.toString())).build();
            Assigned assigned = txn.mutate(mu);
            txn.commit();

            return assigned.getUidsOrThrow("blank-0");
        } catch(Exception e) {
            e.printStackTrace();
        } finally {
            txn.discard();
        }
        return null;
    }

    public Node getNodeByPartialSearch(String nodeName) {
        if (nodeName == null || nodeName.equals("")) {
            return null;
        }
        Node node = null;
        Transaction txn = dgraphClient.newTransaction();
        try {
            String query = "query node($a: string) { node(func: allofterms(name, $a)) {uid name _predicate_ } }";

            Map<String, String> vars = new HashMap<>();
            vars.put("$a", nodeName);

            Response resp = txn.queryWithVars(query, vars);

            JSONArray nodeArr = new JSONObject(resp.getJson().toStringUtf8()).optJSONArray("node");

            if (nodeArr != null && nodeArr.length() > 0) {
                for (int i=0; i<nodeArr.length(); i++) {
                    JSONObject nodeObj = nodeArr.optJSONObject(i);
                    if (nodeObj.optString("name").equalsIgnoreCase(nodeName)) {
                        node = new Node(nodeObj.optString("uid"), nodeObj.optString("name"));
                        break;
                    }
                }
            }
        } catch(Exception e) {
            e.printStackTrace();
        } finally {
            txn.discard();
        }
        return node;
    }

    class Node {
        String name;
        String id;

        public Node(String name) {
            this.name = name;
        }

        public Node(String id, String name) {
            this.id = id;
            this.name = name;
        }

        public String getId() {
            return id;
        }

        public void setId(String id) {
            this.id = id;
        }

        public String getName() {
            return name;
        }

        public void setName(String name) {
            this.name = name;
        }

        @Override
        public String toString() {
            return "Node [name=" + name + ", id=" + id + "]";
        }
    }

    @Test
    public void createDGraphSampleData() {
        DgraphSample dgraph = new DgraphSample();

        dgraph.initializeDgraph();

        List<Node> nodeList = new ArrayList<>();
        for (int i=1; i<=50000; i++) {
            Node node = new Node("India"+i);
            nodeList.add(node);
        }

        long startTime = new Date().getTime();
        dgraph.createNodeListInBatch(nodeList);
        long endTime = new Date().getTime();
        long totalTime = endTime - startTime;
        System.out.println("Total time taken to execute createNodeListInBatch " + totalTime);
        //50k nodes Took 15.16 mins

        startTime = new Date().getTime();
        for (int i=1; i<=50000; i++) {
            Node node = new Node("Canada"+i);
            dgraph.createNode(node);
        }
        endTime = new Date().getTime();
        totalTime = endTime - startTime;
        System.out.println("Total time taken to execute createNode one by one " + totalTime);
        //50k nodes Took 72sec
    }

    @Test
    public void getDGraphSampleData() {
        DgraphSample dgraph = new DgraphSample();

        long startTime = new Date().getTime();
        for (int i=1; i<=50000; i++) {
            dgraph.getNodeByPartialSearch("Canada"+i);
            dgraph.getNodeByPartialSearch("India"+i);
        }
        long endTime = new Date().getTime();
        long totalTime = endTime - startTime;
        System.out.println("For All Node >> Total time taken to execute getNodeByPartialSearch " + totalTime);
        //50k nodes Took 2.17 mins
    }
}