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

Example 16 with CollectionBucket

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

the class MapSideDataCollectOperationTest method testFileUriCollect.

@Test
public void testFileUriCollect() throws Exception {
    ClusterService clusterService = new NoopClusterService();
    Functions functions = getFunctions();
    CollectSourceResolver collectSourceResolver = mock(CollectSourceResolver.class);
    when(collectSourceResolver.getService(any(RoutedCollectPhase.class))).thenReturn(new FileCollectSource(functions, clusterService, Collections.<String, FileInputFactory>emptyMap()));
    MapSideDataCollectOperation collectOperation = new MapSideDataCollectOperation(collectSourceResolver, threadPool);
    File tmpFile = temporaryFolder.newFile("fileUriCollectOperation.json");
    try (OutputStreamWriter writer = new OutputStreamWriter(new FileOutputStream(tmpFile), StandardCharsets.UTF_8)) {
        writer.write("{\"name\": \"Arthur\", \"id\": 4, \"details\": {\"age\": 38}}\n");
        writer.write("{\"id\": 5, \"name\": \"Trillian\", \"details\": {\"age\": 33}}\n");
    }
    FileUriCollectPhase collectNode = new FileUriCollectPhase(UUID.randomUUID(), 0, "test", Collections.singletonList("noop_id"), Literal.of(Paths.get(tmpFile.toURI()).toUri().toString()), Arrays.<Symbol>asList(createReference("name", DataTypes.STRING), createReference(new ColumnIdent("details", "age"), DataTypes.INTEGER)), Collections.emptyList(), null, false);
    String threadPoolName = JobCollectContext.threadPoolName(collectNode, "noop_id");
    TestingBatchConsumer consumer = new TestingBatchConsumer();
    JobCollectContext jobCollectContext = mock(JobCollectContext.class);
    CrateCollector collectors = collectOperation.createCollector(collectNode, consumer, jobCollectContext);
    collectOperation.launchCollector(collectors, threadPoolName);
    assertThat(new CollectionBucket(consumer.getResult()), contains(isRow("Arthur", 38), isRow("Trillian", 33)));
}

Example 17 with CollectionBucket

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

the class MapSideDataCollectOperationTest method testFileUriCollect.

@Test
public void testFileUriCollect() throws Exception {
    FileCollectSource fileCollectSource = new FileCollectSource(createNodeContext(), clusterService, Collections.emptyMap());
    File tmpFile = temporaryFolder.newFile("fileUriCollectOperation.json");
    try (OutputStreamWriter writer = new OutputStreamWriter(new FileOutputStream(tmpFile), StandardCharsets.UTF_8)) {
        writer.write("{\"name\": \"Arthur\", \"id\": 4, \"details\": {\"age\": 38}}\n");
        writer.write("{\"id\": 5, \"name\": \"Trillian\", \"details\": {\"age\": 33}}\n");
    }
    FileUriCollectPhase collectNode = new FileUriCollectPhase(UUID.randomUUID(), 0, "test", Collections.singletonList("noop_id"), Literal.of(Paths.get(tmpFile.toURI()).toUri().toString()), Arrays.asList(createReference("name", DataTypes.STRING), createReference(new ColumnIdent("details", "age"), DataTypes.INTEGER)), Collections.emptyList(), null, false, CopyFromParserProperties.DEFAULT, FileUriCollectPhase.InputFormat.JSON, Settings.EMPTY);
    TestingRowConsumer consumer = new TestingRowConsumer();
    CollectTask collectTask = mock(CollectTask.class);
    BatchIterator<Row> iterator = fileCollectSource.getIterator(CoordinatorTxnCtx.systemTransactionContext(), collectNode, collectTask, false).get(5, TimeUnit.SECONDS);
    consumer.accept(iterator, null);
    assertThat(new CollectionBucket(consumer.getResult()), contains(isRow("Arthur", 38), isRow("Trillian", 33)));
}

Example 18 with CollectionBucket

