Add JanusGraphPropertyMapStep which uses multi-query optimization

Fixes JanusGraph#2444 Signed-off-by: Oleksandr Porunov <alexandr.porunov@gmail.com>
porunov · Apr 16, 2023 · bb9ca1f · bb9ca1f
1 parent 536ac6e
commit bb9ca1f
Show file tree

Hide file tree

Showing 4 changed files with 410 additions and 10 deletions.
diff --git a/janusgraph-backend-testutils/src/main/java/org/janusgraph/graphdb/JanusGraphTest.java b/janusgraph-backend-testutils/src/main/java/org/janusgraph/graphdb/JanusGraphTest.java
@@ -4768,13 +4768,28 @@ public void testLimitBatchSizeForMultiQuery() {
         assertEquals(3, countBackendQueriesOfSize(barrierSize, profile.getMetrics()));
         assertEquals(1, countBackendQueriesOfSize(numV - 3 * barrierSize, profile.getMetrics()));
 
+        // test batching for `valueMap()`
+        traversal = () -> graph.traversal().V(cs).barrier(barrierSize).valueMap("foo");
+        assertEqualResultWithAndWithoutLimitBatchSize(traversal);
+        clopen(option(USE_MULTIQUERY), true, option(LIMIT_BATCH_SIZE), true);
+        profile = traversal.get().profile().next();
+        //assertEquals(3, countBackendQueriesOfSize(barrierSize, profile.getMetrics()));
+        //assertEquals(1, countBackendQueriesOfSize(numV - 3 * barrierSize, profile.getMetrics()));
+
         // test early abort with limit for `values()`
         traversal = () -> graph.traversal().V(cs).barrier(barrierSize).values("foo").limit(limit);
         assertEqualResultWithAndWithoutLimitBatchSize(traversal);
         clopen(option(USE_MULTIQUERY), true, option(LIMIT_BATCH_SIZE), true);
         profile = traversal.get().profile().next();
         assertEquals((int) Math.ceil((double) limit / barrierSize), countBackendQueriesOfSize(barrierSize, profile.getMetrics()));
 
+        // test early abort with limit for `valueMap()`
+        traversal = () -> graph.traversal().V(cs).barrier(barrierSize).valueMap("foo").limit(limit);
+        assertEqualResultWithAndWithoutLimitBatchSize(traversal);
+        clopen(option(USE_MULTIQUERY), true, option(LIMIT_BATCH_SIZE), true);
+        profile = traversal.get().profile().next();
+        //assertEquals((int) Math.ceil((double) limit / barrierSize), countBackendQueriesOfSize(barrierSize, profile.getMetrics()));
+
         // test batching with unlimited batch size
         traversal = () -> graph.traversal().V(bs).barrier(barrierSize).out();
         assertEqualResultWithAndWithoutLimitBatchSize(traversal);

diff --git a/...rc/main/java/org/janusgraph/graphdb/tinkerpop/optimize/step/JanusGraphPropertiesStep.java b/...rc/main/java/org/janusgraph/graphdb/tinkerpop/optimize/step/JanusGraphPropertiesStep.java
@@ -17,6 +17,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Iterators;
 import org.apache.tinkerpop.gremlin.process.traversal.Order;
+import org.apache.tinkerpop.gremlin.process.traversal.Traversal;
 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.TraversalParent;
@@ -26,6 +27,7 @@
 import org.apache.tinkerpop.gremlin.process.traversal.util.TraversalInterruptedException;
 import org.apache.tinkerpop.gremlin.structure.Element;
 import org.apache.tinkerpop.gremlin.structure.Property;
+import org.apache.tinkerpop.gremlin.structure.PropertyType;
 import org.apache.tinkerpop.gremlin.structure.Vertex;
 import org.apache.tinkerpop.gremlin.structure.util.StringFactory;
 import org.apache.tinkerpop.gremlin.structure.util.wrapped.WrappedVertex;
@@ -61,19 +63,33 @@ public class JanusGraphPropertiesStep<E> extends PropertiesStep<E> implements Ha
     private Map<JanusGraphVertex, Iterable<? extends JanusGraphProperty>> multiQueryResults = null;
     private QueryProfiler queryProfiler = QueryProfiler.NO_OP;
 
+    public JanusGraphPropertiesStep(final Traversal.Admin traversal, final PropertyType propertyType,
+                                    ArrayList<HasContainer> hasContainers, final String... propertyKeys) {
+        super(traversal, propertyType, propertyKeys);
+        this.hasContainers = hasContainers;
+    }
+
     public JanusGraphPropertiesStep(PropertiesStep<E> originalStep) {
-        super(originalStep.getTraversal(), originalStep.getReturnType(), originalStep.getPropertyKeys());
+        this(originalStep.getTraversal(),
+            originalStep.getReturnType(),
+            new ArrayList<>(),
+            originalStep.getPropertyKeys());
+
         originalStep.getLabels().forEach(this::addLabel);
 
-        if (originalStep instanceof JanusGraphPropertiesStep) {
-            JanusGraphPropertiesStep originalJanusGraphPropertiesStep = (JanusGraphPropertiesStep) originalStep;
-            this.useMultiQuery = originalJanusGraphPropertiesStep.useMultiQuery;
-            this.hasContainers = originalJanusGraphPropertiesStep.hasContainers;
-            this.limit = originalJanusGraphPropertiesStep.limit;
-        } else {
-            this.hasContainers = new ArrayList<>();
-            this.limit = Query.NO_LIMIT;
-        }
+        this.limit = Query.NO_LIMIT;
+    }
+
+    public JanusGraphPropertiesStep(JanusGraphPropertiesStep<E> originalStep) {
+        this(originalStep.getTraversal(),
+            originalStep.getReturnType(),
+            originalStep.hasContainers,
+            originalStep.getPropertyKeys());
+
+        originalStep.getLabels().forEach(this::addLabel);
+
+        this.useMultiQuery = originalStep.useMultiQuery;
+        this.limit = originalStep.limit;
     }
 
     @Override