org.apache.tinkerpop.gremlin.process.computer.bulkloading.BulkLoaderVertexProgram.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.tinkerpop.gremlin.process.computer.bulkloading.BulkLoaderVertexProgram.java

Source

/*
 * 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 org.apache.tinkerpop.gremlin.process.computer.bulkloading;

import org.apache.commons.configuration.BaseConfiguration;
import org.apache.commons.configuration.Configuration;
import org.apache.commons.configuration.ConfigurationException;
import org.apache.commons.configuration.ConfigurationUtils;
import org.apache.commons.configuration.PropertiesConfiguration;
import org.apache.tinkerpop.gremlin.process.computer.GraphComputer;
import org.apache.tinkerpop.gremlin.process.computer.Memory;
import org.apache.tinkerpop.gremlin.process.computer.MessageScope;
import org.apache.tinkerpop.gremlin.process.computer.Messenger;
import org.apache.tinkerpop.gremlin.process.computer.VertexComputeKey;
import org.apache.tinkerpop.gremlin.process.computer.VertexProgram;
import org.apache.tinkerpop.gremlin.process.computer.util.AbstractVertexProgramBuilder;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.__;
import org.apache.tinkerpop.gremlin.process.traversal.step.util.event.MutationListener;
import org.apache.tinkerpop.gremlin.process.traversal.strategy.decoration.EventStrategy;
import org.apache.tinkerpop.gremlin.structure.Direction;
import org.apache.tinkerpop.gremlin.structure.Edge;
import org.apache.tinkerpop.gremlin.structure.Graph;
import org.apache.tinkerpop.gremlin.structure.Property;
import org.apache.tinkerpop.gremlin.structure.Vertex;
import org.apache.tinkerpop.gremlin.structure.VertexProperty;
import org.apache.tinkerpop.gremlin.structure.util.GraphFactory;
import org.apache.tinkerpop.gremlin.structure.util.StringFactory;
import org.javatuples.Pair;
import org.javatuples.Tuple;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

/**
 * @author Daniel Kuppitz (http://gremlin.guru)
 */
public class BulkLoaderVertexProgram implements VertexProgram<Tuple> {

    private static final Logger LOGGER = LoggerFactory.getLogger(BulkLoaderVertexProgram.class);

    public static final String BULK_LOADER_VERTEX_PROGRAM_CFG_PREFIX = "gremlin.bulkLoaderVertexProgram";
    public static final String BULK_LOADER_CLASS_CFG_KEY = String.join(".", BULK_LOADER_VERTEX_PROGRAM_CFG_PREFIX,
            "class");
    public static final String BULK_LOADER_VERTEX_ID_CFG_KEY = String.join(".",
            BULK_LOADER_VERTEX_PROGRAM_CFG_PREFIX, "vertexIdProperty");
    public static final String INTERMEDIATE_BATCH_SIZE_CFG_KEY = String.join(".",
            BULK_LOADER_VERTEX_PROGRAM_CFG_PREFIX, "intermediateBatchSize");
    public static final String KEEP_ORIGINAL_IDS_CFG_KEY = String.join(".", BULK_LOADER_VERTEX_PROGRAM_CFG_PREFIX,
            "keepOriginalIds");
    public static final String USER_SUPPLIED_IDS_CFG_KEY = String.join(".", BULK_LOADER_VERTEX_PROGRAM_CFG_PREFIX,
            "userSuppliedIds");
    public static final String WRITE_GRAPH_CFG_KEY = String.join(".", BULK_LOADER_VERTEX_PROGRAM_CFG_PREFIX,
            "writeGraph");
    public static final String DEFAULT_BULK_LOADER_VERTEX_ID = "bulkLoader.vertex.id";

    private final MessageScope messageScope;
    private final Set<VertexComputeKey> elementComputeKeys;
    private Configuration configuration;
    private BulkLoader bulkLoader;
    private Graph graph;
    private GraphTraversalSource g;
    private long intermediateBatchSize;

    private BulkLoadingListener listener;

    private BulkLoaderVertexProgram() {
        messageScope = MessageScope.Local.of(__::inE);
        elementComputeKeys = new HashSet<>();
    }

