Merge pull request #764 from bserdar/streams

Streams

Author: bserdar

core-api/src/main/java/com/redhat/lightblue/Response.java

@@ -210,6 +210,24 @@ public void setEntityData(JsonNode node) {
         }
     }
 
+    /**
+     * Adds a document, or array of documents to the entity data
+     */
+    public void addEntityData(JsonNode doc) {
+        if(doc!=null) {
+            if(entityData==null) {
+                entityData=JsonNodeFactory.instance.arrayNode();
+            } 
+            if(doc instanceof ArrayNode) {
+                for(Iterator<JsonNode> itr=doc.elements();itr.hasNext();) {
+                    ((ArrayNode)entityData).add(itr.next());
+                }
+            } else {
+                ((ArrayNode)entityData).add(doc);
+            }
+        }
+    }
+
     /**
      * Metadata list for documents in entityData. If there are more
      * than one documents, the entitydata and metadata indexes match.

crud/src/main/java/com/redhat/lightblue/assoc/CompositeFindImpl.java

@@ -21,6 +21,8 @@
 import java.util.Set;
 import java.util.HashSet;
 import java.util.List;
+import java.util.ArrayList;
+import java.util.Arrays;
 
 import java.util.concurrent.Executors;
 
@@ -32,6 +34,7 @@
 import com.redhat.lightblue.crud.CRUDFindResponse;
 import com.redhat.lightblue.crud.CRUDFindRequest;
 import com.redhat.lightblue.crud.DocCtx;
+import com.redhat.lightblue.crud.ListDocumentStream;
 
 import com.redhat.lightblue.metadata.CompositeMetadata;
 
@@ -48,6 +51,8 @@
 import com.redhat.lightblue.assoc.ep.StepResult;
 import com.redhat.lightblue.assoc.ep.ResultDocument;
 import com.redhat.lightblue.assoc.ep.ExecutionContext;
+import com.redhat.lightblue.assoc.ep.MakeDocCtx;
+import com.redhat.lightblue.assoc.ep.StepResultDocumentStream;
 
 import com.redhat.lightblue.util.JsonDoc;
 
@@ -140,7 +145,7 @@ public void explain(OperationContext ctx,
         initialize(ctx,req);
         ExecutionContext executionContext = new ExecutionContext(ctx,null);
         JsonDoc doc=new JsonDoc(executionPlan.explain(executionContext));
-        ctx.addDocument(doc);        
+        ctx.setDocumentStream(new ListDocumentStream<DocCtx>(Arrays.asList(new DocCtx(doc))));
     }
 
     @Override
@@ -158,7 +163,7 @@ public CRUDFindResponse find(OperationContext ctx,
                 Executors.newWorkStealingPool(parallelism));
         try {
             StepResult<ResultDocument> results = executionPlan.getResults(executionContext);
-            results.stream().map(d -> new DocCtx(d.getDoc())).forEach(d -> ctx.addDocument(d));
+            ctx.setDocumentStream(new StepResultDocumentStream(new MakeDocCtx(results)));            	
             response.setSize(executionContext.getMatchCount());
             LOGGER.debug("Composite find: end");
             return response;

crud/src/main/java/com/redhat/lightblue/assoc/ep/AbstractSearchStep.java

@@ -20,6 +20,11 @@
 
 import java.util.List;
 import java.util.ArrayList;
+import java.util.Spliterator;
+import java.util.Spliterators;
+
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -40,6 +45,7 @@
 
 import com.redhat.lightblue.crud.CRUDFindRequest;
 import com.redhat.lightblue.crud.CRUDFindResponse;
+import com.redhat.lightblue.crud.DocumentStream;
 import com.redhat.lightblue.assoc.Conjunct;
 
 /**
@@ -90,11 +96,15 @@ public void setQueries(List<Conjunct> conjuncts) {
 
     @Override
     public StepResult<ResultDocument> getResults(ExecutionContext ctx) {
-        List<ResultDocument> list = getSearchResults(ctx);
-        return new ListStepResult(list);
+        return new StepResult<ResultDocument>() {
+            @Override
+            public Stream<ResultDocument> stream() {
+                return StreamSupport.stream(Spliterators.spliteratorUnknownSize(getSearchResults(ctx),Spliterator.IMMUTABLE),false);
+            }
+        };
     }
 
-    protected abstract List<ResultDocument> getSearchResults(ExecutionContext ctx);
+    protected abstract DocumentStream<ResultDocument> getSearchResults(ExecutionContext ctx);
 
     public OperationContext search(ExecutionContext ctx, CRUDFindRequest req) {
         return search(block, ctx, req);
@@ -125,7 +135,7 @@ public OperationContext search(ExecutionBlock block, ExecutionContext ctx, CRUDF
             }
             LOGGER.debug("execute {}: returning {} documents",
                     block.getQueryPlanNode().getName(),
-                    searchCtx.getDocuments().size());
+                         response.getSize());
         }
         return searchCtx;
     }

crud/src/main/java/com/redhat/lightblue/assoc/ep/Assemble.java

@@ -46,7 +46,6 @@
 import com.redhat.lightblue.assoc.AnalyzeQuery;
 
 import com.redhat.lightblue.mindex.MemDocIndex;
-import com.redhat.lightblue.mindex.GetIndexKeySpec;
 import com.redhat.lightblue.mindex.GetIndexLookupSpec;
 import com.redhat.lightblue.mindex.KeySpec;
 import com.redhat.lightblue.mindex.LookupSpec;
@@ -350,7 +349,7 @@ private void associateDocsWithIndex(CompositeMetadata childMetadata,
         LOGGER.debug("Lookup spec:"+ls);
         List<ResultDocument> docs=reorder(childDocs,childIndex.find(ls));
         ArrayNode destNode=null;
-        for (ResultDocument childDoc : childDocs) {
+        for (ResultDocument childDoc : docs) {
             if (qeval.evaluate(childDoc.getDoc()).getResult()) {
                 destNode=ensureDestNodeExists(parentDoc,destNode,destFieldName);
                 destNode.add(childDoc.getDoc().getRoot());

crud/src/main/java/com/redhat/lightblue/assoc/ep/Copy.java

@@ -26,6 +26,8 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 
+import com.redhat.lightblue.crud.DocumentStream;
+
 /**
  * Copies the results from another step
  */
