implement nested-loop operation

Author: seut

sql/src/main/java/io/crate/action/job/ContextPreparer.java

@@ -30,6 +30,7 @@
 import io.crate.executor.transport.distributed.SingleBucketBuilder;
 import io.crate.jobs.CountContext;
 import io.crate.jobs.JobExecutionContext;
+import io.crate.jobs.NestedLoopContext;
 import io.crate.jobs.PageDownstreamContext;
 import io.crate.operation.PageDownstream;
 import io.crate.operation.PageDownstreamFactory;
@@ -46,6 +47,7 @@
 import io.crate.planner.node.dql.CollectNode;
 import io.crate.planner.node.dql.CountNode;
 import io.crate.planner.node.dql.MergeNode;
+import io.crate.planner.node.dql.join.NestedLoopNode;
 import io.crate.types.DataTypes;
 import org.elasticsearch.cluster.ClusterService;
 import org.elasticsearch.common.breaker.CircuitBreaker;
@@ -121,8 +123,8 @@ private PreparerContext(UUID jobId,
     private class InnerPreparer extends ExecutionNodeVisitor<PreparerContext, Void> {
 
         @Override
-        public Void visitCountNode(CountNode countNode, PreparerContext context) {
-            Map<String, Map<String, List<Integer>>> locations = countNode.routing().locations();
+        public Void visitCountNode(CountNode node, PreparerContext context) {
+            Map<String, Map<String, List<Integer>>> locations = node.routing().locations();
             if (locations == null) {
                 throw new IllegalArgumentException("locations are empty. Can't start count operation");
             }
@@ -137,15 +139,15 @@ public Void visitCountNode(CountNode countNode, PreparerContext context) {
                     countOperation,
                     singleBucketBuilder,
                     indexShardMap,
-                    countNode.whereClause()
+                    node.whereClause()
             );
             context.directResultFuture = singleBucketBuilder.result();
-            context.contextBuilder.addSubContext(countNode.executionNodeId(), countContext);
+            context.contextBuilder.addSubContext(node.executionNodeId(), countContext);
             return null;
         }
 
         @Override
-        public Void visitMergeNode(final MergeNode node, final PreparerContext context) {
+        public Void visitMergeNode(MergeNode node, PreparerContext context) {
             RamAccountingContext ramAccountingContext = RamAccountingContext.forExecutionNode(circuitBreaker, node);
             ResultProvider downstream = resultProviderFactory.createDownstream(node, node.jobId());
             Tuple<PageDownstream, FlatProjectorChain> pageDownstreamProjectorChain =
@@ -172,9 +174,9 @@ public Void visitMergeNode(final MergeNode node, final PreparerContext context)
         }
 
         @Override
-        public Void visitCollectNode(final CollectNode node, final PreparerContext context) {
+        public Void visitCollectNode(CollectNode node, PreparerContext context) {
             RamAccountingContext ramAccountingContext = RamAccountingContext.forExecutionNode(circuitBreaker, node);
-            ResultProvider downstream = collectOperation.createDownstream(node);
+            ResultProvider downstream = resultProviderFactory.createDownstream(node, node.jobId());
 
             if (ExecutionNodes.hasDirectResponseDownstream(node.downstreamNodes())) {
                 context.directResultFuture = downstream.result();
@@ -189,5 +191,23 @@ public Void visitCollectNode(final CollectNode node, final PreparerContext conte
             context.contextBuilder.addSubContext(node.executionNodeId(), jobCollectContext);
             return null;
         }
+
+        @Override
+        public Void visitNestedLoopNode(NestedLoopNode node, PreparerContext context) {
+            RamAccountingContext ramAccountingContext = RamAccountingContext.forExecutionNode(circuitBreaker, node);
+
+            ResultProvider downstream = resultProviderFactory.createDownstream(node, node.jobId());
+
+            NestedLoopContext nestedLoopContext = new NestedLoopContext(
+                    node,
+                    downstream,
+                    ramAccountingContext,
+                    pageDownstreamFactory,
+                    threadPool,
+                    streamerVisitor);
+
+            context.contextBuilder.addSubContext(node.executionNodeId(), nestedLoopContext);
+            return null;
+        }
     }
 }

sql/src/main/java/io/crate/executor/Job.java

@@ -43,7 +43,7 @@ public UUID id() {
         return id;
     }
 
-    public void addTasks(Collection<Task> tasks) {
+    public void addTasks(Collection<? extends Task> tasks) {
         this.tasks.addAll(tasks);
     }
 

sql/src/main/java/io/crate/executor/transport/ExecutionNodesTask.java

@@ -62,9 +62,6 @@ public class ExecutionNodesTask extends JobTask {
     private static final ESLogger LOGGER = Loggers.getLogger(ExecutionNodesTask.class);
 
     private final TransportJobAction transportJobAction;
-    private final List<List<ExecutionNode>> groupedExecutionNodes;
-    private final List<SettableFuture<TaskResult>> results;
-    private final boolean hasDirectResponse;
     private final ClusterService clusterService;
     private ContextPreparer contextPreparer;
     private final JobContextService jobContextService;
@@ -73,15 +70,13 @@ public class ExecutionNodesTask extends JobTask {
     private TransportCloseContextNodeAction transportCloseContextNodeAction;
     private final StreamerVisitor streamerVisitor;
     private final CircuitBreaker circuitBreaker;
-    private List<MergeNode> mergeNodes;
+
+    private final List<List<ExecutionNode>> groupedExecutionNodes = new ArrayList<>();
+    private final List<MergeNode> finalMergeNodes = new ArrayList<>();
+    private final List<SettableFuture<TaskResult>> results = new ArrayList<>();
+    private boolean hasDirectResponse;
     private boolean rowCountResult = false;
 
-    /**
-     * @param mergeNodes list of mergeNodes for the final merge operation on the handler.
-     *                  This may be null in the constructor but then it must be set using the
-     *                  {@link #mergeNodes(List)} setter before {@link #start()} is called.
-     *                   Multiple merge nodes are only occurring on bulk operations.
-     */
     protected ExecutionNodesTask(UUID jobId,
                                  ClusterService clusterService,
                                  ContextPreparer contextPreparer,
@@ -91,9 +86,7 @@ protected ExecutionNodesTask(UUID jobId,
                                  TransportJobAction transportJobAction,
                                  TransportCloseContextNodeAction transportCloseContextNodeAction,
                                  StreamerVisitor streamerVisitor,
-                                 CircuitBreaker circuitBreaker,
-                                 @Nullable List<MergeNode> mergeNodes,
-                                 List<List<ExecutionNode>> groupedExecutionNodes) {
+                                 CircuitBreaker circuitBreaker) {
         super(jobId);
         this.clusterService = clusterService;
         this.contextPreparer = contextPreparer;
@@ -103,21 +96,31 @@ protected ExecutionNodesTask(UUID jobId,
         this.transportCloseContextNodeAction = transportCloseContextNodeAction;
         this.streamerVisitor = streamerVisitor;
         this.circuitBreaker = circuitBreaker;
-        this.mergeNodes = mergeNodes;
         this.transportJobAction = transportJobAction;
-        this.groupedExecutionNodes = groupedExecutionNodes;
-        hasDirectResponse = hasDirectResponse(groupedExecutionNodes);
+    }
 
-        List<SettableFuture<TaskResult>> results = new ArrayList<>(groupedExecutionNodes.size());
-        for (int i = 0; i < groupedExecutionNodes.size(); i++) {
+
+    /**
+     * @param finalMergeNode a mergeNode for the final merge operation on the handler.
+     *                   Multiple merge nodes are only occurring on bulk operations.
+     */
+    public void addFinalMergeNode(MergeNode finalMergeNode) {
+        finalMergeNode.jobId(jobId());
+        finalMergeNodes.add(finalMergeNode);
+    }
+
+    public void addExecutionNode(int group, ExecutionNode executionNode) {
+        executionNode.jobId(jobId());
+        while (group >= groupedExecutionNodes.size()) {
             results.add(SettableFuture.<TaskResult>create());
+            groupedExecutionNodes.add(new ArrayList<ExecutionNode>());
         }
-        this.results = results;
-    }
+        List<ExecutionNode> executionNodes = groupedExecutionNodes.get(group);
 
-    public void mergeNodes(List<MergeNode> mergeNodes) {
-        assert this.mergeNodes == null : "can only overwrite mergeNodes if it was null";
-        this.mergeNodes = mergeNodes;
+        if (ExecutionNodes.hasDirectResponseDownstream(executionNode.downstreamNodes())) {
+            hasDirectResponse = true;
+        }
+        executionNodes.add(executionNode);
     }
 
     public void rowCountResult(boolean rowCountResult) {
@@ -126,7 +129,7 @@ public void rowCountResult(boolean rowCountResult) {
 
     @Override
     public void start() {
-        assert mergeNodes != null : "mergeNodes must not be null";
+        assert finalMergeNodes.size() == groupedExecutionNodes.size() : "groupedExecutionNodes and finalMergeNodes sizes must match";
 
         Map<String, Collection<ExecutionNode>> nodesByServer = ExecutionNodeGrouper.groupByServer(clusterService.state().nodes().localNodeId(), groupedExecutionNodes);
         RowDownstream rowDownstream;
@@ -135,21 +138,21 @@ public void start() {
         } else {
             rowDownstream = new QueryResultRowDownstream(results);
         }
-        Streamer<?>[] streamers = streamerVisitor.processExecutionNode(mergeNodes.get(0)).inputStreamers();
+        Streamer<?>[] streamers = streamerVisitor.processExecutionNode(finalMergeNodes.get(0)).inputStreamers();
         List<PageDownstreamContext> pageDownstreamContexts = new ArrayList<>(groupedExecutionNodes.size());
 
         for (int i = 0; i < groupedExecutionNodes.size(); i++) {
             RamAccountingContext ramAccountingContext = RamAccountingContext.forExecutionNode(
-                    circuitBreaker, mergeNodes.get(i));
+                    circuitBreaker, finalMergeNodes.get(i));
 
             PageDownstreamContext pageDownstreamContext = createPageDownstreamContext(ramAccountingContext, streamers,
-                    mergeNodes.get(i), groupedExecutionNodes.get(i), rowDownstream);
+                    finalMergeNodes.get(i), groupedExecutionNodes.get(i), rowDownstream);
             if (nodesByServer.size() == 0) {
                 pageDownstreamContext.finish();
                 continue;
             }
             if (!hasDirectResponse) {
-                createLocalContextAndStartOperation(pageDownstreamContext, nodesByServer, mergeNodes.get(i).executionNodeId());
+                createLocalContextAndStartOperation(pageDownstreamContext, nodesByServer, finalMergeNodes.get(i).executionNodeId());
             }
             pageDownstreamContexts.add(pageDownstreamContext);
         }
@@ -262,6 +265,11 @@ public void onFailure(@Nonnull Throwable t) {
 
     @Override
     public List<? extends ListenableFuture<TaskResult>> result() {
+        if (results.size() != groupedExecutionNodes.size()) {
+            for (int i = 0; i < groupedExecutionNodes.size(); i++) {
+                results.add(SettableFuture.<TaskResult>create());
+            }
+        }
         return results;
     }