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

Example 21 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testHandlingNonEmptyEMITAfterOK.

/**
 * Test for the case when 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
 * EMIT outcome from left side with non-empty batch. Then in this case LATERAL should produce the output batch with
 * EMIT outcome if second left and right batches are also consumed entirely in same output batch.
 *
 * @throws Exception
 */
@Test
public void testHandlingNonEmptyEMITAfterOK() throws Exception {
    // Create data for left input
    final RowSet.SingleRowSet leftRowSet2 = fixture.rowSetBuilder(leftSchema).addRow(2, 20, "item20").build();
    // Create data for right input
    final RowSet.SingleRowSet nonEmptyRightRowSet2 = fixture.rowSetBuilder(rightSchema).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.EMIT);
    // 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());
    rightContainer.add(nonEmptyRightRowSet2.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.EMIT);
    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.EMIT == ljBatch.next());
        totalRecordCount += ljBatch.getRecordCount();
        // Compare the total records generated in 2 output batches with expected count.
        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();
        nonEmptyRightRowSet2.clear();
    }
}

Example 22 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testFillingUpOutputBatch.

/**
 * Test to show LATERAL tries to pack the output batch until it's full or all the data is consumed from left and
 * right side. We have multiple left and right batches which fits after join inside the same output batch, hence
 * LATERAL only generates one output batch.
 *
 * @throws Exception
 */
@Test
public void testFillingUpOutputBatch() throws Exception {
    // Create data for left input
    final RowSet.SingleRowSet leftRowSet2 = fixture.rowSetBuilder(leftSchema).addRow(2, 20, "item20").build();
    // Create data for right input
    final RowSet.SingleRowSet nonEmptyRightRowSet2 = fixture.rowSetBuilder(rightSchema).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);
    // 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());
    rightContainer.add(nonEmptyRightRowSet2.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.EMIT);
    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() == (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();
        nonEmptyRightRowSet2.clear();
    }
}

Example 23 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testHandlingUnexpectedSchemaChangeForUnnestField.

/**
 * Verify if there is no schema change on left side and LATERAL still sees an unexpected schema change on right side
 * then it handles it correctly.
 * handle it corr
 * @throws Exception
 */
@Test
public void testHandlingUnexpectedSchemaChangeForUnnestField() 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);
    // 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());
    // second OK_NEW_SCHEMA batch. Right side batch for OK_New_Schema is always empty
    rightContainer.add(emptyRightRowSet2.container());
    rightContainer.add(nonEmptyRightRowSet2.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.OK);
    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 {
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        ljBatch.next();
        fail();
    } catch (AssertionError | Exception error) {
        // Expected since first right batch is supposed to be empty
        assertTrue(error instanceof IllegalStateException);
    } finally {
        // Close all the resources for this test case
        ljBatch.close();
        leftMockBatch.close();
        rightMockBatch.close();
        leftRowSet2.clear();
        emptyRightRowSet2.clear();
        nonEmptyRightRowSet2.clear();
    }
}

Example 24 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testOK_NEW_SCHEMA_WithNoActualSchemaChange_ForUnnestField.

/**
 * When multiple left batch is received with same schema but with OK_NEW_SCHEMA, then LATERAL correctly
 * handles it by re-creating the schema and producing multiple batches of final output
 * The schema change is for columns common on both left and right side.
 *
 * @throws Exception
 */
@Test
public void testOK_NEW_SCHEMA_WithNoActualSchemaChange_ForUnnestField() throws Exception {
    // Create data for left input
    final RowSet.SingleRowSet leftRowSet2 = fixture.rowSetBuilder(leftSchema).addRow(2, 20, "item20").build();
    // Create data for right input
    final RowSet.SingleRowSet nonEmptyRightRowSet2 = fixture.rowSetBuilder(rightSchema).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
    rightContainer.add(emptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet.container());
    rightContainer.add(emptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet2.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();
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        totalRecordCount += ljBatch.getRecordCount();
        assertTrue(totalRecordCount == (nonEmptyLeftRowSet.rowCount() * nonEmptyRightRowSet.rowCount() + leftRowSet2.rowCount() * nonEmptyRightRowSet2.rowCount()));
        // This means only 1 output record batch was received.
        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();
        nonEmptyRightRowSet2.clear();
    }
}

Example 25 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testBasicLeftLateralJoin.

/**
 * Test to check basic left lateral join is working correctly or not. We create a left batch with one and
 * corresponding right batch with zero rows and check if output still get's populated with left side of data or not.
 * Expectation is since it's a left join and even though right batch is empty the left row will be pushed to output
 * batch.
 *
 * @throws Exception
 */
@Test
public void testBasicLeftLateralJoin() 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(emptyRightRowSet.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 LateralJoinPOP popConfig = new LateralJoinPOP(null, null, JoinRelType.LEFT, DrillLateralJoinRelBase.IMPLICIT_COLUMN, Lists.newArrayList());
    final LateralJoinBatch ljBatch = new LateralJoinBatch(popConfig, fixture.getFragmentContext(), leftMockBatch, rightMockBatch);
    try {
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        assertTrue(ljBatch.getRecordCount() == nonEmptyLeftRowSet.container().getRecordCount());
        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();
    }
}