Examples with CollectPhase io.crate.execution.dsl.phases.CollectPhase used on opensource projects

Example 6 with CollectPhase

use of io.crate.execution.dsl.phases.CollectPhase in project crate by crate.

the class GroupByPlannerTest method testCountDistinctWithGroupBy.

@Test
public void testCountDistinctWithGroupBy() throws Exception {
    var e = SQLExecutor.builder(clusterService, 2, RandomizedTest.getRandom(), List.of()).addTable(TableDefinitions.USER_TABLE_DEFINITION).build();
    Merge distributedGroupByMerge = e.plan("select count(distinct id), name from users group by name order by count(distinct id)");
    Merge reducerMerge = (Merge) distributedGroupByMerge.subPlan();
    CollectPhase collectPhase = ((Collect) reducerMerge.subPlan()).collectPhase();
    // collect
    assertThat(collectPhase.toCollect().get(0), instanceOf(Reference.class));
    assertThat(collectPhase.toCollect().size(), is(2));
    assertThat(((Reference) collectPhase.toCollect().get(0)).column().name(), is("id"));
    assertThat(((Reference) collectPhase.toCollect().get(1)).column().name(), is("name"));
    Projection projection = collectPhase.projections().get(0);
    assertThat(projection, instanceOf(GroupProjection.class));
    GroupProjection groupProjection = (GroupProjection) projection;
    Symbol groupKey = groupProjection.keys().get(0);
    assertThat(groupKey, instanceOf(InputColumn.class));
    assertThat(((InputColumn) groupKey).index(), is(1));
    assertThat(groupProjection.values().size(), is(1));
    assertThat(groupProjection.mode(), is(AggregateMode.ITER_PARTIAL));
    Aggregation aggregation = groupProjection.values().get(0);
    Symbol aggregationInput = aggregation.inputs().get(0);
    assertThat(aggregationInput.symbolType(), is(SymbolType.INPUT_COLUMN));
    // reducer
    MergePhase mergePhase = reducerMerge.mergePhase();
    assertThat(mergePhase.projections(), contains(instanceOf(GroupProjection.class), instanceOf(OrderedTopNProjection.class), instanceOf(EvalProjection.class)));
    Projection groupProjection1 = mergePhase.projections().get(0);
    groupProjection = (GroupProjection) groupProjection1;
    assertThat(groupProjection.keys().get(0), instanceOf(InputColumn.class));
    assertThat(((InputColumn) groupProjection.keys().get(0)).index(), is(0));
    assertThat(groupProjection.mode(), is(AggregateMode.PARTIAL_FINAL));
    assertThat(groupProjection.values().get(0), instanceOf(Aggregation.class));
    OrderedTopNProjection topNProjection = (OrderedTopNProjection) mergePhase.projections().get(1);
    Symbol collection_count = topNProjection.outputs().get(0);
    assertThat(collection_count, SymbolMatchers.isInputColumn(0));
    // handler
    MergePhase localMergeNode = distributedGroupByMerge.mergePhase();
    assertThat(localMergeNode.projections(), Matchers.emptyIterable());
}

Example 7 with CollectPhase

use of io.crate.execution.dsl.phases.CollectPhase in project crate by crate.

the class ShardCollectSource method getIterator.

@Override
public CompletableFuture<BatchIterator<Row>> getIterator(TransactionContext txnCtx, CollectPhase phase, CollectTask collectTask, boolean supportMoveToStart) {
    RoutedCollectPhase collectPhase = (RoutedCollectPhase) phase;
    String localNodeId = clusterService.localNode().getId();
    Projectors projectors = new Projectors(collectPhase.projections(), collectPhase.jobId(), collectTask.txnCtx(), collectTask.getRamAccounting(), collectTask.memoryManager(), sharedProjectorFactory);
    boolean requireMoveToStartSupport = supportMoveToStart && !projectors.providesIndependentScroll();
    if (collectPhase.maxRowGranularity() == RowGranularity.SHARD) {
        return CompletableFuture.completedFuture(projectors.wrap(InMemoryBatchIterator.of(getShardsIterator(collectTask.txnCtx(), collectPhase, localNodeId), SentinelRow.SENTINEL, true)));
    }
    OrderBy orderBy = collectPhase.orderBy();
    if (collectPhase.maxRowGranularity() == RowGranularity.DOC && orderBy != null) {
        return createMultiShardScoreDocCollector(collectPhase, requireMoveToStartSupport, collectTask, localNodeId).thenApply(projectors::wrap);
    }
    boolean hasShardProjections = Projections.hasAnyShardProjections(collectPhase.projections());
    Map<String, IntIndexedContainer> indexShards = collectPhase.routing().locations().get(localNodeId);
    List<CompletableFuture<BatchIterator<Row>>> iterators = indexShards == null ? Collections.emptyList() : getIterators(collectTask, collectPhase, requireMoveToStartSupport, indexShards);
    final CompletableFuture<BatchIterator<Row>> result;
    switch(iterators.size()) {
        case 0:
            result = CompletableFuture.completedFuture(InMemoryBatchIterator.empty(SentinelRow.SENTINEL));
            break;
        case 1:
            result = iterators.get(0);
            break;
        default:
            if (hasShardProjections) {
                // use AsyncCompositeBatchIterator for multi-threaded loadNextBatch
                // in order to process shard-based projections concurrently
                // noinspection unchecked
                result = CompletableFutures.allAsList(iterators).thenApply(its -> CompositeBatchIterator.asyncComposite(executor, availableThreads, its.toArray(new BatchIterator[0])));
            } else {
                // noinspection unchecked
                result = CompletableFutures.allAsList(iterators).thenApply(its -> CompositeBatchIterator.seqComposite(its.toArray(new BatchIterator[0])));
            }
    }
    return result.thenApply(it -> projectors.wrap(it));
}

