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

Example 1 with RowConsumer

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

the class InsertFromValues method execute.

@Override
public void execute(DependencyCarrier dependencies, PlannerContext plannerContext, RowConsumer consumer, Row params, SubQueryResults subQueryResults) {
    DocTableInfo tableInfo = dependencies.schemas().getTableInfo(writerProjection.tableIdent(), Operation.INSERT);
    // For instance, the target table of the insert from values
    // statement is the table with the following schema:
    // 
    // CREATE TABLE users (
    // dep_id TEXT,
    // name TEXT,
    // id INT,
    // country_id INT,
    // PRIMARY KEY (dep_id, id, country_id))
    // CLUSTERED BY (dep_id)
    // PARTITIONED BY (country_id)
    // 
    // The insert from values statement below would have the column
    // index writer projection of its plan that contains the column
    // idents and symbols required to create corresponding inputs.
    // The diagram below shows the projection's column symbols used
    // in the plan and relation between symbols sub-/sets.
    // 
    // +------------------------+
    // |          +-------------+  PK symbols
    // cluster by +------+ |          |      +------+
    // symbol            | |          |      |
    // + +          +      +
    // INSERT INTO users (dep_id, name, id, country_id) VALUES (?, ?, ?, ?)
    // +      +    +     +   +
    // +-------+      |    |     |   |
    // all target  +--------------+    |     |   +---+  partitioned by
    // column      +-------------------+     |          symbols
    // symbols     +-------------------------+
    InputFactory inputFactory = new InputFactory(dependencies.nodeContext());
    InputFactory.Context<CollectExpression<Row, ?>> context = inputFactory.ctxForInputColumns(plannerContext.transactionContext());
    var allColumnSymbols = InputColumns.create(writerProjection.allTargetColumns(), new InputColumns.SourceSymbols(writerProjection.allTargetColumns()));
    ArrayList<Input<?>> insertInputs = new ArrayList<>(allColumnSymbols.size());
    for (Symbol symbol : allColumnSymbols) {
        insertInputs.add(context.add(symbol));
    }
    ArrayList<Input<?>> partitionedByInputs = new ArrayList<>(writerProjection.partitionedBySymbols().size());
    for (Symbol partitionedBySymbol : writerProjection.partitionedBySymbols()) {
        partitionedByInputs.add(context.add(partitionedBySymbol));
    }
    ArrayList<Input<?>> primaryKeyInputs = new ArrayList<>(writerProjection.ids().size());
    for (Symbol symbol : writerProjection.ids()) {
        primaryKeyInputs.add(context.add(symbol));
    }
    Input<?> clusterByInput;
    if (writerProjection.clusteredBy() != null) {
        clusterByInput = context.add(writerProjection.clusteredBy());
    } else {
        clusterByInput = null;
    }
    String[] updateColumnNames;
    Symbol[] assignmentSources;
    if (writerProjection.onDuplicateKeyAssignments() == null) {
        updateColumnNames = null;
        assignmentSources = null;
    } else {
        Assignments assignments = Assignments.convert(writerProjection.onDuplicateKeyAssignments(), dependencies.nodeContext());
        assignmentSources = assignments.bindSources(tableInfo, params, subQueryResults);
        updateColumnNames = assignments.targetNames();
    }
    var indexNameResolver = IndexNameResolver.create(writerProjection.tableIdent(), writerProjection.partitionIdent(), partitionedByInputs);
    GroupRowsByShard<ShardUpsertRequest, ShardUpsertRequest.Item> grouper = createRowsByShardGrouper(assignmentSources, insertInputs, indexNameResolver, context, plannerContext, dependencies.clusterService());
    ArrayList<Row> rows = new ArrayList<>();
    evaluateValueTableFunction(tableFunctionRelation.functionImplementation(), tableFunctionRelation.function().arguments(), writerProjection.allTargetColumns(), tableInfo, params, plannerContext, subQueryResults).forEachRemaining(rows::add);
    List<Symbol> returnValues = this.writerProjection.returnValues();
    ShardUpsertRequest.Builder builder = new ShardUpsertRequest.Builder(plannerContext.transactionContext().sessionSettings(), BULK_REQUEST_TIMEOUT_SETTING.get(dependencies.settings()), writerProjection.isIgnoreDuplicateKeys() ? ShardUpsertRequest.DuplicateKeyAction.IGNORE : ShardUpsertRequest.DuplicateKeyAction.UPDATE_OR_FAIL, // continueOnErrors
    rows.size() > 1, updateColumnNames, writerProjection.allTargetColumns().toArray(new Reference[0]), returnValues.isEmpty() ? null : returnValues.toArray(new Symbol[0]), plannerContext.jobId(), false);
    var shardedRequests = new ShardedRequests<>(builder::newRequest, RamAccounting.NO_ACCOUNTING);
    HashMap<String, InsertSourceFromCells> validatorsCache = new HashMap<>();
    for (Row row : rows) {
        grouper.accept(shardedRequests, row);
        try {
            checkPrimaryKeyValuesNotNull(primaryKeyInputs);
            checkClusterByValueNotNull(clusterByInput);
            checkConstraintsOnGeneratedSource(row.materialize(), indexNameResolver.get(), tableInfo, plannerContext, validatorsCache);
        } catch (Throwable t) {
            consumer.accept(null, t);
            return;
        }
    }
    validatorsCache.clear();
    var actionProvider = dependencies.transportActionProvider();
    createIndices(actionProvider.transportBulkCreateIndicesAction(), shardedRequests.itemsByMissingIndex().keySet(), dependencies.clusterService(), plannerContext.jobId()).thenCompose(acknowledgedResponse -> {
        var shardUpsertRequests = resolveAndGroupShardRequests(shardedRequests, dependencies.clusterService()).values();
        return execute(dependencies.nodeLimits(), dependencies.clusterService().state(), shardUpsertRequests, actionProvider.transportShardUpsertAction(), dependencies.scheduler());
    }).whenComplete((response, t) -> {
        if (t == null) {
            if (returnValues.isEmpty()) {
                consumer.accept(InMemoryBatchIterator.of(new Row1((long) response.numSuccessfulWrites()), SENTINEL), null);
            } else {
                consumer.accept(InMemoryBatchIterator.of(new CollectionBucket(response.resultRows()), SENTINEL, false), null);
            }
        } else {
            consumer.accept(null, t);
        }
    });
}

