com.tinkerpop.gremlin.hadoop.process.computer.giraph.GiraphMessenger Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hadoop-gremlin Show documentation
The newest version!
package com.tinkerpop.gremlin.hadoop.process.computer.giraph;

import com.tinkerpop.gremlin.hadoop.structure.io.ObjectWritable;
import com.tinkerpop.gremlin.process.Traversal;
import com.tinkerpop.gremlin.process.computer.MessageScope;
import com.tinkerpop.gremlin.process.computer.Messenger;
import com.tinkerpop.gremlin.process.graph.step.map.VertexStep;
import com.tinkerpop.gremlin.process.graph.step.sideEffect.StartStep;
import com.tinkerpop.gremlin.process.util.TraversalHelper;
import com.tinkerpop.gremlin.structure.Direction;
import com.tinkerpop.gremlin.structure.Edge;
import com.tinkerpop.gremlin.structure.Vertex;
import com.tinkerpop.gremlin.util.iterator.IteratorUtils;
import org.apache.hadoop.io.LongWritable;

/**
 * @author Marko A. Rodriguez (http://markorodriguez.com)
 */
public class GiraphMessenger implements Messenger {

    private GiraphComputeVertex giraphComputeVertex;
    private Iterable> messages;

    public void setCurrentVertex(final GiraphComputeVertex giraphComputeVertex, final Iterable> messages) {
        this.giraphComputeVertex = giraphComputeVertex;
        this.messages = messages;
    }

    @Override
    public Iterable receiveMessages(final MessageScope messageScope) {
        return IteratorUtils.map(this.messages, ObjectWritable::get);
    }

    @Override
    public void sendMessage(final MessageScope messageScope, final M message) {
        if (messageScope instanceof MessageScope.Local) {
            final MessageScope.Local localMessageScope = (MessageScope.Local) messageScope;
            final Traversal.Admin incidentTraversal = GiraphMessenger.setVertexStart(localMessageScope.getIncidentTraversal().get(), this.giraphComputeVertex.getBaseVertex());
            final Direction direction = GiraphMessenger.getOppositeDirection(incidentTraversal);
            incidentTraversal.forEachRemaining(edge ->
                    this.giraphComputeVertex.sendMessage(
                            new LongWritable(Long.valueOf(edge.iterators().vertexIterator(direction).next().id().toString())),
                            new ObjectWritable<>(localMessageScope.getEdgeFunction().apply(message, edge))));
        } else {
            final MessageScope.Global globalMessageScope = (MessageScope.Global) messageScope;
            globalMessageScope.vertices().forEach(vertex ->
                    this.giraphComputeVertex.sendMessage(new LongWritable(Long.valueOf(vertex.id().toString())), new ObjectWritable<>(message)));
        }
    }

    private static > T setVertexStart(final Traversal incidentTraversal, final Vertex vertex) {
        incidentTraversal.asAdmin().addStep(0, new StartStep<>(incidentTraversal, vertex));
        return (T) incidentTraversal;
    }

    private static Direction getOppositeDirection(final Traversal.Admin incidentTraversal) {
        final VertexStep step = TraversalHelper.getLastStepOfAssignableClass(VertexStep.class, incidentTraversal).get();
        return step.getDirection().opposite();
    }
}