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

Example 16 with TestingRowConsumer

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

the class HashInnerJoinBatchIteratorMemoryTest method testReleaseAccountingRows.

@Test
public void testReleaseAccountingRows() throws Exception {
    BatchSimulatingIterator<Row> leftIterator = new BatchSimulatingIterator<>(TestingBatchIterators.range(0, 12), 3, 3, null);
    BatchIterator<Row> rightIterator = new BatchSimulatingIterator<>(TestingBatchIterators.range(0, 10), 2, 4, null);
    when(circuitBreaker.getLimit()).thenReturn(110L);
    when(circuitBreaker.getUsed()).thenReturn(10L);
    RowAccounting<Object[]> rowAccounting = mock(RowAccounting.class);
    BatchIterator<Row> it = new HashInnerJoinBatchIterator(leftIterator, rightIterator, rowAccounting, new CombinedRow(1, 1), getCol0EqCol1JoinCondition(), getHashForLeft(), getHashForRight(), () -> 2);
    TestingRowConsumer consumer = new TestingRowConsumer();
    consumer.accept(it, null);
    consumer.getResult();
    verify(rowAccounting, times(8)).release();
    verify(rowAccounting, times(12)).accountForAndMaybeBreak(Mockito.any(Object[].class));
}

Example 17 with TestingRowConsumer

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

the class ProjectingRowConsumerTest method testErrorHandlingIfProjectorApplicationFails.

@Test
public void testErrorHandlingIfProjectorApplicationFails() throws Exception {
    WriterProjection writerProjection = new WriterProjection(Collections.singletonList(new InputColumn(0, DataTypes.STRING)), Literal.of("/x/y/z/hopefully/invalid/on/your/system/"), null, Collections.emptyMap(), Collections.emptyList(), WriterProjection.OutputFormat.JSON_OBJECT, Settings.EMPTY);
    TestingRowConsumer consumer = new TestingRowConsumer();
    RowConsumer rowConsumer = ProjectingRowConsumer.create(consumer, Collections.singletonList(writerProjection), UUID.randomUUID(), txnCtx, RamAccounting.NO_ACCOUNTING, memoryManager, projectorFactory);
    rowConsumer.accept(InMemoryBatchIterator.empty(SENTINEL), null);
    expectedException.expect(UnhandledServerException.class);
    expectedException.expectMessage("Failed to open output");
    consumer.getResult();
}

Example 18 with TestingRowConsumer

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

the class RootTaskTest method testFailureClosesAllSubContexts.

@Test
public void testFailureClosesAllSubContexts() throws Throwable {
    String localNodeId = "localNodeId";
    RoutedCollectPhase collectPhase = Mockito.mock(RoutedCollectPhase.class);
    Routing routing = Mockito.mock(Routing.class);
    when(routing.containsShards(localNodeId)).thenReturn(false);
    when(collectPhase.phaseId()).thenReturn(1);
    when(collectPhase.routing()).thenReturn(routing);
    when(collectPhase.maxRowGranularity()).thenReturn(RowGranularity.DOC);
    RootTask.Builder builder = new RootTask.Builder(logger, UUID.randomUUID(), "dummy-user", coordinatorNode, Collections.emptySet(), mock(JobsLogs.class));
    CollectTask collectChildTask = new CollectTask(collectPhase, CoordinatorTxnCtx.systemTransactionContext(), mock(MapSideDataCollectOperation.class), RamAccounting.NO_ACCOUNTING, ramAccounting -> new OnHeapMemoryManager(ramAccounting::addBytes), new TestingRowConsumer(), mock(SharedShardContexts.class), Version.CURRENT, 4096);
    TestingRowConsumer batchConsumer = new TestingRowConsumer();
    PageBucketReceiver pageBucketReceiver = new CumulativePageBucketReceiver("n1", 2, Runnable::run, new Streamer[] { IntegerType.INSTANCE.streamer() }, batchConsumer, PassThroughPagingIterator.oneShot(), 1);
    DistResultRXTask distResultRXTask = spy(new DistResultRXTask(2, "dummy", pageBucketReceiver, RamAccounting.NO_ACCOUNTING, 1));
    builder.addTask(collectChildTask);
    builder.addTask(distResultRXTask);
    RootTask rootTask = builder.build();
    Exception failure = new Exception("failure!");
    collectChildTask.kill(failure);
    // other contexts must be killed with same failure
    verify(distResultRXTask, times(1)).kill(failure);
    assertThat(rootTask.getTask(1).completionFuture().isDone(), is(true));
    assertThat(rootTask.getTask(2).completionFuture().isDone(), is(true));
}

Example 19 with TestingRowConsumer

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

the class WindowBatchIteratorTest method testWindowBatchIteratorAccountsUsedMemory.

@Test
public void testWindowBatchIteratorAccountsUsedMemory() {
    RamAccounting ramAccounting = ConcurrentRamAccounting.forCircuitBreaker("test", new NoopCircuitBreaker("dummy"));
    BatchIterator<Row> iterator = WindowFunctionBatchIterator.of(TestingBatchIterators.range(0, 10), new RowAccountingWithEstimators(List.of(DataTypes.INTEGER), ramAccounting, 32), (partitionStart, partitionEnd, currentIndex, sortedRows) -> 0, (partitionStart, partitionEnd, currentIndex, sortedRows) -> currentIndex, null, null, 1, () -> 1, Runnable::run, List.of(rowNumberWindowFunction()), List.of(), new Boolean[] { null }, new Input[][] { new Input[0] });
    TestingRowConsumer consumer = new TestingRowConsumer();
    consumer.accept(iterator, null);
    // should've accounted for 10 integers of 48 bytes each (16 for the integer, 32 for the ArrayList element)
    assertThat(ramAccounting.totalBytes(), is(480L));
}

Example 20 with TestingRowConsumer

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

the class ProjectionToProjectorVisitorTest method testFilterProjection.

@Test
public void testFilterProjection() throws Exception {
    List<Symbol> arguments = Arrays.asList(Literal.of(2), new InputColumn(1));
    EqOperator op = (EqOperator) nodeCtx.functions().get(null, EqOperator.NAME, arguments, SearchPath.pathWithPGCatalogAndDoc());
    Function function = new Function(op.signature(), arguments, EqOperator.RETURN_TYPE);
    FilterProjection projection = new FilterProjection(function, Arrays.asList(new InputColumn(0), new InputColumn(1)));
    Projector projector = visitor.create(projection, txnCtx, RamAccounting.NO_ACCOUNTING, memoryManager, UUID.randomUUID());
    assertThat(projector, instanceOf(FilterProjector.class));
    List<Object[]> rows = new ArrayList<>();
    rows.add($("human", 2));
    rows.add($("vogon", 1));
    BatchIterator<Row> filteredBI = projector.apply(InMemoryBatchIterator.of(new CollectionBucket(rows), SENTINEL, true));
    TestingRowConsumer consumer = new TestingRowConsumer();
    consumer.accept(filteredBI, null);
    Bucket bucket = consumer.getBucket();
    assertThat(bucket.size(), is(1));
}