@@ -59,7 +61,7 @@ public Stream<ResultDocument> stream() {
     }
 
     @Override
-    protected final List<ResultDocument> getSearchResults(ExecutionContext ctx) {
+    protected final DocumentStream<ResultDocument> getSearchResults(ExecutionContext ctx) {
         // This should be called at all
         throw new IllegalStateException();
     }

crud/src/main/java/com/redhat/lightblue/assoc/ep/JoinSearch.java

@@ -20,6 +20,10 @@
 
 import java.util.List;
 import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+import java.util.function.Consumer;
+
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
@@ -32,6 +36,8 @@
 
 import com.redhat.lightblue.mediator.OperationContext;
 import com.redhat.lightblue.crud.CRUDFindRequest;
+import com.redhat.lightblue.crud.DocumentStream;
+import com.redhat.lightblue.crud.DocCtx;
 
 /**
  * Performs searches based on the n-tuple of result documents obtained from the
@@ -51,56 +57,105 @@ public JoinSearch(ExecutionBlock block, Source<JoinTuple> source) {
     }
 
     @Override
-    protected List<ResultDocument> getSearchResults(ExecutionContext ctx) {
-        BatchQueryExecutor executor = new BatchQueryExecutor(256, ctx);
-        source.getStep().getResults(ctx).stream().forEach(x -> executor.add(x));
-        return executor.getResults();
+    protected DocumentStream<ResultDocument> getSearchResults(final ExecutionContext ctx) {
+        return new BatchQueryIterator(256,ctx);
     }
 
-    public class BatchQueryExecutor {
+    /**
+     * Streaming batch query executor/iterator
+     *
+     * When the results are retrieved from the stream, executes a
+     * batch of queries, computes results, and streams them to the
+     * caller
+     */
+    private class BatchQueryIterator implements DocumentStream<ResultDocument> {
         private final int batchSize;
-        private List<JoinTuple> jtBatch;
-        ;
-        private List<QueryExpression> qBatch;
         private final ExecutionContext ctx;
-        private List<ResultDocument> docs = new ArrayList<>();
-
-        public BatchQueryExecutor(int batchSize, ExecutionContext ctx) {
-            this.batchSize = batchSize;
-            this.jtBatch = new ArrayList<>(batchSize);
-            this.qBatch = new ArrayList<>(batchSize);
-            this.ctx = ctx;
+        private final Iterator<JoinTuple> sourceStream;
+        private final ArrayList<Consumer<ResultDocument>> listeners=new ArrayList<>();
+
+        private DocumentStream<DocCtx> currentIterator;
+        private boolean done=false; // Are we still iterating, or are we done?
+        
+        public BatchQueryIterator(int batchSize,ExecutionContext ctx) {
+            this.batchSize=batchSize;
+            this.ctx=ctx;
+            sourceStream=source.getStep().getResults(ctx).stream().iterator();
         }
 
-        public void add(JoinTuple tuple) {
-            jtBatch.add(tuple);
-            qBatch.addAll(Searches.writeQueriesForJoinTuple(tuple, block));
-            if (qBatch.size() >= batchSize) {
-                executeBatch();
-                qBatch = new ArrayList<>(batchSize);
-                jtBatch = new ArrayList<>(batchSize);
+        @Override
+        public boolean hasNext() {
+            if(!done) {
+                if(currentIterator==null||!currentIterator.hasNext())
+                    retrieveNextBatch();
+                if(done)
+                    return false;
+                else
+                    return currentIterator.hasNext();
+            } else {
+                return false;
             }
         }
-
-        public void executeBatch() {
-            if (!qBatch.isEmpty()) {
-                QueryExpression q = Searches.combine(NaryLogicalOperator._or, qBatch);
-                CRUDFindRequest findRequest = new CRUDFindRequest();
-                findRequest.setQuery(Searches.and(q, query));
-                findRequest.setProjection(projection);
-                findRequest.setSort(sort);
-                findRequest.setFrom(from);
-                findRequest.setTo(to);
-                OperationContext opctx = search(ctx, findRequest);
-                opctx.getDocuments().stream().
-                        forEach(doc -> docs.add(new ResultDocument(block, doc.getOutputDocument())));
+        
+        @Override
+        public ResultDocument next() {
+            if(!done) {
+                if(currentIterator==null||!currentIterator.hasNext())
+                    retrieveNextBatch();
+                if(currentIterator!=null) {
+                    ResultDocument doc=new ResultDocument(block,currentIterator.next().getOutputDocument());
+                    for(Consumer<ResultDocument> l:listeners)
+                        l.accept(doc);
+                    return doc;
+                }
             }
+            throw new NoSuchElementException();
+        }
+
+        @Override
+        public void close() {
+            if(currentIterator!=null)
+                currentIterator.close();
         }
 
-        public List<ResultDocument> getResults() {
-            executeBatch();
-            return docs;
+        @Override
+        public void addListener(Consumer<ResultDocument> listener) {
+            listeners.add(listener);
         }
+        
+        private void retrieveNextBatch() {
+            do {
+                int n=0;
+                ArrayList<QueryExpression> qBatch=new ArrayList<>(batchSize);
+                if(currentIterator!=null) {
+                    currentIterator.close();
+                    currentIterator=null;
+                }
+                while(sourceStream.hasNext()&&n<batchSize) {
+                    JoinTuple t=sourceStream.next();
+                    qBatch.addAll(Searches.writeQueriesForJoinTuple(t, block));
+                    n++;
+                }
+                if(!qBatch.isEmpty()) {
+                    QueryExpression q = Searches.combine(NaryLogicalOperator._or, qBatch);
+                    CRUDFindRequest findRequest = new CRUDFindRequest();
+                    findRequest.setQuery(Searches.and(q, query));
+                    findRequest.setProjection(projection);
+                    findRequest.setSort(sort);
+                    findRequest.setFrom(from);
+                    findRequest.setTo(to);
+                    OperationContext opctx = search(ctx, findRequest);
+                    currentIterator=opctx.getDocumentStream();
+                    if(!currentIterator.hasNext()) {
+                        currentIterator.close();
+                        currentIterator=null;
+                    }
+                } else {
+                    done=true;
+                }
+            } while(!done&&currentIterator==null);
+        }   
+            
     }
 
     @Override

crud/src/main/java/com/redhat/lightblue/assoc/ep/MakeDocCtx.java

@@ -0,0 +1,22 @@
+package com.redhat.lightblue.assoc.ep;
+
+import java.util.stream.Stream;
+
+import com.redhat.lightblue.crud.DocCtx;
+
+/**
+ * Decorator interface that gets a StepResult<ResultDocument> and return StepResult<DocCtx>
+ */
+public class MakeDocCtx implements StepResult<DocCtx> {
+
+    private final StepResult<ResultDocument> result;
+
+    public MakeDocCtx(StepResult<ResultDocument> result) {
+        this.result=result;
+    }
+
+    @Override
+    public Stream<DocCtx> stream() {
+        return result.stream().map(d->new DocCtx(d.getDoc()));
+    }
+}