use of io.crate.data.CollectionBucket 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(new Aggregation(avgSignature, avgSignature.getReturnType().createType(), Collections.singletonList(new InputColumn(1))), new Aggregation(CountAggregation.SIGNATURE, CountAggregation.SIGNATURE.getReturnType().createType(), Collections.singletonList(new InputColumn(0))));
    GroupProjection projection = new GroupProjection(keys, aggregations, AggregateMode.ITER_FINAL, RowGranularity.CLUSTER);
    Projector projector = visitor.create(projection, txnCtx, RamAccounting.NO_ACCOUNTING, memoryManager, 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, List.of(new InputColumn(2, DataTypes.DOUBLE)), new boolean[] { false }, new boolean[] { false });
    Projector topNProjector = visitor.create(topNProjection, txnCtx, RamAccounting.NO_ACCOUNTING, memoryManager, UUID.randomUUID());
    String human = "human";
    String vogon = "vogon";
    String male = "male";
    String female = "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<Row> batchIterator = topNProjector.apply(projector.apply(InMemoryBatchIterator.of(new CollectionBucket(rows), SENTINEL, true)));
    TestingRowConsumer consumer = new TestingRowConsumer();
    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)));
}

Example 19 with CollectionBucket

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

the class DistResultRXTaskTest method testCantSetSameBucketTwiceWithoutReceivingFullPage.

@Test
public void testCantSetSameBucketTwiceWithoutReceivingFullPage() throws Throwable {
    TestingRowConsumer batchConsumer = new TestingRowConsumer();
    DistResultRXTask ctx = getPageDownstreamContext(batchConsumer, PassThroughPagingIterator.oneShot(), 3);
    PageResultListener pageResultListener = mock(PageResultListener.class);
    Bucket bucket = new CollectionBucket(Collections.singletonList(new Object[] { "foo" }));
    PageBucketReceiver bucketReceiver = ctx.getBucketReceiver((byte) 0);
    assertThat(bucketReceiver, notNullValue());
    bucketReceiver.setBucket(1, bucket, false, pageResultListener);
    bucketReceiver.setBucket(1, bucket, false, pageResultListener);
    expectedException.expect(IllegalStateException.class);
    expectedException.expectMessage("Same bucket of a page set more than once. node=n1 method=setBucket phaseId=1 bucket=1");
    batchConsumer.getResult();
}

Example 20 with CollectionBucket

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

the class DistResultRXTaskTest method testNonSequentialBucketIds.

@Test
public void testNonSequentialBucketIds() throws Exception {
    TestingRowConsumer batchConsumer = new TestingRowConsumer();
    DistResultRXTask ctx = getPageDownstreamContext(batchConsumer, PassThroughPagingIterator.oneShot(), 3);
    PageBucketReceiver bucketReceiver = ctx.getBucketReceiver((byte) 0);
    assertThat(bucketReceiver, notNullValue());
    final PageResultListener mockListener = mock(PageResultListener.class);
    Bucket b1 = new CollectionBucket(Collections.singletonList(new Object[] { "foo" }));
    bucketReceiver.setBucket(0, b1, true, mockListener);
    Bucket b2 = new CollectionBucket(Collections.singletonList(new Object[] { "bar" }));
    bucketReceiver.setBucket(3, b2, true, mockListener);
    Bucket b3 = new CollectionBucket(Collections.singletonList(new Object[] { "universe" }));
    CheckPageResultListener checkPageResultListener = new CheckPageResultListener();
    bucketReceiver.setBucket(42, b3, false, checkPageResultListener);
    assertThat(checkPageResultListener.needMoreResult, is(true));
    bucketReceiver.setBucket(42, b3, true, checkPageResultListener);
    assertThat(checkPageResultListener.needMoreResult, is(false));
    List<Object[]> result = batchConsumer.getResult();
    assertThat(result.toArray(), arrayContainingInAnyOrder(new Object[] { "foo" }, new Object[] { "bar" }, new Object[] { "universe" }, new Object[] { "universe" }));
}