Examples with CollectingRowConsumer io.crate.data.CollectingRowConsumer used on opensource projects

Example 1 with CollectingRowConsumer

use of io.crate.data.CollectingRowConsumer in project crate by crate.

the class RemoteCollectorFactory method retrieveRows.

private CompletableFuture<List<Row>> retrieveRows(ShardRouting activePrimaryRouting, RoutedCollectPhase collectPhase, CollectTask collectTask, ShardCollectorProviderFactory shardCollectorProviderFactory) {
    Collector<Row, ?, List<Object[]>> listCollector = Collectors.mapping(Row::materialize, Collectors.toList());
    CollectingRowConsumer<?, List<Object[]>> consumer = new CollectingRowConsumer<>(listCollector);
    String nodeId = activePrimaryRouting.currentNodeId();
    String localNodeId = clusterService.localNode().getId();
    if (localNodeId.equalsIgnoreCase(nodeId)) {
        var indexShard = indicesService.indexServiceSafe(activePrimaryRouting.index()).getShard(activePrimaryRouting.shardId().id());
        var collectorProvider = shardCollectorProviderFactory.create(indexShard);
        CompletableFuture<BatchIterator<Row>> it;
        try {
            it = collectorProvider.getFutureIterator(collectPhase, consumer.requiresScroll(), collectTask);
        } catch (Exception e) {
            return Exceptions.rethrowRuntimeException(e);
        }
        it.whenComplete(consumer);
    } else {
        UUID childJobId = UUIDs.dirtyUUID();
        RemoteCollector remoteCollector = new RemoteCollector(childJobId, collectTask.txnCtx().sessionSettings(), localNodeId, nodeId, transportActionProvider.transportJobInitAction(), transportActionProvider.transportKillJobsNodeAction(), searchTp, tasksService, collectTask.getRamAccounting(), consumer, createRemoteCollectPhase(childJobId, collectPhase, activePrimaryRouting.shardId(), nodeId));
        remoteCollector.doCollect();
    }
    return consumer.completionFuture().thenApply(rows -> Lists2.mapLazy(rows, Buckets.arrayToSharedRow()));
}

Example 2 with CollectingRowConsumer

use of io.crate.data.CollectingRowConsumer in project crate by crate.

the class JobLauncher method executeBulk.

public List<CompletableFuture<Long>> executeBulk(TransactionContext txnCtx) {
    Iterable<NodeOperation> nodeOperations = nodeOperationTrees.stream().flatMap(opTree -> opTree.nodeOperations().stream())::iterator;
    Map<String, Collection<NodeOperation>> operationByServer = NodeOperationGrouper.groupByServer(nodeOperations);
    List<ExecutionPhase> handlerPhases = new ArrayList<>(nodeOperationTrees.size());
    List<RowConsumer> handlerConsumers = new ArrayList<>(nodeOperationTrees.size());
    List<CompletableFuture<Long>> results = new ArrayList<>(nodeOperationTrees.size());
    for (NodeOperationTree nodeOperationTree : nodeOperationTrees) {
        CollectingRowConsumer<?, Long> consumer = new CollectingRowConsumer<>(Collectors.collectingAndThen(Collectors.summingLong(r -> ((long) r.get(0))), sum -> sum));
        handlerConsumers.add(consumer);
        results.add(consumer.completionFuture());
        handlerPhases.add(nodeOperationTree.leaf());
    }
    try {
        setupTasks(txnCtx, operationByServer, handlerPhases, handlerConsumers);
    } catch (Throwable throwable) {
        return Collections.singletonList(CompletableFuture.failedFuture(throwable));
    }
    return results;
}

Example 3 with CollectingRowConsumer

use of io.crate.data.CollectingRowConsumer in project crate by crate.

the class MultiPhaseExecutor method execute.

public static CompletableFuture<SubQueryResults> execute(Map<LogicalPlan, SelectSymbol> dependencies, DependencyCarrier executor, PlannerContext plannerContext, Row params) {
    List<CompletableFuture<?>> dependencyFutures = new ArrayList<>(dependencies.size());
    IdentityHashMap<SelectSymbol, Object> valueBySubQuery = new IdentityHashMap<>();
    for (Map.Entry<LogicalPlan, SelectSymbol> entry : dependencies.entrySet()) {
        LogicalPlan depPlan = entry.getKey();
        SelectSymbol selectSymbol = entry.getValue();
        CollectingRowConsumer<?, ?> rowConsumer = getConsumer(selectSymbol.getResultType());
        depPlan.execute(executor, PlannerContext.forSubPlan(plannerContext), rowConsumer, params, SubQueryResults.EMPTY);
        dependencyFutures.add(rowConsumer.completionFuture().thenAccept(val -> {
            synchronized (valueBySubQuery) {
                valueBySubQuery.put(selectSymbol, val);
            }
        }));
    }
    return CompletableFuture.allOf(dependencyFutures.toArray(new CompletableFuture[0])).thenApply(ignored -> new SubQueryResults(valueBySubQuery));
}