Examples with CompositeCollector io.crate.operation.collect.collectors.CompositeCollector used on opensource projects

Example 1 with CompositeCollector

use of io.crate.operation.collect.collectors.CompositeCollector in project crate by crate.

the class ShardCollectSource method getCollector.

@Override
public CrateCollector getCollector(CollectPhase phase, BatchConsumer lastConsumer, JobCollectContext jobCollectContext) {
    RoutedCollectPhase collectPhase = (RoutedCollectPhase) phase;
    RoutedCollectPhase normalizedPhase = collectPhase.normalize(nodeNormalizer, null);
    String localNodeId = clusterService.localNode().getId();
    BatchConsumer firstConsumer = ProjectingBatchConsumer.create(lastConsumer, Projections.nodeProjections(normalizedPhase.projections()), collectPhase.jobId(), jobCollectContext.queryPhaseRamAccountingContext(), sharedProjectorFactory);
    if (normalizedPhase.maxRowGranularity() == RowGranularity.SHARD) {
        // The getShardsCollector method always only uses a single RowReceiver and not one per shard)
        return getShardsCollector(collectPhase, normalizedPhase, localNodeId, firstConsumer);
    }
    OrderBy orderBy = normalizedPhase.orderBy();
    if (normalizedPhase.maxRowGranularity() == RowGranularity.DOC && orderBy != null) {
        return createMultiShardScoreDocCollector(normalizedPhase, firstConsumer, jobCollectContext, localNodeId);
    }
    // actual shards might be less if table is partitioned and a partition has been deleted meanwhile
    final int maxNumShards = normalizedPhase.routing().numShards(localNodeId);
    boolean hasShardProjections = Projections.hasAnyShardProjections(normalizedPhase.projections());
    Map<String, Map<String, List<Integer>>> locations = normalizedPhase.routing().locations();
    final List<CrateCollector.Builder> builders = new ArrayList<>(maxNumShards);
    Map<String, List<Integer>> indexShards = locations.get(localNodeId);
    if (indexShards != null) {
        builders.addAll(getDocCollectors(jobCollectContext, normalizedPhase, lastConsumer.requiresScroll(), indexShards));
    }
    switch(builders.size()) {
        case 0:
            return RowsCollector.empty(firstConsumer);
        case 1:
            CrateCollector.Builder collectorBuilder = builders.iterator().next();
            return collectorBuilder.build(collectorBuilder.applyProjections(firstConsumer));
        default:
            if (hasShardProjections) {
                // in order to process shard-based projections concurrently
                return new CompositeCollector(builders, firstConsumer, iterators -> new AsyncCompositeBatchIterator(executor, iterators));
            } else {
                return new CompositeCollector(builders, firstConsumer, CompositeBatchIterator::new);
            }
    }
}