Examples with Collect io.crate.planner.node.dql.Collect used on opensource projects

Example 1 with Collect

use of io.crate.planner.node.dql.Collect in project crate by crate.

the class DeleteStatementPlanner method collectWithDeleteProjection.

private static Plan collectWithDeleteProjection(TableInfo tableInfo, WhereClause whereClause, Planner.Context plannerContext) {
    // for delete, we always need to collect the `_uid`
    Reference idReference = tableInfo.getReference(DocSysColumns.ID);
    DeleteProjection deleteProjection = new DeleteProjection(new InputColumn(0, DataTypes.STRING));
    Routing routing = plannerContext.allocateRouting(tableInfo, whereClause, Preference.PRIMARY.type());
    RoutedCollectPhase collectPhase = new RoutedCollectPhase(plannerContext.jobId(), plannerContext.nextExecutionPhaseId(), "collect", routing, tableInfo.rowGranularity(), ImmutableList.of(idReference), ImmutableList.of(deleteProjection), whereClause, DistributionInfo.DEFAULT_BROADCAST);
    Collect collect = new Collect(collectPhase, TopN.NO_LIMIT, 0, 1, 1, null);
    return Merge.ensureOnHandler(collect, plannerContext, Collections.singletonList(MergeCountProjection.INSTANCE));
}

Example 2 with Collect

use of io.crate.planner.node.dql.Collect 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;
}

Example 3 with Collect

use of io.crate.planner.node.dql.Collect in project crate by crate.

the class NodeOperationTreeGenerator method visitMerge.

@Override
public Void visitMerge(Merge merge, NodeOperationTreeContext context) {
    ExecutionPlan subExecutionPlan = merge.subPlan();
    boolean useDirectResponse = context.noPreviousPhases() && subExecutionPlan instanceof Collect && !Paging.shouldPage(subExecutionPlan.resultDescription().maxRowsPerNode());
    context.addPhase(merge.mergePhase());
    if (useDirectResponse) {
        context.addPhase(((Collect) subExecutionPlan).collectPhase(), true);
    } else {
        process(subExecutionPlan, context);
    }
    return null;
}

Example 4 with Collect

use of io.crate.planner.node.dql.Collect in project crate by crate.

the class GroupByPlannerTest method testNonDistributedGroupByOnClusteredColumn.

@Test
public void testNonDistributedGroupByOnClusteredColumn() throws Exception {
    var e = SQLExecutor.builder(clusterService, 2, RandomizedTest.getRandom(), List.of()).addTable(TableDefinitions.USER_TABLE_DEFINITION).build();
    Merge merge = e.plan("select count(*), id from users group by id limit 20");
    Collect collect = ((Collect) merge.subPlan());
    RoutedCollectPhase collectPhase = ((RoutedCollectPhase) collect.collectPhase());
    assertThat(collectPhase.projections(), contains(instanceOf(GroupProjection.class), instanceOf(TopNProjection.class), // swaps id, count(*) output from group by to count(*), id
    instanceOf(EvalProjection.class)));
    assertThat(collectPhase.projections().get(0).requiredGranularity(), is(RowGranularity.SHARD));
    MergePhase mergePhase = merge.mergePhase();
    assertThat(mergePhase.projections(), contains(instanceOf(TopNProjection.class)));
}

Example 5 with Collect

use of io.crate.planner.node.dql.Collect in project crate by crate.

the class GroupByPlannerTest method testDistributedGroupByProjectionHasShardLevelGranularity.

@Test
public void testDistributedGroupByProjectionHasShardLevelGranularity() throws Exception {
    var e = SQLExecutor.builder(clusterService, 2, RandomizedTest.getRandom(), List.of()).addTable(TableDefinitions.USER_TABLE_DEFINITION).build();
    Merge distributedGroupByMerge = e.plan("select count(*) from users group by name");
    Merge reduceMerge = (Merge) distributedGroupByMerge.subPlan();
    CollectPhase collectPhase = ((Collect) reduceMerge.subPlan()).collectPhase();
    assertThat(collectPhase.projections().size(), is(1));
    assertThat(collectPhase.projections().get(0), instanceOf(GroupProjection.class));
    assertThat(collectPhase.projections().get(0).requiredGranularity(), is(RowGranularity.SHARD));
}