org.apache.tinkerpop.gremlin.process.traversal.strategy.decoration.SubgraphStrategy Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of gremlin-core
There is a newer version: 4.0.0-beta.1
/*
 * 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.traversal.strategy.decoration;

import org.apache.tinkerpop.gremlin.process.traversal.Step;
import org.apache.tinkerpop.gremlin.process.traversal.Traversal;
import org.apache.tinkerpop.gremlin.process.traversal.TraversalStrategy;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.__;
import org.apache.tinkerpop.gremlin.process.traversal.step.filter.TraversalFilterStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.AddEdgeStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.AddVertexStartStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.AddVertexStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.EdgeOtherVertexStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.EdgeVertexStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.VertexStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.GraphStep;
import org.apache.tinkerpop.gremlin.process.traversal.strategy.AbstractTraversalStrategy;
import org.apache.tinkerpop.gremlin.process.traversal.util.TraversalHelper;
import org.apache.tinkerpop.gremlin.structure.Direction;
import org.apache.tinkerpop.gremlin.structure.Edge;
import org.apache.tinkerpop.gremlin.structure.Vertex;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * This {@link TraversalStrategy} provides a way to limit the view of a {@link Traversal}.  By providing
 * {@link Traversal} representations that represent a form of filtering criterion for vertices and/or edges,
 * this strategy will inject that criterion into the appropriate places of a traversal thus restricting what
 * it traverses and returns.
 *
 * @author Stephen Mallette (http://stephen.genoprime.com)
 */
public final class SubgraphStrategy extends AbstractTraversalStrategy
        implements TraversalStrategy.DecorationStrategy {

    private final Traversal vertexCriterion;
    private final Traversal edgeCriterion;

    private SubgraphStrategy(final Traversal vertexCriterion, final Traversal edgeCriterion) {
        this.vertexCriterion = vertexCriterion;

        // if there is no vertex predicate there is no need to test either side of the edge
        if (null == vertexCriterion) {
            this.edgeCriterion = edgeCriterion;
        } else {
            final Traversal inVertexPredicate = __.inV().filter(vertexCriterion);
            final Traversal outVertexPredicate = __.outV().filter(vertexCriterion);

            // if there is a vertex predicate then there is an implied edge filter on vertices even if there is no
            // edge predicate provided by the user.
            if (null == edgeCriterion)
                this.edgeCriterion = __.and(inVertexPredicate.asAdmin(), outVertexPredicate.asAdmin());
            else
                this.edgeCriterion = edgeCriterion.asAdmin().addStep(new TraversalFilterStep<>(edgeCriterion.asAdmin(), __.and(inVertexPredicate.asAdmin(), outVertexPredicate.asAdmin())));
        }
    }

    @Override
    public void apply(final Traversal.Admin traversal) {
        final List graphSteps = TraversalHelper.getStepsOfAssignableClass(GraphStep.class, traversal);
        final List vertexSteps = TraversalHelper.getStepsOfAssignableClass(VertexStep.class, traversal);

        if (vertexCriterion != null) {
            final List vertexStepsToInsertFilterAfter = new ArrayList<>();
            vertexStepsToInsertFilterAfter.addAll(TraversalHelper.getStepsOfAssignableClass(EdgeOtherVertexStep.class, traversal));
            vertexStepsToInsertFilterAfter.addAll(TraversalHelper.getStepsOfAssignableClass(EdgeVertexStep.class, traversal));
            vertexStepsToInsertFilterAfter.addAll(TraversalHelper.getStepsOfAssignableClass(AddVertexStep.class, traversal));
            vertexStepsToInsertFilterAfter.addAll(TraversalHelper.getStepsOfAssignableClass(AddVertexStartStep.class, traversal));
            vertexStepsToInsertFilterAfter.addAll(graphSteps.stream().filter(GraphStep::returnsVertex).collect(Collectors.toList()));

            vertexStepsToInsertFilterAfter.forEach(s -> TraversalHelper.insertAfterStep(new TraversalFilterStep<>(traversal, vertexCriterion.asAdmin().clone()), s, traversal));
        }

        if (edgeCriterion != null) {
            final List edgeStepsToInsertFilterAfter = new ArrayList<>();
            edgeStepsToInsertFilterAfter.addAll(TraversalHelper.getStepsOfAssignableClass(AddEdgeStep.class, traversal));
            edgeStepsToInsertFilterAfter.addAll(graphSteps.stream().filter(GraphStep::returnsEdge).collect(Collectors.toList()));
            edgeStepsToInsertFilterAfter.addAll(vertexSteps.stream().filter(VertexStep::returnsEdge).collect(Collectors.toList()));

            edgeStepsToInsertFilterAfter.forEach(s -> TraversalHelper.insertAfterStep(new TraversalFilterStep<>(traversal, edgeCriterion.asAdmin().clone()), s, traversal));
        }

        // explode g.V().out() to g.V().outE().inV() only if there is an edge predicate otherwise
        vertexSteps.stream().filter(VertexStep::returnsVertex).forEach(s -> {
            if (null == edgeCriterion)
                TraversalHelper.insertAfterStep(new TraversalFilterStep<>(traversal, vertexCriterion.asAdmin().clone()), s, traversal);
            else {
                final VertexStep replacementVertexStep = new VertexStep(traversal, Edge.class, s.getDirection(), s.getEdgeLabels());
                Step intermediateFilterStep = null;
                if (s.getDirection() == Direction.BOTH)
                    intermediateFilterStep = new EdgeOtherVertexStep(traversal);
                else
                    intermediateFilterStep = new EdgeVertexStep(traversal, s.getDirection().opposite());

                TraversalHelper.replaceStep(s, replacementVertexStep, traversal);
                TraversalHelper.insertAfterStep(intermediateFilterStep, replacementVertexStep, traversal);
                TraversalHelper.insertAfterStep(new TraversalFilterStep<>(traversal, edgeCriterion.asAdmin().clone()), replacementVertexStep, traversal);

                if (vertexCriterion != null)
                    TraversalHelper.insertAfterStep(new TraversalFilterStep<>(traversal, vertexCriterion.asAdmin().clone()), intermediateFilterStep, traversal);
            }
        });
    }

    public Traversal getVertexCriterion() {
        return vertexCriterion;
    }

    public Traversal getEdgeCriterion() {
        return edgeCriterion;
    }

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

    public final static class Builder {

        private Traversal vertexCriterion = null;
        private Traversal edgeCriterion = null;

        private Builder() {
        }

        public Builder vertexCriterion(final Traversal predicate) {
            vertexCriterion = predicate;
            return this;
        }

        public Builder edgeCriterion(final Traversal predicate) {
            edgeCriterion = predicate;
            return this;
        }

        public SubgraphStrategy create() {
            if (null == edgeCriterion && null == vertexCriterion)
                throw new IllegalStateException("A subgraph must be filtered by an edge or vertex criterion");
            return new SubgraphStrategy(vertexCriterion, edgeCriterion);
        }
    }
}