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

Example 36 with BatchIterator

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

the class FileReadingIteratorTest method testIteratorContract_givenCSVInputFormat__AndNoRelevantFileExtension_thenWritesAsMap.

@Test
public void testIteratorContract_givenCSVInputFormat__AndNoRelevantFileExtension_thenWritesAsMap() throws Exception {
    tempFilePath = createTempFile("tempfile", ".any-suffix");
    tmpFile = tempFilePath.toFile();
    try (OutputStreamWriter writer = new OutputStreamWriter(new FileOutputStream(tmpFile), StandardCharsets.UTF_8)) {
        writer.write("name,id,age\n");
        writer.write("Arthur,4,38\n");
        writer.write("Trillian,5,33\n");
    }
    fileUri = tempFilePath.toUri().toString();
    Supplier<BatchIterator<Row>> batchIteratorSupplier = () -> createBatchIterator(Collections.singletonList(fileUri), CSV);
    List<Object[]> expectedResult = Arrays.asList(new Object[] { CSV_AS_MAP_FIRST_LINE }, new Object[] { CSV_AS_MAP_SECOND_LINE });
    BatchIteratorTester tester = new BatchIteratorTester(batchIteratorSupplier);
    tester.verifyResultAndEdgeCaseBehaviour(expectedResult);
}

Example 37 with BatchIterator

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

the class NestedLoopBatchIteratorsTest method testSemiJoinBatchedSource.

@Test
public void testSemiJoinBatchedSource() throws Exception {
    Supplier<BatchIterator<Row>> batchIteratorSupplier = () -> new SemiJoinNLBatchIterator<>(new BatchSimulatingIterator<>(TestingBatchIterators.range(0, 5), 2, 2, null), new BatchSimulatingIterator<>(TestingBatchIterators.range(2, 6), 2, 2, null), new CombinedRow(1, 1), getCol0EqCol1JoinCondition());
    BatchIteratorTester tester = new BatchIteratorTester(batchIteratorSupplier);
    tester.verifyResultAndEdgeCaseBehaviour(semiJoinResult);
}

Example 38 with BatchIterator

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

the class NestedLoopBatchIteratorsTest method testLeftJoin.

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

Example 39 with BatchIterator

use of io.crate.data.BatchIterator 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 40 with BatchIterator

use of io.crate.data.BatchIterator 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;
}