use of io.crate.data.SentinelRow.SENTINEL 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);
}
});
}
use of io.crate.data.SentinelRow.SENTINEL 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())));
}
use of io.crate.data.SentinelRow.SENTINEL 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);
}
use of io.crate.data.SentinelRow.SENTINEL in project crate by crate.
the class CreateTablePlan method executeOrFail.
@Override
public void executeOrFail(DependencyCarrier dependencies, PlannerContext plannerContext, RowConsumer consumer, Row params, SubQueryResults subQueryResults) {
BoundCreateTable boundCreateTable = bind(createTable, plannerContext.transactionContext(), dependencies.nodeContext(), params, subQueryResults, numberOfShards, schemas, dependencies.fulltextAnalyzerResolver());
if (boundCreateTable.noOp()) {
consumer.accept(InMemoryBatchIterator.empty(SENTINEL), null);
return;
}
tableCreator.create(boundCreateTable).whenComplete(new OneRowActionListener<>(consumer, rCount -> new Row1(rCount == null ? -1 : rCount)));
}
use of io.crate.data.SentinelRow.SENTINEL in project crate by crate.
the class AbstractWindowFunctionTest method assertEvaluate.
@SuppressWarnings("unchecked")
protected <T> void assertEvaluate(String functionExpression, Matcher<T> expectedValue, List<ColumnIdent> rowsColumnDescription, Object[]... inputRows) throws Throwable {
performInputSanityChecks(inputRows);
Symbol normalizedFunctionSymbol = sqlExpressions.normalize(sqlExpressions.asSymbol(functionExpression));
assertThat(normalizedFunctionSymbol, instanceOf(io.crate.expression.symbol.WindowFunction.class));
var windowFunctionSymbol = (io.crate.expression.symbol.WindowFunction) normalizedFunctionSymbol;
ReferenceResolver<InputCollectExpression> referenceResolver = r -> new InputCollectExpression(rowsColumnDescription.indexOf(r.column()));
var sourceSymbols = Lists2.map(rowsColumnDescription, x -> sqlExpressions.normalize(sqlExpressions.asSymbol(x.sqlFqn())));
ensureInputRowsHaveCorrectType(sourceSymbols, inputRows);
var argsCtx = inputFactory.ctxForRefs(txnCtx, referenceResolver);
argsCtx.add(windowFunctionSymbol.arguments());
FunctionImplementation impl = sqlExpressions.nodeCtx.functions().getQualified(windowFunctionSymbol, txnCtx.sessionSettings().searchPath());
assert impl instanceof WindowFunction || impl instanceof AggregationFunction : "Got " + impl + " but expected a window function";
WindowFunction windowFunctionImpl;
if (impl instanceof AggregationFunction) {
windowFunctionImpl = new AggregateToWindowFunctionAdapter((AggregationFunction) impl, new ExpressionsInput<>(Literal.BOOLEAN_TRUE, List.of()), Version.CURRENT, RamAccounting.NO_ACCOUNTING, memoryManager, Version.CURRENT);
} else {
windowFunctionImpl = (WindowFunction) impl;
}
int numCellsInSourceRows = inputRows[0].length;
var windowDef = windowFunctionSymbol.windowDefinition();
var partitionOrderBy = windowDef.partitions().isEmpty() ? null : new OrderBy(windowDef.partitions());
Comparator<Object[]> cmpOrderBy = createComparator(() -> inputFactory.ctxForRefs(txnCtx, referenceResolver), windowDef.orderBy());
InputColumns.SourceSymbols inputColSources = new InputColumns.SourceSymbols(sourceSymbols);
var mappedWindowDef = windowDef.map(s -> InputColumns.create(s, inputColSources));
BatchIterator<Row> iterator = WindowFunctionBatchIterator.of(InMemoryBatchIterator.of(Arrays.stream(inputRows).map(RowN::new).collect(Collectors.toList()), SENTINEL, true), new IgnoreRowAccounting(), WindowProjector.createComputeStartFrameBoundary(numCellsInSourceRows, txnCtx, sqlExpressions.nodeCtx, mappedWindowDef, cmpOrderBy), WindowProjector.createComputeEndFrameBoundary(numCellsInSourceRows, txnCtx, sqlExpressions.nodeCtx, mappedWindowDef, cmpOrderBy), createComparator(() -> inputFactory.ctxForRefs(txnCtx, referenceResolver), partitionOrderBy), cmpOrderBy, numCellsInSourceRows, () -> 1, Runnable::run, List.of(windowFunctionImpl), argsCtx.expressions(), new Boolean[] { windowFunctionSymbol.ignoreNulls() }, argsCtx.topLevelInputs().toArray(new Input[0]));
List<Object> actualResult;
try {
actualResult = BatchIterators.collect(iterator, Collectors.mapping(row -> row.get(numCellsInSourceRows), Collectors.toList())).get(5, TimeUnit.SECONDS);
} catch (ExecutionException e) {
throw e.getCause();
}
assertThat((T) actualResult, expectedValue);
}
Aggregations