Example 8 with CollectPhase

use of io.crate.execution.dsl.phases.CollectPhase in project crate by crate.

the class SystemCollectSource method getIterator.

@Override
public CompletableFuture<BatchIterator<Row>> getIterator(TransactionContext txnCtx, CollectPhase phase, CollectTask collectTask, boolean supportMoveToStart) {
    RoutedCollectPhase collectPhase = (RoutedCollectPhase) phase;
    Map<String, Map<String, IntIndexedContainer>> locations = collectPhase.routing().locations();
    String table = Iterables.getOnlyElement(locations.get(clusterService.localNode().getId()).keySet());
    RelationName relationName = RelationName.fromIndexName(table);
    StaticTableDefinition<?> tableDefinition = tableDefinition(relationName);
    User user = requireNonNull(userLookup.findUser(txnCtx.sessionSettings().userName()), "User who invoked a statement must exist");
    return CompletableFuture.completedFuture(CollectingBatchIterator.newInstance(() -> {
    }, // If data is already local, then `CollectingBatchIterator` takes care of kill handling.
    t -> {
    }, () -> tableDefinition.retrieveRecords(txnCtx, user).thenApply(records -> recordsToRows(collectPhase, collectTask.txnCtx(), tableDefinition.getReferenceResolver(), supportMoveToStart, records)), tableDefinition.involvesIO()));
}

Example 9 with CollectPhase

use of io.crate.execution.dsl.phases.CollectPhase in project crate by crate.

the class SubQueryPlannerTest method testNestedSimpleSelectContainsGroupProjectionWithFunction.

@Test
public void testNestedSimpleSelectContainsGroupProjectionWithFunction() throws Exception {
    Collect collect = e.plan("select c + 100, max(max) from " + "    (select x + 10::int as c, max(i) as max from t1 group by x + 10::int) t " + "group by c + 100 order by c + 100 " + "limit 100");
    CollectPhase collectPhase = collect.collectPhase();
    assertThat(collectPhase.toCollect(), contains(isReference("i"), isFunction("add", isReference("x"), isLiteral(10))));
    assertThat(collectPhase.projections(), contains(instanceOf(GroupProjection.class), instanceOf(GroupProjection.class), instanceOf(EvalProjection.class), instanceOf(GroupProjection.class), instanceOf(OrderedTopNProjection.class), instanceOf(TopNProjection.class)));
}

Example 10 with CollectPhase

use of io.crate.execution.dsl.phases.CollectPhase in project crate by crate.

the class SubQueryPlannerTest method testJoinWithAggregationOnSubSelectsWithAggregations.

@Test
public void testJoinWithAggregationOnSubSelectsWithAggregations() throws Exception {
    Join nl = e.plan("select t1.a, count(*) from " + " (select a, count(*) as cnt from t1 group by a) t1 " + "join" + " (select distinct i from t2) t2 " + "on t1.cnt = t2.i::long " + "group by t1.a");
    assertThat(nl.joinPhase().projections(), contains(instanceOf(EvalProjection.class), instanceOf(GroupProjection.class)));
    assertThat(nl.left(), instanceOf(Collect.class));
    Collect leftPlan = (Collect) nl.left();
    CollectPhase leftCollectPhase = leftPlan.collectPhase();
    assertThat(leftCollectPhase.projections(), contains(instanceOf(GroupProjection.class), instanceOf(GroupProjection.class), instanceOf(EvalProjection.class)));
    Collect rightPlan = (Collect) nl.right();
    assertThat(rightPlan.collectPhase().projections(), contains(instanceOf(GroupProjection.class), instanceOf(GroupProjection.class)));
}