Examples with ShardCollectorProvider io.crate.execution.engine.collect.ShardCollectorProvider used on opensource projects

Example 1 with ShardCollectorProvider

use of io.crate.execution.engine.collect.ShardCollectorProvider in project crate by crate.

the class ShardCollectSource method createMultiShardScoreDocCollector.

private CompletableFuture<BatchIterator<Row>> createMultiShardScoreDocCollector(RoutedCollectPhase collectPhase, boolean supportMoveToStart, CollectTask collectTask, String localNodeId) {
    Map<String, Map<String, IntIndexedContainer>> locations = collectPhase.routing().locations();
    SharedShardContexts sharedShardContexts = collectTask.sharedShardContexts();
    Map<String, IntIndexedContainer> indexShards = locations.get(localNodeId);
    List<CompletableFuture<OrderedDocCollector>> orderedDocCollectors = new ArrayList<>();
    Metadata metadata = clusterService.state().metadata();
    for (Map.Entry<String, IntIndexedContainer> entry : indexShards.entrySet()) {
        String indexName = entry.getKey();
        Index index = metadata.index(indexName).getIndex();
        for (IntCursor shard : entry.getValue()) {
            ShardId shardId = new ShardId(index, shard.value);
            try {
                SharedShardContext context = sharedShardContexts.getOrCreateContext(shardId);
                ShardCollectorProvider shardCollectorProvider = getCollectorProviderSafe(shardId);
                orderedDocCollectors.add(shardCollectorProvider.getFutureOrderedCollector(collectPhase, context, collectTask, supportMoveToStart));
            } catch (ShardNotFoundException | IllegalIndexShardStateException e) {
                throw e;
            } catch (IndexNotFoundException e) {
                if (IndexParts.isPartitioned(indexName)) {
                    break;
                }
                throw e;
            }
        }
    }
    List<DataType<?>> columnTypes = Symbols.typeView(collectPhase.toCollect());
    OrderBy orderBy = collectPhase.orderBy();
    assert orderBy != null : "orderBy must not be null";
    return CompletableFutures.allAsList(orderedDocCollectors).thenApply(collectors -> OrderedLuceneBatchIteratorFactory.newInstance(collectors, OrderingByPosition.rowOrdering(OrderByPositionVisitor.orderByPositions(orderBy.orderBySymbols(), collectPhase.toCollect()), orderBy.reverseFlags(), orderBy.nullsFirst()), new RowAccountingWithEstimators(columnTypes, collectTask.getRamAccounting()), executor, availableThreads, supportMoveToStart));
}

Example 2 with ShardCollectorProvider

use of io.crate.execution.engine.collect.ShardCollectorProvider in project crate by crate.

the class ShardCollectSource method getShardsIterator.

private Iterable<Row> getShardsIterator(TransactionContext txnCtx, RoutedCollectPhase collectPhase, String localNodeId) {
    Map<String, Map<String, IntIndexedContainer>> locations = collectPhase.routing().locations();
    List<UnassignedShard> unassignedShards = new ArrayList<>();
    List<ShardRowContext> shardRowContexts = new ArrayList<>();
    Map<String, IntIndexedContainer> indexShardsMap = locations.get(localNodeId);
    Metadata metadata = clusterService.state().metadata();
    for (Map.Entry<String, IntIndexedContainer> indexShards : indexShardsMap.entrySet()) {
        String indexName = indexShards.getKey();
        IndexMetadata indexMetadata = metadata.index(indexName);
        if (indexMetadata == null) {
            continue;
        }
        Index index = indexMetadata.getIndex();
        IntIndexedContainer shards = indexShards.getValue();
        IndexService indexService = indicesService.indexService(index);
        if (indexService == null) {
            for (IntCursor shard : shards) {
                unassignedShards.add(toUnassignedShard(index.getName(), UnassignedShard.markAssigned(shard.value)));
            }
            continue;
        }
        for (IntCursor shard : shards) {
            if (UnassignedShard.isUnassigned(shard.value)) {
                unassignedShards.add(toUnassignedShard(index.getName(), UnassignedShard.markAssigned(shard.value)));
                continue;
            }
            ShardId shardId = new ShardId(index, shard.value);
            try {
                ShardCollectorProvider shardCollectorProvider = getCollectorProviderSafe(shardId);
                shardRowContexts.add(shardCollectorProvider.shardRowContext());
            } catch (ShardNotFoundException | IllegalIndexShardStateException e) {
                unassignedShards.add(toUnassignedShard(index.getName(), shard.value));
            }
        }
    }
    Iterable<Row> assignedShardRows = RowsTransformer.toRowsIterable(txnCtx, inputFactory, shardReferenceResolver, collectPhase, shardRowContexts, false);
    Iterable<Row> rows;
    if (unassignedShards.size() > 0) {
        Iterable<Row> unassignedShardRows = RowsTransformer.toRowsIterable(txnCtx, inputFactory, unassignedShardReferenceResolver, collectPhase, unassignedShards, false);
        rows = Iterables.concat(assignedShardRows, unassignedShardRows);
    } else {
        rows = assignedShardRows;
    }
    if (collectPhase.orderBy() != null) {
        return RowsTransformer.sortRows(Iterables.transform(rows, Row::materialize), collectPhase);
    }
    return rows;
}

