Examples with BatchIterator io.crate.data.BatchIterator used on opensource projects

Example 6 with BatchIterator

use of io.crate.data.BatchIterator in project crate by crate.

the class RemoteCollectorFactory method retrieveRows.

private CompletableFuture<List<Row>> retrieveRows(ShardRouting activePrimaryRouting, RoutedCollectPhase collectPhase, CollectTask collectTask, ShardCollectorProviderFactory shardCollectorProviderFactory) {
    Collector<Row, ?, List<Object[]>> listCollector = Collectors.mapping(Row::materialize, Collectors.toList());
    CollectingRowConsumer<?, List<Object[]>> consumer = new CollectingRowConsumer<>(listCollector);
    String nodeId = activePrimaryRouting.currentNodeId();
    String localNodeId = clusterService.localNode().getId();
    if (localNodeId.equalsIgnoreCase(nodeId)) {
        var indexShard = indicesService.indexServiceSafe(activePrimaryRouting.index()).getShard(activePrimaryRouting.shardId().id());
        var collectorProvider = shardCollectorProviderFactory.create(indexShard);
        CompletableFuture<BatchIterator<Row>> it;
        try {
            it = collectorProvider.getFutureIterator(collectPhase, consumer.requiresScroll(), collectTask);
        } catch (Exception e) {
            return Exceptions.rethrowRuntimeException(e);
        }
        it.whenComplete(consumer);
    } else {
        UUID childJobId = UUIDs.dirtyUUID();
        RemoteCollector remoteCollector = new RemoteCollector(childJobId, collectTask.txnCtx().sessionSettings(), localNodeId, nodeId, transportActionProvider.transportJobInitAction(), transportActionProvider.transportKillJobsNodeAction(), searchTp, tasksService, collectTask.getRamAccounting(), consumer, createRemoteCollectPhase(childJobId, collectPhase, activePrimaryRouting.shardId(), nodeId));
        remoteCollector.doCollect();
    }
    return consumer.completionFuture().thenApply(rows -> Lists2.mapLazy(rows, Buckets.arrayToSharedRow()));
}

Example 7 with BatchIterator

use of io.crate.data.BatchIterator in project crate by crate.

the class PKLookupOperation method lookup.

public BatchIterator<Row> lookup(UUID jobId, TransactionContext txnCtx, Supplier<RamAccounting> ramAccountingSupplier, Supplier<MemoryManager> memoryManagerSupplier, boolean ignoreMissing, Map<ShardId, List<PKAndVersion>> idsByShard, Collection<? extends Projection> projections, boolean requiresScroll, Function<Doc, Row> resultToRow) {
    ArrayList<BatchIterator<Row>> iterators = new ArrayList<>(idsByShard.size());
    for (Map.Entry<ShardId, List<PKAndVersion>> idsByShardEntry : idsByShard.entrySet()) {
        ShardId shardId = idsByShardEntry.getKey();
        IndexService indexService = indicesService.indexService(shardId.getIndex());
        if (indexService == null) {
            if (ignoreMissing) {
                continue;
            }
            throw new IndexNotFoundException(shardId.getIndex());
        }
        IndexShard shard = indexService.getShardOrNull(shardId.id());
        if (shard == null) {
            if (ignoreMissing) {
                continue;
            }
            throw new ShardNotFoundException(shardId);
        }
        Stream<Row> rowStream = idsByShardEntry.getValue().stream().map(pkAndVersion -> lookupDoc(shard, pkAndVersion.id(), pkAndVersion.version(), VersionType.EXTERNAL, pkAndVersion.seqNo(), pkAndVersion.primaryTerm())).filter(Objects::nonNull).map(resultToRow);
        if (projections.isEmpty()) {
            final Iterable<Row> rowIterable = requiresScroll ? rowStream.map(row -> new RowN(row.materialize())).collect(Collectors.toList()) : rowStream::iterator;
            iterators.add(InMemoryBatchIterator.of(rowIterable, SentinelRow.SENTINEL, true));
        } else {
            ProjectorFactory projectorFactory;
            try {
                projectorFactory = shardCollectSource.getProjectorFactory(shardId);
            } catch (ShardNotFoundException e) {
                if (ignoreMissing) {
                    continue;
                }
                throw e;
            }
            Projectors projectors = new Projectors(projections, jobId, txnCtx, ramAccountingSupplier.get(), memoryManagerSupplier.get(), projectorFactory);
            final Iterable<Row> rowIterable = requiresScroll && !projectors.providesIndependentScroll() ? rowStream.map(row -> new RowN(row.materialize())).collect(Collectors.toList()) : rowStream::iterator;
            iterators.add(projectors.wrap(InMemoryBatchIterator.of(rowIterable, SentinelRow.SENTINEL, true)));
        }
    }
    // noinspection unchecked
    return CompositeBatchIterator.seqComposite(iterators.toArray(new BatchIterator[0]));
}

Example 8 with BatchIterator

use of io.crate.data.BatchIterator in project crate by crate.

the class FetchProjector method create.

public static Projector create(FetchProjection projection, RamAccounting ramAccounting, LongSupplier getBucketsBytesThreshold, TransactionContext txnCtx, NodeContext nodeCtx, FetchOperation fetchOperation) {
    final FetchRows fetchRows = FetchRows.create(txnCtx, nodeCtx, projection.fetchSources(), projection.outputSymbols());
    EstimateCellsSize estimateRowSize = new EstimateCellsSize(projection.inputTypes());
    return (BatchIterator<Row> source) -> {
        final long maxBucketsSizeInBytes = getBucketsBytesThreshold.getAsLong();
        BatchIterator<ReaderBuckets> buckets = BatchIterators.partition(source, projection.getFetchSize(), () -> new ReaderBuckets(fetchRows, projection::getFetchSourceByReader, estimateRowSize, ramAccounting), ReaderBuckets::add, readerBuckets -> readerBuckets.ramBytesUsed() > maxBucketsSizeInBytes);
        return new AsyncFlatMapBatchIterator<>(buckets, new FetchMapper(fetchOperation, projection.nodeReaders()));
    };
}

Example 9 with BatchIterator

use of io.crate.data.BatchIterator in project crate by crate.

the class NestedLoopBatchIteratorsTest method testAntiJoinRightEmpty.

@Test
public void testAntiJoinRightEmpty() throws Exception {
    Supplier<BatchIterator<Row>> batchIteratorSupplier = () -> new AntiJoinNLBatchIterator<>(new BatchSimulatingIterator<>(TestingBatchIterators.range(0, 3), 2, 2, null), InMemoryBatchIterator.empty(SENTINEL), new CombinedRow(1, 1), getCol0EqCol1JoinCondition());
    BatchIteratorTester tester = new BatchIteratorTester(batchIteratorSupplier);
    tester.verifyResultAndEdgeCaseBehaviour(Arrays.asList(new Object[] { 0 }, new Object[] { 1 }, new Object[] { 2 }));
}

Example 10 with BatchIterator

use of io.crate.data.BatchIterator in project crate by crate.

the class NestedLoopBatchIteratorsTest method testFullOuterJoin.

@Test
public void testFullOuterJoin() throws Exception {
    Supplier<BatchIterator<Row>> batchIteratorSupplier = () -> new FullOuterJoinNLBatchIterator<>(TestingBatchIterators.range(0, 4), TestingBatchIterators.range(2, 6), new CombinedRow(1, 1), getCol0EqCol1JoinCondition());
    BatchIteratorTester tester = new BatchIteratorTester(batchIteratorSupplier);
    tester.verifyResultAndEdgeCaseBehaviour(fullJoinResult);
}