Merge pull request #159 from twilmes/issue_155_hasId

Issue 155: Updated JanusGraphStepStrategy to fold hasId steps into JanusGraphStep

janusgraph-core/src/main/java/org/janusgraph/graphdb/tinkerpop/optimize/HasStepFolder.java

@@ -14,6 +14,8 @@
 
 package org.janusgraph.graphdb.tinkerpop.optimize;
 
+import org.apache.tinkerpop.gremlin.process.traversal.step.map.GraphStep;
+import org.apache.tinkerpop.gremlin.process.traversal.step.map.NoOpBarrierStep;
 import org.janusgraph.core.Cardinality;
 import org.janusgraph.core.PropertyKey;
 import org.janusgraph.core.JanusGraphTransaction;
@@ -33,6 +35,7 @@
 
 import org.javatuples.Pair;
 
+import java.util.Arrays;
 import java.util.Comparator;
 import java.util.List;
 
@@ -41,26 +44,26 @@
  */
 public interface HasStepFolder<S, E> extends Step<S, E> {
 
-    public void addAll(Iterable<HasContainer> hasContainers);
+    void addAll(Iterable<HasContainer> hasContainers);
 
-    public void orderBy(String key, Order order);
+    void orderBy(String key, Order order);
 
-    public void setLimit(int limit);
+    void setLimit(int limit);
 
-    public int getLimit();
+    int getLimit();
 
-    public static boolean validJanusGraphHas(HasContainer has) {
+    static boolean validJanusGraphHas(HasContainer has) {
         return JanusGraphPredicate.Converter.supports(has.getBiPredicate());
     }
 
-    public static boolean validJanusGraphHas(Iterable<HasContainer> has) {
+    static boolean validJanusGraphHas(Iterable<HasContainer> has) {
         for (HasContainer h : has) {
             if (!validJanusGraphHas(h)) return false;
         }
         return true;
     }
 
-    public static boolean validJanusGraphOrder(OrderGlobalStep ostep, Traversal rootTraversal,
+    static boolean validJanusGraphOrder(OrderGlobalStep ostep, Traversal rootTraversal,
                                           boolean isVertexOrder) {
         for (Pair<Traversal.Admin<Object, Comparable>, Comparator<Comparable>> comp : (List<Pair<Traversal.Admin<Object, Comparable>, Comparator<Comparable>>>) ostep.getComparators()) {
             if (!(comp.getValue1() instanceof ElementValueComparator)) return false;
@@ -76,8 +79,29 @@ public static boolean validJanusGraphOrder(OrderGlobalStep ostep, Traversal root
         return true;
     }
 
-    public static void foldInHasContainer(final HasStepFolder janusgraphStep, final Traversal.Admin<?, ?> traversal) {
+    static void foldInIds(final HasStepFolder janusgraphStep, final Traversal.Admin<?, ?> traversal) {
+        Step<?, ?> currentStep = janusgraphStep.getNextStep();
+        while (true) {
+            if (currentStep instanceof HasContainerHolder) {
+                for (final HasContainer hasContainer : ((HasContainerHolder) currentStep).getHasContainers()) {
+                    if (GraphStep.processHasContainerIds((GraphStep) janusgraphStep, hasContainer)) {
+                        currentStep.getLabels().forEach(janusgraphStep::addLabel);
+                        traversal.removeStep(currentStep);
+                    }
+                }
+            }
+            else if (currentStep instanceof IdentityStep) {
+                // do nothing, has no impact
+            } else if (currentStep instanceof NoOpBarrierStep) {
+                // do nothing, has no impact
+            } else {
+                break;
+            }
+            currentStep = currentStep.getNextStep();
+        }
+    }
 
+    static void foldInHasContainer(final HasStepFolder janusgraphStep, final Traversal.Admin<?, ?> traversal) {
         Step<?, ?> currentStep = janusgraphStep.getNextStep();
         while (true) {
             if (currentStep instanceof HasContainerHolder) {
@@ -89,6 +113,8 @@ public static void foldInHasContainer(final HasStepFolder janusgraphStep, final
                 }
             } else if (currentStep instanceof IdentityStep) {
                 // do nothing, has no impact
+            } else if (currentStep instanceof NoOpBarrierStep) {
+                // do nothing, has no impact
             } else {
                 break;
             }
@@ -105,7 +131,7 @@ public static void foldInHasContainer(final HasStepFolder janusgraphStep, final
 //        } else return false;
 //    }
 
-    public static void foldInOrder(final HasStepFolder janusgraphStep, final Traversal.Admin<?, ?> traversal,
+    static void foldInOrder(final HasStepFolder janusgraphStep, final Traversal.Admin<?, ?> traversal,
                                    final Traversal<?, ?> rootTraversal, boolean isVertexOrder) {
         Step<?, ?> currentStep = janusgraphStep.getNextStep();
         OrderGlobalStep<?, ?> lastOrder = null;
@@ -120,6 +146,8 @@ public static void foldInOrder(final HasStepFolder janusgraphStep, final Travers
                 // do nothing, can be skipped
             } else if (currentStep instanceof HasStep) {
                 // do nothing, can be skipped
+            } else if (currentStep instanceof NoOpBarrierStep) {
+                // do nothing, can be skipped
             } else {
                 break;
             }
@@ -139,7 +167,7 @@ public static void foldInOrder(final HasStepFolder janusgraphStep, final Travers
         }
     }
 
-    public static class OrderEntry {
+    class OrderEntry {
 
         public final String key;
         public final Order order;
@@ -150,7 +178,7 @@ public OrderEntry(String key, Order order) {
         }
     }
 
-    public static <E extends Ranging> void foldInRange(final HasStepFolder janusgraphStep, final Traversal.Admin<?, ?> traversal) {
+    static <E extends Ranging> void foldInRange(final HasStepFolder janusgraphStep, final Traversal.Admin<?, ?> traversal) {
         Step<?, ?> nextStep = JanusGraphTraversalUtil.getNextNonIdentityStep(janusgraphStep);
 
         if (nextStep instanceof RangeGlobalStep) {

janusgraph-core/src/main/java/org/janusgraph/graphdb/tinkerpop/optimize/JanusGraphStep.java

@@ -15,6 +15,7 @@
 package org.janusgraph.graphdb.tinkerpop.optimize;
 
 import com.google.common.collect.Iterables;
+import org.apache.tinkerpop.gremlin.structure.Graph;
 import org.janusgraph.core.JanusGraphQuery;
 import org.janusgraph.core.JanusGraphTransaction;
 import org.janusgraph.graphdb.query.BaseQuery;
@@ -33,7 +34,9 @@
 import org.apache.tinkerpop.gremlin.structure.util.StringFactory;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
 
 /**
@@ -51,6 +54,10 @@ public JanusGraphStep(final GraphStep<S, E> originalStep) {
         super(originalStep.getTraversal(), originalStep.getReturnClass(), originalStep.isStartStep(), originalStep.getIds());
         originalStep.getLabels().forEach(this::addLabel);
         this.setIteratorSupplier(() -> {
+            if (this.ids != null && this.ids.length > 0) {
+                final Graph graph = (Graph)traversal.asAdmin().getGraph().get();
+                return iteratorList((Iterator)graph.vertices(this.ids));
+            }
             JanusGraphTransaction tx = JanusGraphTraversalUtil.getTx(traversal);
             JanusGraphQuery query = tx.query();
             for (HasContainer condition : hasContainers) {
@@ -65,7 +72,8 @@ public JanusGraphStep(final GraphStep<S, E> originalStep) {
 
     @Override
     public String toString() {
-        return this.hasContainers.isEmpty() ? super.toString() : StringFactory.stepString(this, this.hasContainers);
+        return this.hasContainers.isEmpty() ?
+                super.toString() : StringFactory.stepString(this, Arrays.toString(this.ids), this.hasContainers);
     }
 
     @Override
@@ -102,5 +110,15 @@ public List<HasContainer> getHasContainers() {
     public void addHasContainer(final HasContainer hasContainer) {
         this.addAll(Collections.singleton(hasContainer));
     }
+
+    private <E extends Element> Iterator<E> iteratorList(final Iterator<E> iterator) {
+        final List<E> list = new ArrayList<E>();
+        while (iterator.hasNext()) {
+            final E e = iterator.next();
+            if (HasContainer.testAll(e, this.hasContainers))
+                list.add(e);
+        }
+        return list.iterator();
+    }
 }
 

janusgraph-core/src/main/java/org/janusgraph/graphdb/tinkerpop/optimize/JanusGraphStepStrategy.java

@@ -46,6 +46,7 @@ public void apply(final Traversal.Admin<?, ?> traversal) {
                 //Try to optimize for index calls
                 final JanusGraphStep<?, ?> janusGraphStep = new JanusGraphStep<>(originalGraphStep);
                 TraversalHelper.replaceStep(originalGraphStep, (Step) janusGraphStep, traversal);
+                HasStepFolder.foldInIds(janusGraphStep, traversal);
                 HasStepFolder.foldInHasContainer(janusGraphStep, traversal);
                 HasStepFolder.foldInOrder(janusGraphStep, traversal, traversal, janusGraphStep.returnsVertex());
                 HasStepFolder.foldInRange(janusGraphStep, traversal);