Examples with StreamingAggBatch org.apache.drill.exec.physical.impl.aggregate.StreamingAggBatch used on opensource projects

Example 1 with StreamingAggBatch

use of org.apache.drill.exec.physical.impl.aggregate.StreamingAggBatch in project drill by apache.

the class TestStreamingAggEmitOutcome method t2_testStreamingAggrNonEmptyBatchEmitOutcome.

/**
 * Verifies that if StreamingAgg receives a RecordBatch with EMIT outcome post build schema phase then it produces
 * output for those input batch correctly. The first output batch will always be returned with OK_NEW_SCHEMA
 * outcome followed by EMIT with empty batch. The test verifies the output order with the expected baseline.
 */
@Test
public void t2_testStreamingAggrNonEmptyBatchEmitOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(13, 130, "item13").addRow(13, 130, "item13").addRow(2, 20, "item2").addRow(2, 20, "item2").addRow(4, 40, "item4").build();
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(resultSchema).addRow("item1", (long) 11).addRow("item13", (long) 286).addRow("item2", (long) 44).addRow("item4", (long) 44).build();
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, parseExprs("name_left", "name"), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    // Data before EMIT is returned with an OK_NEW_SCHEMA.
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertEquals(4, strAggBatch.getRecordCount());
    RowSet actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    // EMIT comes with an empty batch
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
    // Release memory for row sets
    nonEmptyInputRowSet2.clear();
    expectedRowSet.clear();
}

Example 2 with StreamingAggBatch

use of org.apache.drill.exec.physical.impl.aggregate.StreamingAggBatch in project drill by apache.

the class TestStreamingAggEmitOutcome method t17_testStreamingAggrMultipleEMITOutcome.

/**
 *   Repeats t7_testStreamingAggrMultipleEMITOutcome with no group by
 */
@Test
public void t17_testStreamingAggrMultipleEMITOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item2").addRow(3, 30, "item3").build();
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputContainer.add(emptyInputRowSet.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, new ArrayList<NamedExpression>(), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertEquals(1, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(1, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(1, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
    nonEmptyInputRowSet2.clear();
}

Example 3 with StreamingAggBatch

use of org.apache.drill.exec.physical.impl.aggregate.StreamingAggBatch in project drill by apache.

the class TestStreamingAggEmitOutcome method t10_1_testStreamingAggr_InputSplitToMultipleOutputBatch.

@Test
public void t10_1_testStreamingAggr_InputSplitToMultipleOutputBatch() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(1, 20, "item1").build();
    final RowSet.SingleRowSet nonEmptyInputRowSet3 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 30, "item2").build();
    final RowSet.SingleRowSet nonEmptyInputRowSet4 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 40, "item2").addRow(2, 50, "item2").addRow(2, 60, "item2").addRow(2, 70, "item2").addRow(3, 100, "item3").addRow(3, 200, "item3").addRow(3, 300, "item3").addRow(3, 400, "item3").build();
    TupleMetadata resultSchema2 = new SchemaBuilder().add("name", TypeProtos.MinorType.VARCHAR).add("id", TypeProtos.MinorType.INT).add("total_count", TypeProtos.MinorType.BIGINT).buildSchema();
    final RowSet.SingleRowSet expectedRowSet1 = operatorFixture.rowSetBuilder(resultSchema2).addRow("item1", 1, (long) 2).addRow("item2", 2, (long) 5).build();
    final RowSet.SingleRowSet expectedRowSet2 = operatorFixture.rowSetBuilder(resultSchema2).addRow("item3", 3, (long) 4).build();
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputContainer.add(nonEmptyInputRowSet3.container());
    inputContainer.add(nonEmptyInputRowSet4.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, parseExprs("name_left", "name", "id_left", "id"), parseExprs("count(cost_left)", "total_count"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    strAggBatch.setMaxOutputRowCount(2);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertEquals(2, strAggBatch.getRecordCount());
    RowSet actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet1).verify(actualRowSet);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK);
    assertEquals(1, strAggBatch.getRecordCount());
    actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet2).verify(actualRowSet);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
    nonEmptyInputRowSet2.clear();
    nonEmptyInputRowSet3.clear();
    nonEmptyInputRowSet4.clear();
    expectedRowSet1.clear();
    expectedRowSet2.clear();
}

Example 4 with StreamingAggBatch

use of org.apache.drill.exec.physical.impl.aggregate.StreamingAggBatch in project drill by apache.

the class TestStreamingAggEmitOutcome method t13_testStreamingAggrEmptyBatchFollowedByNonEmptyBatchEmitOutcome.

/**
 *   Repeats t3_testStreamingAggrEmptyBatchFollowedByNonEmptyBatchEmitOutcome with no group by
 */
@Test
public void t13_testStreamingAggrEmptyBatchFollowedByNonEmptyBatchEmitOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(13, 130, "item13").addRow(0, 1300, "item13").addRow(2, 20, "item2").addRow(0, 2000, "item2").addRow(4, 40, "item4").addRow(0, 4000, "item4").build();
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(resultSchemaNoGroupBy).addRow((long) 7509).build();
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, new ArrayList<NamedExpression>(), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(1, strAggBatch.getRecordCount());
    RowSet actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
    // Release memory for row sets
    nonEmptyInputRowSet2.clear();
    expectedRowSet.clear();
}

Example 5 with StreamingAggBatch

use of org.apache.drill.exec.physical.impl.aggregate.StreamingAggBatch in project drill by apache.

the class TestStreamingAggEmitOutcome method t21_testStreamingAggrRunsofEmpty_NonEmpty_Sv2.

/**
 *   Repeats t21_testStreamingAggrRunsofEmpty_NonEmpty_Sv2 with no group by
 */
@Test
public void t21_testStreamingAggrRunsofEmpty_NonEmpty_Sv2() {
    TupleMetadata inputSchema_sv2 = new SchemaBuilder().add("id_left", TypeProtos.MinorType.INT).add("cost_left", TypeProtos.MinorType.INT).add("name_left", TypeProtos.MinorType.VARCHAR).buildSchema();
    final RowSet.SingleRowSet emptyRowSet_Sv2 = operatorFixture.rowSetBuilder(inputSchema_sv2).withSv2().build();
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema_sv2).addSelection(false, 2, 20, "item2").addSelection(true, 3, 30, "item3").withSv2().build();
    inputContainer.add(emptyRowSet_Sv2.container());
    inputContainer.add(emptyRowSet_Sv2.container());
    inputContainer.add(emptyRowSet_Sv2.container());
    inputContainer.add(emptyRowSet_Sv2.container());
    inputContainer.add(emptyRowSet_Sv2.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    inputContainerSv2.add(emptyRowSet_Sv2.getSv2());
    inputContainerSv2.add(emptyRowSet_Sv2.getSv2());
    inputContainerSv2.add(emptyRowSet_Sv2.getSv2());
    inputContainerSv2.add(emptyRowSet_Sv2.getSv2());
    inputContainerSv2.add(emptyRowSet_Sv2.getSv2());
    inputContainerSv2.add(nonEmptyInputRowSet2.getSv2());
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, inputContainerSv2, inputContainer.get(0).getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, new ArrayList<>(), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(resultSchemaNoGroupBy).addRow((long) 33).build();
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertEquals(0, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    // For special batch.
    assertEquals(1, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(1, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(1, strAggBatch.getRecordCount());
    RowSet actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    nonEmptyInputRowSet2.clear();
    emptyRowSet_Sv2.clear();
    expectedRowSet.clear();
}