    private BulkLoader createBulkLoader() {
        final BulkLoader loader;
        final Configuration config = configuration.subset(BULK_LOADER_VERTEX_PROGRAM_CFG_PREFIX);
        if (config.containsKey("class")) {
            final String className = config.getString("class");
            try {
                final Class<?> bulkLoaderClass = Class.forName(className);
                loader = (BulkLoader) bulkLoaderClass.getConstructor().newInstance();
            } catch (ClassNotFoundException e) {
                LOGGER.error("Unable to find custom bulk loader class: {}", className);
                throw new IllegalStateException(e);
            } catch (Exception e) {
                LOGGER.error("Unable to create an instance of the given bulk loader class: {}", className);
                throw new IllegalStateException(e);
            }
        } else {
            loader = new IncrementalBulkLoader();
        }
        loader.configure(configuration);
        return loader;
    }

    /**
     * Eventually commits the current transaction and closes the current graph instance. commit() will be called
     * if close is set true, otherwise it will only be called if the intermediate batch size is set and reached.
     *
     * @param close Whether to close the current graph instance after calling commit() or not.
     */
    private void commit(final boolean close) {
        if (!close && (intermediateBatchSize == 0L || listener.mutations() < intermediateBatchSize))
            return;
        if (null != graph) {
            if (graph.features().graph().supportsTransactions()) {
                LOGGER.info("Committing transaction on Graph instance: {} [{} mutations]", graph,
                        listener.mutations());
                try {
                    graph.tx().commit();
                    LOGGER.debug("Committed transaction on Graph instance: {}", graph);
                    listener.resetCounter();
                } catch (Exception e) {
                    LOGGER.error("Failed to commit transaction on Graph instance: {}", graph);
                    graph.tx().rollback();
                    listener.resetCounter();
                    throw e;
                }
            }
            if (close) {
                try {
                    graph.close();
                    LOGGER.info("Closed Graph instance: {}", graph);
                    graph = null;
                } catch (Exception e) {
                    LOGGER.warn("Failed to close Graph instance", e);
                }
            }
        }
    }

    @Override
    public void setup(final Memory memory) {
    }

    @Override
    public void loadState(final Graph graph, final Configuration config) {
        configuration = new BaseConfiguration();
        if (config != null) {
            ConfigurationUtils.copy(config, configuration);
        }
        intermediateBatchSize = configuration.getLong(INTERMEDIATE_BATCH_SIZE_CFG_KEY, 0L);
        elementComputeKeys.add(VertexComputeKey.of(DEFAULT_BULK_LOADER_VERTEX_ID, true));
        bulkLoader = createBulkLoader();
    }

    @Override
    public void storeState(final Configuration config) {
        VertexProgram.super.storeState(config);
        if (configuration != null) {
            ConfigurationUtils.copy(configuration, config);
        }
    }

    @Override
    public void workerIterationStart(final Memory memory) {
        if (null == graph) {
            graph = GraphFactory.open(configuration.subset(WRITE_GRAPH_CFG_KEY));
            LOGGER.info("Opened Graph instance: {}", graph);
            try {
                listener = new BulkLoadingListener();
                g = graph.traversal().withStrategies(EventStrategy.build().addListener(listener).create());
            } catch (Exception e) {
                try {
                    graph.close();
                } catch (Exception e2) {
                    LOGGER.warn("Failed to close Graph instance", e2);
                }
                throw e;
            }
        } else {
            LOGGER.warn("Leaked Graph instance: {}", graph);
        }
    }

    @Override
    public void workerIterationEnd(final Memory memory) {
        this.commit(true);
    }

    @Override
    public void execute(final Vertex sourceVertex, final Messenger<Tuple> messenger, final Memory memory) {
        try {
            executeInternal(sourceVertex, messenger, memory);
        } catch (Exception e) {
            if (graph.features().graph().supportsTransactions()) {
                graph.tx().rollback();
            }
            throw e;
        }
    }