Example 2 with RowConsumer

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

the class RefreshTablePlan method executeOrFail.

@Override
public void executeOrFail(DependencyCarrier dependencies, PlannerContext plannerContext, RowConsumer consumer, Row parameters, SubQueryResults subQueryResults) {
    if (analysis.tables().isEmpty()) {
        consumer.accept(InMemoryBatchIterator.empty(SENTINEL), null);
        return;
    }
    Function<? super Symbol, Object> eval = x -> SymbolEvaluator.evaluate(plannerContext.transactionContext(), plannerContext.nodeContext(), x, parameters, subQueryResults);
    ArrayList<String> toRefresh = new ArrayList<>();
    for (Map.Entry<Table<Symbol>, DocTableInfo> table : analysis.tables().entrySet()) {
        var tableInfo = table.getValue();
        var tableSymbol = table.getKey();
        if (tableSymbol.partitionProperties().isEmpty()) {
            toRefresh.addAll(Arrays.asList(tableInfo.concreteOpenIndices()));
        } else {
            var partitionName = toPartitionName(tableInfo, Lists2.map(tableSymbol.partitionProperties(), p -> p.map(eval)));
            if (!tableInfo.partitions().contains(partitionName)) {
                throw new PartitionUnknownException(partitionName);
            }
            toRefresh.add(partitionName.asIndexName());
        }
    }
    RefreshRequest request = new RefreshRequest(toRefresh.toArray(String[]::new));
    request.indicesOptions(IndicesOptions.lenientExpandOpen());
    var transportRefreshAction = dependencies.transportActionProvider().transportRefreshAction();
    transportRefreshAction.execute(request, new OneRowActionListener<>(consumer, response -> new Row1(toRefresh.isEmpty() ? -1L : (long) toRefresh.size())));
}

Example 3 with RowConsumer

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

the class SetSessionPlan method executeOrFail.

@Override
public void executeOrFail(DependencyCarrier executor, PlannerContext plannerContext, RowConsumer consumer, Row params, SubQueryResults subQueryResults) throws Exception {
    Function<? super Symbol, Object> eval = x -> SymbolEvaluator.evaluate(plannerContext.transactionContext(), plannerContext.nodeContext(), x, params, subQueryResults);
    SessionContext sessionContext = plannerContext.transactionContext().sessionContext();
    Assignment<Symbol> assignment = settings.get(0);
    String settingName = eval.apply(assignment.columnName()).toString();
    validateSetting(settingName);
    SessionSetting<?> sessionSetting = sessionSettingRegistry.settings().get(settingName);
    if (sessionSetting == null) {
        LOGGER.info("SET SESSION STATEMENT WILL BE IGNORED: {}", settingName);
    } else {
        sessionSetting.apply(sessionContext, assignment.expressions(), eval);
    }
    consumer.accept(InMemoryBatchIterator.empty(SENTINEL), null);
}

Example 4 with RowConsumer

use of io.crate.data.RowConsumer 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 5 with RowConsumer

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

the class JobLauncher method createHandlerPhaseAndReceivers.

private List<Tuple<ExecutionPhase, RowConsumer>> createHandlerPhaseAndReceivers(List<ExecutionPhase> handlerPhases, List<RowConsumer> handlerReceivers, InitializationTracker initializationTracker) {
    List<Tuple<ExecutionPhase, RowConsumer>> handlerPhaseAndReceiver = new ArrayList<>();
    ListIterator<RowConsumer> consumerIt = handlerReceivers.listIterator();
    for (ExecutionPhase handlerPhase : handlerPhases) {
        InterceptingRowConsumer interceptingBatchConsumer = new InterceptingRowConsumer(jobId, consumerIt.next(), initializationTracker, executor, transportKillJobsNodeAction);
        handlerPhaseAndReceiver.add(new Tuple<>(handlerPhase, interceptingBatchConsumer));
    }
    return handlerPhaseAndReceiver;
}