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

Example 36 with MockRecordBatch

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

the class TestLateralJoinCorrectnessBatchProcessing method testLeftAndRightWithInitialMissingRows_MultipleBatch.

@Test
public void testLeftAndRightWithInitialMissingRows_MultipleBatch() throws Exception {
    leftContainer.add(nonEmptyLeftRowSet.container());
    // Get the left IterOutcomes for Lateral Join
    leftOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    // Create Left MockRecordBatch
    final CloseableRecordBatch leftMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, leftContainer, leftOutcomes, leftContainer.get(0).getSchema());
    // Get the right container with dummy data
    final RowSet.SingleRowSet nonEmptyRightRowSet2 = fixture.rowSetBuilder(rightSchema).addRow(2, 22, 220, "item22").addRow(3, 33, 330, "item33").build();
    final RowSet.SingleRowSet nonEmptyRightRowSet3 = fixture.rowSetBuilder(rightSchema).addRow(4, 44, 440, "item44_1").addRow(4, 44, 440, "item44_2").build();
    final RowSet.SingleRowSet expectedRowSet = fixture.rowSetBuilder(expectedSchema).addRow(2, 20, "item2", 22, 220, "item22").addRow(3, 30, "item3", 33, 330, "item33").addRow(4, 40, "item4", 44, 440, "item44_1").addRow(4, 40, "item4", 44, 440, "item44_2").build();
    rightContainer.add(emptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet2.container());
    rightContainer.add(nonEmptyRightRowSet3.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.OK);
    rightOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final CloseableRecordBatch rightMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, rightContainer, rightOutcomes, rightContainer.get(0).getSchema());
    final LateralJoinBatch ljBatch = new LateralJoinBatch(ljPopConfig, fixture.getFragmentContext(), leftMockBatch, rightMockBatch);
    try {
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        assertTrue(ljBatch.getRecordCount() == (nonEmptyRightRowSet2.rowCount() + nonEmptyRightRowSet3.rowCount()));
        // verify results
        RowSet actualRowSet = DirectRowSet.fromContainer(ljBatch.getContainer());
        new RowSetComparison(expectedRowSet).verify(actualRowSet);
        assertTrue(RecordBatch.IterOutcome.NONE == ljBatch.next());
    } finally {
        // Close all the resources for this test case
        ljBatch.close();
        leftMockBatch.close();
        rightMockBatch.close();
        nonEmptyRightRowSet2.clear();
        expectedRowSet.clear();
    }
}

Example 37 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testPostBuildSchema_OK_NEW_SCHEMA_NonEmptyRightBatch.

/**
 * Verifies that if a non-empty batch with OK_NEW_SCHEMA is received from right side post buildSchema phase then it
 * is handled correctly by sending an empty batch with OK_NEW_SCHEMA and later consuming it to produce actual
 * output batch with some data
 */
@Test
public void testPostBuildSchema_OK_NEW_SCHEMA_NonEmptyRightBatch() throws Exception {
    // Create left input schema 2
    TupleMetadata leftSchema2 = new SchemaBuilder().add("id_left", TypeProtos.MinorType.INT).add("cost_left", TypeProtos.MinorType.VARCHAR).add("name_left", TypeProtos.MinorType.VARCHAR).buildSchema();
    // Create right input schema
    TupleMetadata rightSchema2 = new SchemaBuilder().add(ljPopConfig.getImplicitRIDColumn(), TypeProtos.MinorType.INT).add("id_right", TypeProtos.MinorType.INT).add("cost_right", TypeProtos.MinorType.VARCHAR).add("name_right", TypeProtos.MinorType.VARCHAR).buildSchema();
    // Create data for left input
    final RowSet.SingleRowSet leftRowSet2 = fixture.rowSetBuilder(leftSchema2).addRow(2, "20", "item20").build();
    // Create data for right input
    final RowSet.SingleRowSet emptyRightRowSet2 = fixture.rowSetBuilder(rightSchema2).build();
    final RowSet.SingleRowSet nonEmptyRightRowSet2 = fixture.rowSetBuilder(rightSchema2).addRow(1, 4, "41", "item41").addRow(1, 5, "51", "item51").build();
    // Get the left container with dummy data for Lateral Join
    leftContainer.add(nonEmptyLeftRowSet.container());
    leftContainer.add(leftRowSet2.container());
    // Get the left IterOutcomes for Lateral Join
    leftOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    leftOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    // Create Left MockRecordBatch
    final CloseableRecordBatch leftMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, leftContainer, leftOutcomes, leftContainer.get(0).getSchema());
    // Get the right container with dummy data
    // first OK_NEW_SCHEMA batch
    rightContainer.add(emptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet.container());
    // non-empty OK_NEW_SCHEMA batch
    rightContainer.add(nonEmptyRightRowSet2.container());
    rightContainer.add(emptyRightRowSet2.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.EMIT);
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final CloseableRecordBatch rightMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, rightContainer, rightOutcomes, rightContainer.get(0).getSchema());
    final LateralJoinBatch ljBatch = new LateralJoinBatch(ljPopConfig, fixture.getFragmentContext(), leftMockBatch, rightMockBatch);
    try {
        int totalRecordCount = 0;
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        totalRecordCount += ljBatch.getRecordCount();
        // This means 2 output record batches were received because of Schema change
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        assertEquals(0, ljBatch.getRecordCount());
        totalRecordCount += ljBatch.getRecordCount();
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        totalRecordCount += ljBatch.getRecordCount();
        assertTrue(totalRecordCount == (nonEmptyLeftRowSet.rowCount() * nonEmptyRightRowSet.rowCount() + leftRowSet2.rowCount() * nonEmptyRightRowSet2.rowCount()));
        assertTrue(RecordBatch.IterOutcome.NONE == ljBatch.next());
    } catch (AssertionError | Exception error) {
        fail();
    } finally {
        // Close all the resources for this test case
        ljBatch.close();
        leftMockBatch.close();
        rightMockBatch.close();
        leftRowSet2.clear();
        emptyRightRowSet2.clear();
        nonEmptyRightRowSet2.clear();
    }
}