    private void executeInternal(final Vertex sourceVertex, final Messenger<Tuple> messenger, final Memory memory) {
        if (memory.isInitialIteration()) {
            this.listener.resetStats();
            // get or create the vertex
            final Vertex targetVertex = bulkLoader.getOrCreateVertex(sourceVertex, graph, g);
            // write all the properties of the vertex to the newly created vertex
            final Iterator<VertexProperty<Object>> vpi = sourceVertex.properties();
            if (this.listener.isNewVertex()) {
                vpi.forEachRemaining(vp -> bulkLoader.createVertexProperty(vp, targetVertex, graph, g));
            } else {
                vpi.forEachRemaining(vp -> bulkLoader.getOrCreateVertexProperty(vp, targetVertex, graph, g));
            }
            this.commit(false);
            if (!bulkLoader.useUserSuppliedIds()) {
                // create an id pair and send it to all the vertex's incoming adjacent vertices
                sourceVertex.property(DEFAULT_BULK_LOADER_VERTEX_ID, targetVertex.id());
                messenger.sendMessage(messageScope, Pair.with(sourceVertex.id(), targetVertex.id()));
            }
        } else if (memory.getIteration() == 1) {
            if (bulkLoader.useUserSuppliedIds()) {
                final Vertex outV = bulkLoader.getVertex(sourceVertex, graph, g);
                final boolean incremental = outV.edges(Direction.OUT).hasNext();
                sourceVertex.edges(Direction.OUT).forEachRemaining(edge -> {
                    final Vertex inV = bulkLoader.getVertex(edge.inVertex(), graph, g);
                    if (incremental) {
                        bulkLoader.getOrCreateEdge(edge, outV, inV, graph, g);
                    } else {
                        bulkLoader.createEdge(edge, outV, inV, graph, g);
                    }
                    this.commit(false);
                });
            } else {
                // create an id map and populate it with all the incoming messages
                final Map<Object, Object> idPairs = new HashMap<>();
                final Iterator<Tuple> idi = messenger.receiveMessages();
                while (idi.hasNext()) {
                    final Tuple idPair = idi.next();
                    idPairs.put(idPair.getValue(0), idPair.getValue(1));
                }
                // get the vertex with given the dummy id property
                final Object outVId = sourceVertex.value(DEFAULT_BULK_LOADER_VERTEX_ID);
                final Vertex outV = bulkLoader.getVertexById(outVId, graph, g);
                // for all the incoming edges of the vertex, get the incoming adjacent vertex and write the edge and its properties
                sourceVertex.edges(Direction.OUT).forEachRemaining(edge -> {
                    final Object inVId = idPairs.get(edge.inVertex().id());
                    final Vertex inV = bulkLoader.getVertexById(inVId, graph, g);
                    bulkLoader.getOrCreateEdge(edge, outV, inV, graph, g);
                    this.commit(false);
                });
            }
        } else if (memory.getIteration() == 2) {
            final Object vertexId = sourceVertex.value(DEFAULT_BULK_LOADER_VERTEX_ID);
            bulkLoader.getVertexById(vertexId, graph, g).property(bulkLoader.getVertexIdProperty()).remove();
            this.commit(false);
        }
    }

    @Override
    public boolean terminate(final Memory memory) {
        switch (memory.getIteration()) {
        case 1:
            return bulkLoader.keepOriginalIds() || bulkLoader.getVertexIdProperty() == null;
        case 2:
            return true;
        }
        return false;
    }

    @Override
    public Set<VertexComputeKey> getVertexComputeKeys() {
        return elementComputeKeys;
    }

    @Override
    public Set<MessageScope> getMessageScopes(final Memory memory) {
        return Collections.singleton(messageScope);
    }

    @SuppressWarnings({ "CloneDoesntDeclareCloneNotSupportedException", "CloneDoesntCallSuperClone" })
    @Override
    public VertexProgram<Tuple> clone() {
        return this;
    }

    @Override
    public GraphComputer.ResultGraph getPreferredResultGraph() {
        return GraphComputer.ResultGraph.ORIGINAL;
    }

    @Override
    public GraphComputer.Persist getPreferredPersist() {
        return GraphComputer.Persist.NOTHING;
    }

    @Override
    public String toString() {
        final StringBuilder sb = new StringBuilder();
        if (bulkLoader != null) {
            sb.append("bulkLoader=").append(bulkLoader.getClass().getSimpleName()).append(", ");
            sb.append("vertexIdProperty=").append(bulkLoader.getVertexIdProperty()).append(", ");
            sb.append("userSuppliedIds=").append(bulkLoader.useUserSuppliedIds()).append(", ");
            sb.append("keepOriginalIds=").append(bulkLoader.keepOriginalIds()).append(", ");
        } else {
            sb.append("bulkLoader=").append(bulkLoader).append(", ");
        }
        sb.append("batchSize=").append(intermediateBatchSize);
        return StringFactory.vertexProgramString(this, sb.toString());
    }

    public static Builder build() {
        return new Builder();
    }

    public static class Builder extends AbstractVertexProgramBuilder<Builder> {

        private Builder() {
            super(BulkLoaderVertexProgram.class);
        }

        @SuppressWarnings("unchecked")
        @Override
        public BulkLoaderVertexProgram create(final Graph graph) {
            ConfigurationUtils.append(graph.configuration().subset(BULK_LOADER_VERTEX_PROGRAM_CFG_PREFIX),
                    configuration);
            return (BulkLoaderVertexProgram) VertexProgram.createVertexProgram(graph, configuration);
        }

