[GIE-IR] support sending back a page of results if batch size met (#2076

)
alibaba · Sep 27, 2022 · 2d57a3d · 2d57a3d
1 parent a7462ae
commit 2d57a3d
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 14 deletions.
diff --git a/...in/java/com/alibaba/graphscope/gremlin/integration/result/GremlinTestResultProcessor.java b/...in/java/com/alibaba/graphscope/gremlin/integration/result/GremlinTestResultProcessor.java
@@ -20,7 +20,6 @@
 import com.alibaba.graphscope.gremlin.result.GremlinResultProcessor;
 import com.google.common.collect.ImmutableMap;
 
-import org.apache.tinkerpop.gremlin.driver.message.ResponseStatusCode;
 import org.apache.tinkerpop.gremlin.process.remote.traversal.DefaultRemoteTraverser;
 import org.apache.tinkerpop.gremlin.server.Context;
 import org.apache.tinkerpop.gremlin.structure.Edge;
@@ -48,18 +47,6 @@ public GremlinTestResultProcessor(
         this.cachedProperties = testGraph.getProperties();
     }
 
-    @Override
-    public void finish() {
-        synchronized (this) {
-            if (!locked) {
-                logger.debug("process finish");
-                formatResultIfNeed();
-                writeResultList(writeResult, resultCollectors, ResponseStatusCode.SUCCESS);
-                locked = true;
-            }
-        }
-    }
-
     @Override
     protected void formatResultIfNeed() {
         super.formatResultIfNeed();

diff --git a/.../compiler/src/main/java/com/alibaba/graphscope/gremlin/result/GremlinResultProcessor.java b/.../compiler/src/main/java/com/alibaba/graphscope/gremlin/result/GremlinResultProcessor.java
@@ -24,10 +24,12 @@
 import io.netty.channel.ChannelHandlerContext;
 
 import org.apache.tinkerpop.gremlin.driver.MessageSerializer;
+import org.apache.tinkerpop.gremlin.driver.Tokens;
 import org.apache.tinkerpop.gremlin.driver.message.RequestMessage;
 import org.apache.tinkerpop.gremlin.driver.message.ResponseMessage;
 import org.apache.tinkerpop.gremlin.driver.message.ResponseStatusCode;
 import org.apache.tinkerpop.gremlin.server.Context;
+import org.apache.tinkerpop.gremlin.server.Settings;
 import org.apache.tinkerpop.gremlin.server.handler.Frame;
 import org.apache.tinkerpop.gremlin.server.handler.StateKey;
 import org.apache.tinkerpop.gremlin.server.op.standard.StandardOpProcessor;
@@ -40,20 +42,38 @@
 public class GremlinResultProcessor extends StandardOpProcessor implements ResultProcessor {
     private static Logger logger = LoggerFactory.getLogger(GremlinResultProcessor.class);
     protected Context writeResult;
-    protected List<Object> resultCollectors = new ArrayList<>();
+    protected List<Object> resultCollectors;
+    protected int resultCollectorsBatchSize;
     protected boolean locked = false;
     protected ResultParser resultParser;
 
     public GremlinResultProcessor(Context writeResult, ResultParser resultParser) {
         this.writeResult = writeResult;
         this.resultParser = resultParser;
+        RequestMessage msg = writeResult.getRequestMessage();
+        Settings settings = writeResult.getSettings();
+        // init batch size from resultIterationBatchSize in conf/gremlin-server.yaml,
+        // or args in RequestMessage which is originate from gremlin client
+        this.resultCollectorsBatchSize =
+                (Integer)
+                        msg.optionalArgs(Tokens.ARGS_BATCH_SIZE)
+                                .orElse(settings.resultIterationBatchSize);
+        this.resultCollectors = new ArrayList<>(this.resultCollectorsBatchSize);
     }
 
     @Override
     public void process(PegasusClient.JobResponse response) {
         synchronized (this) {
             try {
                 if (!locked) {
+                    // send back a page of results if batch size is met and then reset the
+                    // resultCollectors
+                    if (this.resultCollectors.size() >= this.resultCollectorsBatchSize) {
+                        formatResultIfNeed();
+                        writeResultList(
+                                writeResult, resultCollectors, ResponseStatusCode.PARTIAL_CONTENT);
+                        this.resultCollectors = new ArrayList<>(this.resultCollectorsBatchSize);
+                    }
                     resultCollectors.addAll(resultParser.parseFrom(response));
                 }
             } catch (Exception e) {