Examples with TestingBatchConsumer io.crate.testing.TestingBatchConsumer used on opensource projects

Example 31 with TestingBatchConsumer

use of io.crate.testing.TestingBatchConsumer in project crate by crate.

the class DistributingConsumerTest method testDistributingConsumerForwardsFailure.

@Test
public void testDistributingConsumerForwardsFailure() throws Exception {
    Streamer<?>[] streamers = { DataTypes.INTEGER.streamer() };
    TestingBatchConsumer collectingConsumer = new TestingBatchConsumer();
    PageDownstreamContext pageDownstreamContext = createPageDownstreamContext(streamers, collectingConsumer);
    TransportDistributedResultAction distributedResultAction = createFakeTransport(streamers, pageDownstreamContext);
    DistributingConsumer distributingConsumer = createDistributingConsumer(streamers, distributedResultAction);
    distributingConsumer.accept(null, new CompletionException(new IllegalArgumentException("foobar")));
    expectedException.expect(IllegalArgumentException.class);
    expectedException.expectMessage("foobar");
    collectingConsumer.getResult();
}

Example 32 with TestingBatchConsumer

use of io.crate.testing.TestingBatchConsumer in project crate by crate.

the class FileWriterProjectorTest method testDirectoryAsFile.

@Test
public void testDirectoryAsFile() throws Exception {
    expectedException.expect(UnhandledServerException.class);
    expectedException.expectMessage("Failed to open output: 'Output path is a directory: ");
    Path directory = createTempDir();
    FileWriterProjector fileWriterProjector = new FileWriterProjector(executorService, directory.toUri().toString(), null, null, ImmutableSet.of(), new HashMap<>(), null, WriterProjection.OutputFormat.JSON_OBJECT);
    new TestingBatchConsumer().accept(fileWriterProjector.apply(sourceSupplier.get()), null);
}

Example 33 with TestingBatchConsumer

use of io.crate.testing.TestingBatchConsumer in project crate by crate.

the class FileWriterProjectorTest method testFileAsDirectory.

@Test
public void testFileAsDirectory() throws Exception {
    expectedException.expect(UnhandledServerException.class);
    expectedException.expectMessage("Failed to open output");
    String uri = Paths.get(folder.newFile().toURI()).resolve("out.json").toUri().toString();
    FileWriterProjector fileWriterProjector = new FileWriterProjector(executorService, uri, null, null, ImmutableSet.of(), new HashMap<>(), null, WriterProjection.OutputFormat.JSON_OBJECT);
    new TestingBatchConsumer().accept(fileWriterProjector.apply(sourceSupplier.get()), null);
}

Example 34 with TestingBatchConsumer

use of io.crate.testing.TestingBatchConsumer in project crate by crate.

the class IndexWriterProjectorTest method testIndexWriter.

@Test
public void testIndexWriter() throws Throwable {
    execute("create table bulk_import (id int primary key, name string) with (number_of_replicas=0)");
    ensureGreen();
    InputCollectExpression sourceInput = new InputCollectExpression(1);
    List<CollectExpression<Row, ?>> collectExpressions = Collections.<CollectExpression<Row, ?>>singletonList(sourceInput);
    IndexWriterProjector writerProjector = new IndexWriterProjector(internalCluster().getInstance(ClusterService.class), internalCluster().getInstance(Functions.class), new IndexNameExpressionResolver(Settings.EMPTY), Settings.EMPTY, internalCluster().getInstance(TransportBulkCreateIndicesAction.class), internalCluster().getInstance(TransportShardUpsertAction.class)::execute, IndexNameResolver.forTable(new TableIdent(null, "bulk_import")), internalCluster().getInstance(BulkRetryCoordinatorPool.class), new Reference(new ReferenceIdent(bulkImportIdent, DocSysColumns.RAW), RowGranularity.DOC, DataTypes.STRING), Arrays.asList(ID_IDENT), Arrays.<Symbol>asList(new InputColumn(0)), null, null, sourceInput, collectExpressions, 20, null, null, false, false, UUID.randomUUID());
    BatchIterator rowsIterator = RowsBatchIterator.newInstance(IntStream.range(0, 100).mapToObj(i -> new RowN(new Object[] { i, new BytesRef("{\"id\": " + i + ", \"name\": \"Arthur\"}") })).collect(Collectors.toList()), 2);
    TestingBatchConsumer consumer = new TestingBatchConsumer();
    consumer.accept(writerProjector.apply(rowsIterator), null);
    Bucket objects = consumer.getBucket();
    assertThat(objects, contains(isRow(100L)));
    execute("refresh table bulk_import");
    execute("select count(*) from bulk_import");
    assertThat(response.rowCount(), is(1L));
    assertThat(response.rows()[0][0], is(100L));
}

Example 35 with TestingBatchConsumer

use of io.crate.testing.TestingBatchConsumer in project crate by crate.

the class ProjectionToProjectorVisitorTest method testGroupProjector.

@Test
public void testGroupProjector() throws Exception {
    //         in(0)  in(1)      in(0),      in(2)
    // select  race, avg(age), count(race), gender  ... group by race, gender
    List<Symbol> keys = Arrays.asList(new InputColumn(0, DataTypes.STRING), new InputColumn(2, DataTypes.STRING));
    List<Aggregation> aggregations = Arrays.asList(Aggregation.finalAggregation(avgInfo, Collections.singletonList(new InputColumn(1)), Aggregation.Step.ITER), Aggregation.finalAggregation(countInfo, Collections.singletonList(new InputColumn(0)), Aggregation.Step.ITER));
    GroupProjection projection = new GroupProjection(keys, aggregations, RowGranularity.CLUSTER);
    Projector projector = visitor.create(projection, RAM_ACCOUNTING_CONTEXT, UUID.randomUUID());
    assertThat(projector, instanceOf(GroupingProjector.class));
    // use a topN projection in order to get sorted outputs
    List<Symbol> outputs = Arrays.asList(new InputColumn(0, DataTypes.STRING), new InputColumn(1, DataTypes.STRING), new InputColumn(2, DataTypes.DOUBLE), new InputColumn(3, DataTypes.LONG));
    OrderedTopNProjection topNProjection = new OrderedTopNProjection(10, 0, outputs, ImmutableList.of(new InputColumn(2, DataTypes.DOUBLE)), new boolean[] { false }, new Boolean[] { null });
    Projector topNProjector = visitor.create(topNProjection, RAM_ACCOUNTING_CONTEXT, UUID.randomUUID());
    BytesRef human = new BytesRef("human");
    BytesRef vogon = new BytesRef("vogon");
    BytesRef male = new BytesRef("male");
    BytesRef female = new BytesRef("female");
    List<Object[]> rows = new ArrayList<>();
    rows.add($(human, 34, male));
    rows.add($(human, 22, female));
    rows.add($(vogon, 40, male));
    rows.add($(vogon, 48, male));
    rows.add($(human, 34, male));
    BatchIterator batchIterator = topNProjector.apply(projector.apply(RowsBatchIterator.newInstance(new CollectionBucket(rows), 3)));
    TestingBatchConsumer consumer = new TestingBatchConsumer();
    consumer.accept(batchIterator, null);
    Bucket bucket = consumer.getBucket();
    assertThat(bucket, contains(isRow(human, female, 22.0, 1L), isRow(human, male, 34.0, 2L), isRow(vogon, male, 44.0, 2L)));
}