Example 38 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testHandlingNoneAfterOK.

/**
 * Test for the case where LATERAL received a left batch with OK outcome and then populate the Join output in the
 * outgoing batch. There is still some space left in output batch so LATERAL call's next() on left side and receive
 * NONE outcome from left side. Then in this case LATERAL should produce the previous output batch with OK outcome
 * and then handle NONE outcome in future next() call.
 *
 * @throws Exception
 */
@Test
public void testHandlingNoneAfterOK() throws Exception {
    // Get the left container with dummy data for Lateral Join
    leftContainer.add(nonEmptyLeftRowSet.container());
    // Get the left IterOutcomes for Lateral Join
    leftOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    // Create Left MockRecordBatch
    final CloseableRecordBatch leftMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, leftContainer, leftOutcomes, leftContainer.get(0).getSchema());
    // Get the right container with dummy data
    rightContainer.add(emptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final CloseableRecordBatch rightMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, rightContainer, rightOutcomes, rightContainer.get(0).getSchema());
    final LateralJoinBatch ljBatch = new LateralJoinBatch(ljPopConfig, fixture.getFragmentContext(), leftMockBatch, rightMockBatch);
    try {
        int totalRecordCount = 0;
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        // 1st output batch is received for first EMIT from LEFT side
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        totalRecordCount += ljBatch.getRecordCount();
        // Compare the total records generated in 2 output batches with expected count.
        assertTrue(totalRecordCount == (nonEmptyLeftRowSet.rowCount() * nonEmptyRightRowSet.rowCount()));
        assertTrue(RecordBatch.IterOutcome.NONE == ljBatch.next());
    } catch (AssertionError | Exception error) {
        fail();
    } finally {
        // Close all the resources for this test case
        ljBatch.close();
        leftMockBatch.close();
        rightMockBatch.close();
    }
}

Example 39 with MockRecordBatch

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

the class TestPartitionLimitBatch method testPartitionLimitCommon.

/**
 * Common method used by all the tests for {@link PartitionLimitRecordBatch} below. It creates the MockRecordBatch
 * and {@link PartitionLimitRecordBatch} with the populated containers and outcomes list in the test. It also
 * verifies the expected outcomes list and record count populated by each test against each next() call to
 * {@link PartitionLimitRecordBatch}. For cases when the expected record count is >0 it verifies the actual output
 * returned by {@link PartitionLimitRecordBatch} with expected output rows.
 * @param start - Start offset for {@link PartitionLimit} PopConfig
 * @param end - End offset for {@link PartitionLimit} PopConfig
 */
private void testPartitionLimitCommon(Integer start, Integer end) {
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, inputContainerSv2, inputContainer.get(0).getSchema());
    final PartitionLimit limitConf = new PartitionLimit(null, start, end, PARTITION_COLUMN);
    limitBatch = new PartitionLimitRecordBatch(limitConf, operatorFixture.getFragmentContext(), mockInputBatch);
    int i = 0;
    int expectedRowSetIndex = 0;
    while (i < expectedOutcomes.size()) {
        try {
            assertTrue(expectedOutcomes.get(i) == limitBatch.next());
            assertTrue(expectedRecordCounts.get(i++) == limitBatch.getRecordCount());
            if (limitBatch.getRecordCount() > 0) {
                final RowSet actualRowSet = IndirectRowSet.fromSv2(limitBatch.getContainer(), limitBatch.getSelectionVector2());
                new RowSetComparison(expectedRowSets.get(expectedRowSetIndex++)).verify(actualRowSet);
            }
        } finally {
            limitBatch.getSelectionVector2().clear();
            limitBatch.getContainer().zeroVectors();
        }
    }
}

Example 40 with MockRecordBatch

use of org.apache.drill.exec.physical.impl.MockRecordBatch 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();
}