/*
* 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.spark.process.computer.traversal.strategy.optimization.interceptor;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.tinkerpop.gremlin.hadoop.structure.io.VertexWritable;
import org.apache.tinkerpop.gremlin.process.computer.ProgramPhase;
import org.apache.tinkerpop.gremlin.process.computer.traversal.MemoryTraversalSideEffects;
import org.apache.tinkerpop.gremlin.process.computer.traversal.TraversalVertexProgram;
import org.apache.tinkerpop.gremlin.process.computer.traversal.strategy.finalization.ComputerFinalizationStrategy;
import org.apache.tinkerpop.gremlin.process.traversal.Scope;
import org.apache.tinkerpop.gremlin.process.traversal.Step;
import org.apache.tinkerpop.gremlin.process.traversal.Traversal;
import org.apache.tinkerpop.gremlin.process.traversal.Traverser;
import org.apache.tinkerpop.gremlin.process.traversal.step.Barrier;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.CountGlobalStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.FoldStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.GraphStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.GroupCountStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.GroupStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.MaxGlobalStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.MeanGlobalStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.MinGlobalStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.SumGlobalStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.util.ReducingBarrierStep;
import org.apache.tinkerpop.gremlin.process.traversal.strategy.decoration.SubgraphStrategy;
import org.apache.tinkerpop.gremlin.process.traversal.strategy.verification.ComputerVerificationStrategy;
import org.apache.tinkerpop.gremlin.process.traversal.traverser.TraverserRequirement;
import org.apache.tinkerpop.gremlin.process.traversal.traverser.util.TraverserSet;
import org.apache.tinkerpop.gremlin.process.traversal.util.TraversalHelper;
import org.apache.tinkerpop.gremlin.spark.process.computer.SparkMemory;
import org.apache.tinkerpop.gremlin.spark.process.computer.traversal.strategy.SparkVertexProgramInterceptor;
import org.apache.tinkerpop.gremlin.structure.Vertex;
import org.apache.tinkerpop.gremlin.structure.util.ElementHelper;
import org.apache.tinkerpop.gremlin.util.NumberHelper;
import org.apache.tinkerpop.gremlin.util.function.ArrayListSupplier;
import org.apache.tinkerpop.gremlin.util.iterator.IteratorUtils;
import java.util.ArrayList;
import java.util.List;
import java.util.function.BinaryOperator;
/**
* @author Marko A. Rodriguez (http://markorodriguez.com)
*/
public final class SparkStarBarrierInterceptor implements SparkVertexProgramInterceptor<TraversalVertexProgram> {
public SparkStarBarrierInterceptor() {
}
@Override
public JavaPairRDD<Object, VertexWritable> apply(final TraversalVertexProgram vertexProgram, final JavaPairRDD<Object, VertexWritable> inputRDD, final SparkMemory memory) {
vertexProgram.setup(memory);
final Traversal.Admin<Vertex, Object> traversal = (Traversal.Admin) vertexProgram.getTraversal().getPure().clone();
final GraphStep<Vertex, Vertex> graphStep = ((GraphStep) traversal.getStartStep());
final Object[] graphStepIds = graphStep.getIds(); // any V(1,2,3)-style ids to filter on
final ReducingBarrierStep endStep = (ReducingBarrierStep) traversal.getEndStep(); // needed for the final traverser generation
traversal.removeStep(0); // remove GraphStep
traversal.removeStep(traversal.getSteps().size() - 1); // remove ReducingBarrierStep
traversal.setStrategies(traversal.clone().getStrategies().removeStrategies(ComputerVerificationStrategy.class, ComputerFinalizationStrategy.class)); // no longer a computer job, but parallel standard jobs
traversal.applyStrategies(); // compile
boolean identityTraversal = traversal.getSteps().isEmpty(); // if the traversal is empty, just return the vertex (fast)
///////////////////////////////
MemoryTraversalSideEffects.setMemorySideEffects(traversal, memory, ProgramPhase.EXECUTE); // any intermediate sideEffect steps are backed by SparkMemory
memory.setInExecute(true);
final JavaRDD<Traverser.Admin<Object>> nextRDD = inputRDD.values()
.filter(vertexWritable -> ElementHelper.idExists(vertexWritable.get().id(), graphStepIds)) // ensure vertex ids are in V(x)
.flatMap(vertexWritable -> {
if (identityTraversal) // g.V.count()-style (identity)
return IteratorUtils.of(traversal.getTraverserGenerator().generate(vertexWritable.get(), (Step) graphStep, 1l));
else { // add the vertex to head of the traversal
final Traversal.Admin<Vertex, ?> clone = traversal.clone(); // need a unique clone for each vertex to isolate the computation
clone.getStartStep().addStart(clone.getTraverserGenerator().generate(vertexWritable.get(), graphStep, 1l));
return (Step) clone.getEndStep();
}
});
// USE SPARK DSL FOR THE RESPECTIVE END REDUCING BARRIER STEP OF THE TRAVERSAL
final Object result;
if (endStep instanceof CountGlobalStep)
result = nextRDD.map(Traverser::bulk).fold(0l, (a, b) -> a + b);
else if (endStep instanceof SumGlobalStep)
result = nextRDD
.map(traverser -> NumberHelper.mul(traverser.bulk(), (Number) traverser.get()))
.fold(0, NumberHelper::add);
else if (endStep instanceof MeanGlobalStep)
result = nextRDD
.map(traverser -> new MeanGlobalStep.MeanNumber((Number) traverser.get(), traverser.bulk()))
.fold(new MeanGlobalStep.MeanNumber(), MeanGlobalStep.MeanNumber::add)
.getFinal();
else if (endStep instanceof MinGlobalStep)
result = nextRDD
.map(traverser -> (Number) traverser.get())
.fold(Integer.MAX_VALUE, NumberHelper::min);
else if (endStep instanceof MaxGlobalStep)
result = nextRDD
.map(traverser -> (Number) traverser.get())
.fold(Integer.MIN_VALUE, NumberHelper::max);
else if (endStep instanceof FoldStep) {
final BinaryOperator biOperator = endStep.getBiOperator();
result = nextRDD.map(traverser -> {
if (endStep.getSeedSupplier() instanceof ArrayListSupplier) {
final List list = new ArrayList<>();
for (long i = 0; i < traverser.bulk(); i++) {
list.add(traverser.get());
}
return list;
} else {
return traverser.get();
}
}).fold(endStep.getSeedSupplier().get(), biOperator::apply);
} else if (endStep instanceof GroupStep) {
final GroupStep.GroupBiOperator<Object, Object> biOperator = (GroupStep.GroupBiOperator) endStep.getBiOperator();
result = ((GroupStep) endStep).generateFinalResult(nextRDD.
mapPartitions(partitions -> {
final GroupStep<Object, Object, Object> clone = (GroupStep) endStep.clone();
return IteratorUtils.map(partitions, clone::projectTraverser);
}).fold(((GroupStep<Object, Object, Object>) endStep).getSeedSupplier().get(), biOperator::apply));
} else if (endStep instanceof GroupCountStep) {
final GroupCountStep.GroupCountBiOperator<Object> biOperator = GroupCountStep.GroupCountBiOperator.instance();
result = nextRDD
.mapPartitions(partitions -> {
final GroupCountStep<Object, Object> clone = (GroupCountStep) endStep.clone();
return IteratorUtils.map(partitions, clone::projectTraverser);
})
.fold(((GroupCountStep<Object, Object>) endStep).getSeedSupplier().get(), biOperator::apply);
} else
throw new IllegalArgumentException("The end step is an unsupported barrier: " + endStep);
memory.setInExecute(false);
///////////////////////////////
// generate the HALTED_TRAVERSERS for the memory
final TraverserSet<Long> haltedTraversers = new TraverserSet<>();
haltedTraversers.add(traversal.getTraverserGenerator().generate(result, endStep, 1l)); // all reducing barrier steps produce a result of bulk 1
memory.set(TraversalVertexProgram.HALTED_TRAVERSERS, haltedTraversers);
memory.incrIteration(); // any local star graph reduction takes a single iteration
return inputRDD;
}
public static boolean isLegal(final Traversal.Admin<?, ?> traversal) {
final Step<?, ?> startStep = traversal.getStartStep();
final Step<?, ?> endStep = traversal.getEndStep();
// right now this is not supported because of how the SparkStarBarrierInterceptor mutates the traversal prior to local evaluation
if (traversal.getStrategies().toList().stream().filter(strategy -> strategy instanceof SubgraphStrategy).findAny().isPresent())
return false;
if (!startStep.getClass().equals(GraphStep.class) || ((GraphStep) startStep).returnsEdge())
return false;
if (!endStep.getClass().equals(CountGlobalStep.class) &&
!endStep.getClass().equals(SumGlobalStep.class) &&
!endStep.getClass().equals(MeanGlobalStep.class) &&
!endStep.getClass().equals(MaxGlobalStep.class) &&
!endStep.getClass().equals(MinGlobalStep.class) &&
!endStep.getClass().equals(FoldStep.class) &&
!endStep.getClass().equals(GroupStep.class) &&
!endStep.getClass().equals(GroupCountStep.class))
// TODO: tree()
return false;
if (TraversalHelper.getStepsOfAssignableClassRecursively(Scope.global, Barrier.class, traversal).size() != 1)
return false;
if (traversal.getTraverserRequirements().contains(TraverserRequirement.SACK))
return false;
return TraversalHelper.isLocalStarGraph(traversal);
}
}