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

Example 16 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testBuildSchemaEmptyLRBatch.

/**
 * With both left and right batch being empty, the {@link LateralJoinBatch#buildSchema()} phase
 * should still build the output container schema and return empty batch
 *
 * @throws Exception
 */
@Test
public void testBuildSchemaEmptyLRBatch() throws Exception {
    // Get the left container with dummy data for Lateral Join
    leftContainer.add(emptyLeftRowSet.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, emptyLeftRowSet.container().getSchema());
    // Get the right container with dummy data
    rightContainer.add(emptyRightRowSet.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    final CloseableRecordBatch rightMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, rightContainer, rightOutcomes, emptyRightRowSet.container().getSchema());
    final LateralJoinBatch ljBatch = new LateralJoinBatch(ljPopConfig, fixture.getFragmentContext(), leftMockBatch, rightMockBatch);
    try {
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        assertTrue(ljBatch.getRecordCount() == 0);
        while (!isTerminal(ljBatch.next())) {
        // do nothing
        }
    } catch (AssertionError | Exception error) {
        fail();
    } finally {
        // Close all the resources for this test case
        ljBatch.close();
        leftMockBatch.close();
        rightMockBatch.close();
    }
}

Example 17 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testLateral_SchemaChange_Left_EmptyRightBatchForFirst.

/**
 * Test to verify that for first left incoming if there is no right side incoming batch and then second left
 * incoming comes with schema change, then the schema change with empty output batch for first incoming is handled
 * properly.
 * @throws Exception
 */
@Test
public void testLateral_SchemaChange_Left_EmptyRightBatchForFirst() 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 data for left input
    final RowSet.SingleRowSet leftRowSet2 = fixture.rowSetBuilder(leftSchema2).addRow(2, "20", "item20").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(emptyRightRowSet.container());
    // non-empty OK_NEW_SCHEMA batch
    rightContainer.add(nonEmptyRightRowSet.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_NEW_SCHEMA == ljBatch.next());
        // This means 2 output record batches were received because of Schema change
        assertEquals(3, ljBatch.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();
        leftRowSet2.clear();
    }
}

Example 18 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testFillingUpOutputBatch_WithExcludedColumns.

@Test
public void testFillingUpOutputBatch_WithExcludedColumns() 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();
    TupleMetadata expectedSchema = new SchemaBuilder().add("id_left", TypeProtos.MinorType.INT).add("name_left", TypeProtos.MinorType.VARCHAR).add("id_right", TypeProtos.MinorType.INT).add("cost_right", TypeProtos.MinorType.INT).add("name_right", TypeProtos.MinorType.VARCHAR).buildSchema();
    final RowSet.SingleRowSet expectedRowSet = fixture.rowSetBuilder(expectedSchema).addRow(1, "item1", 1, 11, "item11").addRow(1, "item1", 2, 21, "item21").addRow(1, "item1", 3, 31, "item31").addRow(2, "item20", 4, 41, "item41").addRow(2, "item20", 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());
    List<SchemaPath> excludedCols = new ArrayList<>();
    excludedCols.add(SchemaPath.getSimplePath("cost_left"));
    try {
        testExcludedColumns(excludedCols, leftMockBatch, rightMockBatch, expectedRowSet);
    } finally {
        // Close all the resources for this test case
        leftRowSet2.clear();
        nonEmptyRightRowSet2.clear();
    }
}

Example 19 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testHandlingSchemaChangeForUnnestField.

/**
 * When multiple left and right batches are received with different schema, then LATERAL produces output for each
 * schema type separately (even though output batch is not filled completely) and handles the schema change correctly.
 * Moreover in this case the schema change was for both left and right branch (produced by UNNEST with empty batch).
 *
 * @throws Exception
 */
@Test
public void testHandlingSchemaChangeForUnnestField() 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());
    // 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.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());
        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();
        throw error;
    } finally {
        // Close all the resources for this test case
        ljBatch.close();
        leftMockBatch.close();
        rightMockBatch.close();
        leftRowSet2.clear();
        emptyRightRowSet2.clear();
        nonEmptyRightRowSet2.clear();
    }
}

Example 20 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testHandlingNonEmpty_EMITAfterOK_WithMultipleOutput.

/**
 * Temporary test to validate LATERAL handling output batch getting filled without consuming full output from left
 * and right batch join.
 * <p>
 * For this test we are updating {@link LateralJoinBatch#MAX_BATCH_ROW_COUNT} by making it public, which might not expected
 * after including the BatchSizing logic
 * TODO: Update the test after incorporating the BatchSizing change.
 *
 * @throws Exception
 */
@Test
public void testHandlingNonEmpty_EMITAfterOK_WithMultipleOutput() 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);
    final int maxBatchSize = 2;
    ljBatch.setUseMemoryManager(false);
    ljBatch.setMaxOutputRowCount(maxBatchSize);
    try {
        int totalRecordCount = 0;
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        // 1st output batch
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        totalRecordCount += ljBatch.getRecordCount();
        assertTrue(ljBatch.getRecordCount() == maxBatchSize);
        // 2nd output batch
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        assertTrue(ljBatch.getRecordCount() == maxBatchSize);
        totalRecordCount += ljBatch.getRecordCount();
        // 3rd output batch
        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();
    }
}