        private void setGraphConfigurationProperty(final String key, final Object value) {
            configuration.setProperty(String.join(".", WRITE_GRAPH_CFG_KEY, key), value);
        }

        /**
         * Sets the class name of the BulkLoader implementation to be used.
         */
        public Builder bulkLoader(final String className) {
            configuration.setProperty(BULK_LOADER_CLASS_CFG_KEY, className);
            return this;
        }

        /**
         * Sets the class of the BulkLoader implementation to be used.
         */
        public Builder bulkLoader(final Class<? extends BulkLoader> clazz) {
            return bulkLoader(clazz.getCanonicalName());
        }

        /**
         * Sets the name of the property that is used to store the original vertex identifiers in the target graph.
         */
        public Builder vertexIdProperty(final String name) {
            configuration.setProperty(BULK_LOADER_VERTEX_ID_CFG_KEY, name);
            return this;
        }

        /**
         * Specifies whether user supplied identifiers should be used when the bulk loader creates vertices in the
         * target graph.
         */
        public Builder userSuppliedIds(final boolean useUserSuppliedIds) {
            configuration.setProperty(USER_SUPPLIED_IDS_CFG_KEY, useUserSuppliedIds);
            return this;
        }

        /**
         * Specifies whether the original vertex identifiers should be kept in the target graph or not. In case of false
         * BulkLoaderVertexProgram will add another iteration to remove the properties and it won't be possible to use
         * the data for further incremental bulk loads.
         */
        public Builder keepOriginalIds(final boolean keepOriginalIds) {
            configuration.setProperty(KEEP_ORIGINAL_IDS_CFG_KEY, keepOriginalIds);
            return this;
        }

        /**
         * The batch size for a single transaction (number of vertices in the vertex loading stage; number of edges in
         * the edge loading stage).
         */
        public Builder intermediateBatchSize(final int batchSize) {
            configuration.setProperty(INTERMEDIATE_BATCH_SIZE_CFG_KEY, batchSize);
            return this;
        }

        /**
         * A configuration for the target graph that can be passed to GraphFactory.open().
         */
        public Builder writeGraph(final String configurationFile) throws ConfigurationException {
            return writeGraph(new PropertiesConfiguration(configurationFile));
        }

        /**
         * A configuration for the target graph that can be passed to GraphFactory.open().
         */
        public Builder writeGraph(final Configuration configuration) {
            configuration.getKeys()
                    .forEachRemaining(key -> setGraphConfigurationProperty(key, configuration.getProperty(key)));
            return this;
        }
    }

    @Override
    public Features getFeatures() {
        return new Features() {
            @Override
            public boolean requiresLocalMessageScopes() {
                return true;
            }

            @Override
            public boolean requiresVertexPropertyAddition() {
                return true;
            }
        };
    }

    static class BulkLoadingListener implements MutationListener {

        private long counter;
        private boolean isNewVertex;

        public BulkLoadingListener() {
            this.counter = 0L;
            this.isNewVertex = false;
        }

        public boolean isNewVertex() {
            return this.isNewVertex;
        }

        public long mutations() {
            return this.counter;
        }

        public void resetStats() {
            this.isNewVertex = false;
        }

        public void resetCounter() {
            this.counter = 0L;
        }

        @Override
        public void vertexAdded(final Vertex vertex) {
            this.isNewVertex = true;
            this.counter++;
        }

        @Override
        public void vertexRemoved(final Vertex vertex) {
            this.counter++;
        }

        @Override
        public void vertexPropertyChanged(final Vertex element, final Property oldValue, final Object setValue,
                final Object... vertexPropertyKeyValues) {
            this.counter++;
        }

        @Override
        public void vertexPropertyRemoved(final VertexProperty vertexProperty) {
            this.counter++;
        }

        @Override
        public void edgeAdded(final Edge edge) {
            this.counter++;
        }

        @Override
        public void edgeRemoved(final Edge edge) {
            this.counter++;
        }

        @Override
        public void edgePropertyChanged(final Edge element, final Property oldValue, final Object setValue) {
            this.counter++;
        }

        @Override
        public void edgePropertyRemoved(final Edge element, final Property property) {
            this.counter++;
        }

        @Override
        public void vertexPropertyPropertyChanged(final VertexProperty element, final Property oldValue,
                final Object setValue) {
            this.counter++;
        }

        @Override
        public void vertexPropertyPropertyRemoved(final VertexProperty element, final Property property) {
            this.counter++;
        }
    }
}