use of io.crate.execution.dsl.projection.EvalProjection in project crate by crate.
the class Get method build.
@Override
public ExecutionPlan build(PlannerContext plannerContext, Set<PlanHint> hints, ProjectionBuilder projectionBuilder, int limitHint, int offsetHint, @Nullable OrderBy order, @Nullable Integer pageSizeHint, Row params, SubQueryResults subQueryResults) {
HashMap<String, Map<ShardId, List<PKAndVersion>>> idsByShardByNode = new HashMap<>();
DocTableInfo docTableInfo = tableRelation.tableInfo();
for (DocKeys.DocKey docKey : docKeys) {
String id = docKey.getId(plannerContext.transactionContext(), plannerContext.nodeContext(), params, subQueryResults);
if (id == null) {
continue;
}
List<String> partitionValues = docKey.getPartitionValues(plannerContext.transactionContext(), plannerContext.nodeContext(), params, subQueryResults);
String indexName = indexName(docTableInfo, partitionValues);
String routing = docKey.getRouting(plannerContext.transactionContext(), plannerContext.nodeContext(), params, subQueryResults);
ShardRouting shardRouting;
try {
shardRouting = plannerContext.resolveShard(indexName, id, routing);
} catch (IndexNotFoundException e) {
if (docTableInfo.isPartitioned()) {
continue;
}
throw e;
}
String currentNodeId = shardRouting.currentNodeId();
if (currentNodeId == null) {
// If relocating is fast enough this will work, otherwise it will result in a shard failure which
// will cause a statement retry
currentNodeId = shardRouting.relocatingNodeId();
if (currentNodeId == null) {
throw new ShardNotFoundException(shardRouting.shardId());
}
}
Map<ShardId, List<PKAndVersion>> idsByShard = idsByShardByNode.get(currentNodeId);
if (idsByShard == null) {
idsByShard = new HashMap<>();
idsByShardByNode.put(currentNodeId, idsByShard);
}
List<PKAndVersion> pkAndVersions = idsByShard.get(shardRouting.shardId());
if (pkAndVersions == null) {
pkAndVersions = new ArrayList<>();
idsByShard.put(shardRouting.shardId(), pkAndVersions);
}
long version = docKey.version(plannerContext.transactionContext(), plannerContext.nodeContext(), params, subQueryResults).orElse(Versions.MATCH_ANY);
long sequenceNumber = docKey.sequenceNo(plannerContext.transactionContext(), plannerContext.nodeContext(), params, subQueryResults).orElse(SequenceNumbers.UNASSIGNED_SEQ_NO);
long primaryTerm = docKey.primaryTerm(plannerContext.transactionContext(), plannerContext.nodeContext(), params, subQueryResults).orElse(SequenceNumbers.UNASSIGNED_PRIMARY_TERM);
pkAndVersions.add(new PKAndVersion(id, version, sequenceNumber, primaryTerm));
}
var docKeyColumns = new ArrayList<>(docTableInfo.primaryKey());
docKeyColumns.addAll(docTableInfo.partitionedBy());
docKeyColumns.add(docTableInfo.clusteredBy());
docKeyColumns.add(DocSysColumns.VERSION);
docKeyColumns.add(DocSysColumns.SEQ_NO);
docKeyColumns.add(DocSysColumns.PRIMARY_TERM);
var binder = new SubQueryAndParamBinder(params, subQueryResults);
List<Symbol> boundOutputs = Lists2.map(outputs, binder);
var boundQuery = binder.apply(query);
// Collect all columns which are used inside the query
// If the query contains only DocKeys, no filter is needed as all DocKeys are handled by the PKLookupOperation
AtomicBoolean requiresAdditionalFilteringOnNonDocKeyColumns = new AtomicBoolean(false);
var toCollectSet = new LinkedHashSet<>(boundOutputs);
Consumer<Reference> addRefIfMatch = ref -> {
toCollectSet.add(ref);
if (docKeyColumns.contains(ref.column()) == false) {
requiresAdditionalFilteringOnNonDocKeyColumns.set(true);
}
};
RefVisitor.visitRefs(boundQuery, addRefIfMatch);
var toCollect = boundOutputs;
ArrayList<Projection> projections = new ArrayList<>();
if (requiresAdditionalFilteringOnNonDocKeyColumns.get()) {
toCollect = List.copyOf(toCollectSet);
var filterProjection = ProjectionBuilder.filterProjection(toCollect, boundQuery);
filterProjection.requiredGranularity(RowGranularity.SHARD);
projections.add(filterProjection);
// reduce outputs which have been added for the filter projection
var evalProjection = new EvalProjection(InputColumn.mapToInputColumns(boundOutputs), RowGranularity.SHARD);
projections.add(evalProjection);
}
var collect = new Collect(new PKLookupPhase(plannerContext.jobId(), plannerContext.nextExecutionPhaseId(), docTableInfo.partitionedBy(), toCollect, idsByShardByNode), TopN.NO_LIMIT, 0, toCollect.size(), docKeys.size(), null);
for (var projection : projections) {
collect.addProjection(projection);
}
return collect;
}
use of io.crate.execution.dsl.projection.EvalProjection in project crate by crate.
the class TopNDistinct method build.
@Override
public ExecutionPlan build(PlannerContext plannerContext, Set<PlanHint> planHints, ProjectionBuilder projectionBuilder, int limitHint, int offsetHint, @Nullable OrderBy order, @Nullable Integer pageSizeHint, Row params, SubQueryResults subQueryResults) {
var executionPlan = source.build(plannerContext, planHints, projectionBuilder, TopN.NO_LIMIT, TopN.NO_OFFSET, null, null, params, subQueryResults);
if (executionPlan.resultDescription().hasRemainingLimitOrOffset()) {
executionPlan = Merge.ensureOnHandler(executionPlan, plannerContext);
}
if (!source.outputs().equals(outputs)) {
EvalProjection evalProjection = new EvalProjection(InputColumns.create(outputs, new InputColumns.SourceSymbols(source.outputs())));
executionPlan.addProjection(evalProjection);
}
int limit = DataTypes.INTEGER.sanitizeValue(evaluate(plannerContext.transactionContext(), plannerContext.nodeContext(), this.limit, params, subQueryResults));
int offset = DataTypes.INTEGER.sanitizeValue(evaluate(plannerContext.transactionContext(), plannerContext.nodeContext(), this.offset, params, subQueryResults));
var inputColOutputs = InputColumn.mapToInputColumns(outputs);
executionPlan.addProjection(new TopNDistinctProjection(limit + offset, inputColOutputs, source.preferShardProjections() ? RowGranularity.SHARD : RowGranularity.CLUSTER));
boolean onHandler = ExecutionPhases.executesOnHandler(plannerContext.handlerNode(), executionPlan.resultDescription().nodeIds());
if (!onHandler || source.preferShardProjections()) {
if (!onHandler) {
executionPlan = Merge.ensureOnHandler(executionPlan, plannerContext);
}
TopNDistinctProjection topNDistinct = new TopNDistinctProjection(limit + offset, inputColOutputs, RowGranularity.CLUSTER);
executionPlan.addProjection(topNDistinct);
}
if (offset > 0) {
// TopNDistinctProjection outputs a distinct result-set,
// That allows us to use the TopNProjection to apply the offset
executionPlan.addProjection(new TopNProjection(limit, offset, Symbols.typeView(inputColOutputs)));
}
return executionPlan;
}
use of io.crate.execution.dsl.projection.EvalProjection in project crate by crate.
the class InsertFromSubQueryPlanner method plan.
public static LogicalPlan plan(AnalyzedInsertStatement statement, PlannerContext plannerContext, LogicalPlanner logicalPlanner, SubqueryPlanner subqueryPlanner) {
if (statement.outputs() != null && !plannerContext.clusterState().getNodes().getMinNodeVersion().onOrAfter(Version.V_4_2_0)) {
throw new UnsupportedFeatureException(RETURNING_VERSION_ERROR_MSG);
}
List<Reference> targetColsExclPartitionCols = new ArrayList<>(statement.columns().size() - statement.tableInfo().partitionedBy().size());
for (Reference column : statement.columns()) {
if (statement.tableInfo().partitionedBy().contains(column.column())) {
continue;
}
targetColsExclPartitionCols.add(column);
}
List<Symbol> columnSymbols = InputColumns.create(targetColsExclPartitionCols, new InputColumns.SourceSymbols(statement.columns()));
// if fields are null default to number of rows imported
var outputs = statement.outputs() == null ? List.of(new InputColumn(0, DataTypes.LONG)) : statement.outputs();
ColumnIndexWriterProjection indexWriterProjection = new ColumnIndexWriterProjection(statement.tableInfo().ident(), null, statement.tableInfo().primaryKey(), statement.columns(), targetColsExclPartitionCols, columnSymbols, statement.isIgnoreDuplicateKeys(), statement.onDuplicateKeyAssignments(), statement.primaryKeySymbols(), statement.partitionedBySymbols(), statement.tableInfo().clusteredBy(), statement.clusteredBySymbol(), Settings.EMPTY, statement.tableInfo().isPartitioned(), outputs, statement.outputs() == null ? List.of() : statement.outputs());
LogicalPlan plannedSubQuery = logicalPlanner.plan(statement.subQueryRelation(), plannerContext, subqueryPlanner, true);
EvalProjection castOutputs = EvalProjection.castValues(Symbols.typeView(statement.columns()), plannedSubQuery.outputs());
return new Insert(plannedSubQuery, indexWriterProjection, castOutputs);
}
use of io.crate.execution.dsl.projection.EvalProjection in project crate by crate.
the class GroupByPlannerTest method testGroupByWithAggregationPlan.
@Test
public void testGroupByWithAggregationPlan() throws Exception {
var e = SQLExecutor.builder(clusterService, 2, RandomizedTest.getRandom(), List.of()).addTable(TableDefinitions.USER_TABLE_DEFINITION).build();
Merge distributedGroupByMerge = e.plan("select count(*), name from users group by name");
Merge reducerMerge = (Merge) distributedGroupByMerge.subPlan();
// distributed collect
RoutedCollectPhase collectPhase = ((RoutedCollectPhase) ((Collect) reducerMerge.subPlan()).collectPhase());
assertThat(collectPhase.maxRowGranularity(), is(RowGranularity.DOC));
assertThat(collectPhase.nodeIds().size(), is(2));
assertThat(collectPhase.toCollect().size(), is(1));
assertThat(collectPhase.projections().size(), is(1));
assertThat(collectPhase.projections().get(0), instanceOf(GroupProjection.class));
assertThat(collectPhase.outputTypes().size(), is(2));
assertEquals(DataTypes.STRING, collectPhase.outputTypes().get(0));
assertEquals(CountAggregation.LongStateType.INSTANCE, collectPhase.outputTypes().get(1));
MergePhase mergePhase = reducerMerge.mergePhase();
assertThat(mergePhase.numUpstreams(), is(2));
assertThat(mergePhase.nodeIds().size(), is(2));
assertEquals(mergePhase.inputTypes(), collectPhase.outputTypes());
// for function evaluation and column-reordering there is always a EvalProjection
assertThat(mergePhase.projections().size(), is(2));
assertThat(mergePhase.projections().get(1), instanceOf(EvalProjection.class));
assertThat(mergePhase.projections().get(0), instanceOf(GroupProjection.class));
GroupProjection groupProjection = (GroupProjection) mergePhase.projections().get(0);
InputColumn inputColumn = (InputColumn) groupProjection.values().get(0).inputs().get(0);
assertThat(inputColumn.index(), is(1));
assertThat(mergePhase.outputTypes().size(), is(2));
assertEquals(DataTypes.LONG, mergePhase.outputTypes().get(0));
assertEquals(DataTypes.STRING, mergePhase.outputTypes().get(1));
MergePhase localMerge = distributedGroupByMerge.mergePhase();
assertThat(localMerge.numUpstreams(), is(2));
assertThat(localMerge.nodeIds().size(), is(1));
assertThat(localMerge.nodeIds().iterator().next(), is(NODE_ID));
assertEquals(mergePhase.outputTypes(), localMerge.inputTypes());
assertThat(localMerge.projections(), empty());
}
use of io.crate.execution.dsl.projection.EvalProjection in project crate by crate.
the class InsertPlannerTest method testInsertFromSubQueryReduceOnCollectorGroupByWithCast.
@Test
public void testInsertFromSubQueryReduceOnCollectorGroupByWithCast() {
Merge merge = e.plan("insert into users (id, name) (select id, count(*) from users group by id)");
Collect nonDistributedGroupBy = (Collect) merge.subPlan();
RoutedCollectPhase collectPhase = ((RoutedCollectPhase) nonDistributedGroupBy.collectPhase());
assertThat(collectPhase.projections(), contains(instanceOf(GroupProjection.class), instanceOf(EvalProjection.class), instanceOf(ColumnIndexWriterProjection.class)));
EvalProjection collectTopN = (EvalProjection) collectPhase.projections().get(1);
assertThat(collectTopN.outputs(), contains(isInputColumn(0), isFunction(ImplicitCastFunction.NAME, List.of(DataTypes.LONG, DataTypes.STRING))));
ColumnIndexWriterProjection columnIndexWriterProjection = (ColumnIndexWriterProjection) collectPhase.projections().get(2);
assertThat(columnIndexWriterProjection.columnReferencesExclPartition(), contains(isReference("id"), isReference("name")));
MergePhase mergePhase = merge.mergePhase();
assertThat(mergePhase.projections(), contains(instanceOf(MergeCountProjection.class)));
}
Aggregations