package com.thinkaurelius.titan.graphdb.tinkerpop.optimize; import com.google.common.collect.Iterables; import com.thinkaurelius.titan.core.BaseVertexQuery; import com.thinkaurelius.titan.core.TitanElement; import com.thinkaurelius.titan.core.TitanMultiVertexQuery; import com.thinkaurelius.titan.core.TitanVertex; import com.thinkaurelius.titan.core.TitanVertexQuery; import com.thinkaurelius.titan.graphdb.query.BaseQuery; import com.thinkaurelius.titan.graphdb.query.Query; import com.thinkaurelius.titan.graphdb.query.TitanPredicate; import com.thinkaurelius.titan.graphdb.query.profile.QueryProfiler; import com.thinkaurelius.titan.graphdb.query.vertex.BasicVertexCentricQueryBuilder; import com.thinkaurelius.titan.graphdb.tinkerpop.profile.TP3ProfileWrapper; import org.apache.tinkerpop.gremlin.process.traversal.Order; import org.apache.tinkerpop.gremlin.process.traversal.Traverser; import org.apache.tinkerpop.gremlin.process.traversal.step.Profiling; import org.apache.tinkerpop.gremlin.process.traversal.step.map.VertexStep; import org.apache.tinkerpop.gremlin.process.traversal.step.util.HasContainer; import org.apache.tinkerpop.gremlin.process.traversal.util.FastNoSuchElementException; import org.apache.tinkerpop.gremlin.process.traversal.util.MutableMetrics; import org.apache.tinkerpop.gremlin.structure.Element; import org.apache.tinkerpop.gremlin.structure.Vertex; import org.apache.tinkerpop.gremlin.structure.util.StringFactory; import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.Map; /** * @author Matthias Broecheler ([email protected]) */ public class TitanVertexStep<E extends Element> extends VertexStep<E> implements HasStepFolder<Vertex, E>, Profiling, MultiQueriable<Vertex,E> { public TitanVertexStep(VertexStep<E> originalStep) { super(originalStep.getTraversal(), originalStep.getReturnClass(), originalStep.getDirection(), originalStep.getEdgeLabels()); originalStep.getLabels().forEach(this::addLabel); this.hasContainers = new ArrayList<>(); this.limit = Query.NO_LIMIT; } private boolean initialized = false; private boolean useMultiQuery = false; private Map<TitanVertex, Iterable<? extends TitanElement>> multiQueryResults = null; private QueryProfiler queryProfiler = QueryProfiler.NO_OP; @Override public void setUseMultiQuery(boolean useMultiQuery) { this.useMultiQuery = useMultiQuery; } public <Q extends BaseVertexQuery> Q makeQuery(Q query) { query.labels(getEdgeLabels()); query.direction(getDirection()); for (HasContainer condition : hasContainers) { query.has(condition.getKey(), TitanPredicate.Converter.convert(condition.getBiPredicate()), condition.getValue()); } for (OrderEntry order : orders) query.orderBy(order.key, order.order); if (limit != BaseQuery.NO_LIMIT) query.limit(limit); ((BasicVertexCentricQueryBuilder) query).profiler(queryProfiler); return query; } @SuppressWarnings("deprecation") private void initialize() { assert !initialized; initialized = true; if (useMultiQuery) { if (!starts.hasNext()) throw FastNoSuchElementException.instance(); TitanMultiVertexQuery mquery = TitanTraversalUtil.getTx(traversal).multiQuery(); List<Traverser.Admin<Vertex>> vertices = new ArrayList<>(); starts.forEachRemaining(v -> { vertices.add(v); mquery.addVertex(v.get()); }); starts.add(vertices.iterator()); assert vertices.size() > 0; makeQuery(mquery); multiQueryResults = (Vertex.class.isAssignableFrom(getReturnClass())) ? mquery.vertices() : mquery.edges(); } } @Override protected Traverser<E> processNextStart() { if (!initialized) initialize(); return super.processNextStart(); } @Override protected Iterator<E> flatMap(final Traverser.Admin<Vertex> traverser) { if (useMultiQuery) { assert multiQueryResults != null; return (Iterator<E>) multiQueryResults.get(traverser.get()).iterator(); } else { TitanVertexQuery query = makeQuery((TitanTraversalUtil.getTitanVertex(traverser)).query()); return (Vertex.class.isAssignableFrom(getReturnClass())) ? query.vertices().iterator() : query.edges().iterator(); } } @Override public void reset() { super.reset(); this.initialized = false; } @Override public TitanVertexStep<E> clone() { final TitanVertexStep<E> clone = (TitanVertexStep<E>) super.clone(); clone.initialized = false; return clone; } /* ===== HOLDER ===== */ private final List<HasContainer> hasContainers; private int limit = BaseQuery.NO_LIMIT; private List<OrderEntry> orders = new ArrayList<>(); @Override public void addAll(Iterable<HasContainer> has) { Iterables.addAll(hasContainers, has); } @Override public void orderBy(String key, Order order) { orders.add(new OrderEntry(key, order)); } @Override public void setLimit(int limit) { this.limit = limit; } @Override public int getLimit() { return this.limit; } @Override public String toString() { return this.hasContainers.isEmpty() ? super.toString() : StringFactory.stepString(this, this.hasContainers); } @Override public void setMetrics(MutableMetrics metrics) { queryProfiler = new TP3ProfileWrapper(metrics); } }