Example 3 with ShardCollectorProvider

use of io.crate.execution.engine.collect.ShardCollectorProvider in project crate by crate.

the class ShardCollectSource method getCollectorProviderSafe.

private ShardCollectorProvider getCollectorProviderSafe(ShardId shardId) {
    Supplier<ShardCollectorProvider> supplier = shards.get(shardId);
    if (supplier == null) {
        throw new ShardNotFoundException(shardId);
    }
    ShardCollectorProvider shardCollectorProvider = supplier.get();
    if (shardCollectorProvider == null) {
        throw new ShardNotFoundException(shardId);
    }
    return shardCollectorProvider;
}

Example 4 with ShardCollectorProvider

use of io.crate.execution.engine.collect.ShardCollectorProvider in project crate by crate.

the class ShardCollectSource method getIterators.

private List<CompletableFuture<BatchIterator<Row>>> getIterators(CollectTask collectTask, RoutedCollectPhase collectPhase, boolean requiresScroll, Map<String, IntIndexedContainer> indexShards) {
    Metadata metadata = clusterService.state().metadata();
    List<CompletableFuture<BatchIterator<Row>>> iterators = new ArrayList<>();
    for (Map.Entry<String, IntIndexedContainer> entry : indexShards.entrySet()) {
        String indexName = entry.getKey();
        IndexMetadata indexMD = metadata.index(indexName);
        if (indexMD == null) {
            if (IndexParts.isPartitioned(indexName)) {
                continue;
            }
            throw new IndexNotFoundException(indexName);
        }
        Index index = indexMD.getIndex();
        try {
            indicesService.indexServiceSafe(index);
        } catch (IndexNotFoundException e) {
            if (IndexParts.isPartitioned(indexName)) {
                continue;
            }
            throw e;
        }
        for (IntCursor shardCursor : entry.getValue()) {
            ShardId shardId = new ShardId(index, shardCursor.value);
            try {
                ShardCollectorProvider shardCollectorProvider = getCollectorProviderSafe(shardId);
                CompletableFuture<BatchIterator<Row>> iterator = shardCollectorProvider.getFutureIterator(collectPhase, requiresScroll, collectTask);
                iterators.add(iterator);
            } catch (ShardNotFoundException | IllegalIndexShardStateException e) {
                // and the reader required in the fetchPhase would be missing.
                if (Symbols.containsColumn(collectPhase.toCollect(), DocSysColumns.FETCHID)) {
                    throw e;
                }
                iterators.add(remoteCollectorFactory.createCollector(shardId, collectPhase, collectTask, shardCollectorProviderFactory));
            } catch (IndexNotFoundException e) {
                // Prevent wrapping this to not break retry-detection
                throw e;
            } catch (Throwable t) {
                Exceptions.rethrowRuntimeException(t);
            }
        }
    }
    return iterators;
}