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

Example 11 with MockRecordBatch

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

the class TestSortEmitOutcome method testSortEmptyBatchFollowedByNonEmptyBatchEmitOutcome.

/**
 * Verifies ExternalSortBatch behavior when it receives first incoming batch post buildSchema phase as empty batch
 * with EMIT outcome followed by non-empty batch with EMIT outcome. Expectation is sort will handle the EMIT
 * boundary correctly and produce 2 empty output batch for first EMIT outcome and 1 non-empty output batch for second
 * EMIT outcome.
 */
@Test
public void testSortEmptyBatchFollowedByNonEmptyBatchEmitOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item2").addRow(13, 130, "item13").addRow(4, 40, "item4").build();
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item2").addRow(4, 40, "item4").addRow(13, 130, "item13").build();
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputOutcomes.add(OK_NEW_SCHEMA);
    inputOutcomes.add(EMIT);
    inputOutcomes.add(EMIT);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    sortBatch = new ExternalSortBatch(sortPopConfig, operatorFixture.getFragmentContext(), mockInputBatch);
    // BuildSchema phase
    assertTrue(sortBatch.next() == OK_NEW_SCHEMA);
    outputRecordCount += sortBatch.getRecordCount();
    assertEquals(0, outputRecordCount);
    // Output for first empty EMIT batch
    assertTrue(sortBatch.next() == OK_NEW_SCHEMA);
    assertTrue(sortBatch.next() == EMIT);
    outputRecordCount += sortBatch.getRecordCount();
    assertEquals(0, outputRecordCount);
    // Output for second non-empty EMIT batch
    assertTrue(sortBatch.next() == EMIT);
    outputRecordCount += sortBatch.getRecordCount();
    assertEquals(3, outputRecordCount);
    // verify results
    RowSet actualRowSet = HyperRowSetImpl.fromContainer(sortBatch.getContainer(), sortBatch.getSelectionVector4());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    // Release memory for row sets
    nonEmptyInputRowSet2.clear();
    expectedRowSet.clear();
}

Example 12 with MockRecordBatch

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

the class TestSortEmitOutcome method testSort_NonEmptyFirst_EmptyOKEmitOutcome.

/**
 * Verifies ExternalSortBatch behavior when it receives incoming batches with different IterOutcomes like
 * OK_NEW_SCHEMA / OK / EMIT / NONE
 */
@Test
public void testSort_NonEmptyFirst_EmptyOKEmitOutcome() {
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(inputSchema).addRow(1, 10, "item1").build();
    inputContainer.add(nonEmptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputOutcomes.add(OK_NEW_SCHEMA);
    inputOutcomes.add(OK);
    inputOutcomes.add(EMIT);
    inputOutcomes.add(NONE);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    sortBatch = new ExternalSortBatch(sortPopConfig, operatorFixture.getFragmentContext(), mockInputBatch);
    // BuildSchema phase output
    assertTrue(sortBatch.next() == OK_NEW_SCHEMA);
    assertEquals(0, sortBatch.getRecordCount());
    // Output batch 1 for first 3 input batches with OK_NEW_SCHEMA/OK/EMIT outcome
    assertTrue(sortBatch.next() == OK_NEW_SCHEMA);
    assertEquals(1, sortBatch.getRecordCount());
    // verify results
    RowSet actualRowSet = HyperRowSetImpl.fromContainer(sortBatch.getContainer(), sortBatch.getSelectionVector4());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    // Output batch 2 for first 3 input batches with OK_NEW_SCHEMA/OK/EMIT outcome
    assertTrue(sortBatch.next() == EMIT);
    assertEquals(0, sortBatch.getRecordCount());
    // Output batch for NONE outcome
    assertTrue(sortBatch.next() == NONE);
    // Release memory for row set
    expectedRowSet.clear();
}

Example 13 with MockRecordBatch

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

the class TestSortEmitOutcome method testTopNMultipleOutputBatch.

@Test
public void testTopNMultipleOutputBatch() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(4, 40, "item4").addRow(2, 20, "item2").addRow(5, 50, "item5").addRow(3, 30, "item3").build();
    final RowSet.SingleRowSet expectedRowSet1 = operatorFixture.rowSetBuilder(inputSchema).addRow(1, 10, "item1").build();
    final RowSet.SingleRowSet expectedRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item2").addRow(3, 30, "item3").addRow(4, 40, "item4").addRow(5, 50, "item5").build();
    inputContainer.add(nonEmptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputOutcomes.add(OK_NEW_SCHEMA);
    inputOutcomes.add(EMIT);
    inputOutcomes.add(OK);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    sortBatch = new ExternalSortBatch(sortPopConfig, operatorFixture.getFragmentContext(), mockInputBatch);
    // BuildSchema phase output
    assertTrue(sortBatch.next() == OK_NEW_SCHEMA);
    // Output batch 1 for first EMIT outcome
    assertTrue(sortBatch.next() == OK_NEW_SCHEMA);
    assertEquals(1, sortBatch.getRecordCount());
    // verify results
    RowSet actualRowSet1 = HyperRowSetImpl.fromContainer(sortBatch.getContainer(), sortBatch.getSelectionVector4());
    new RowSetComparison(expectedRowSet1).verify(actualRowSet1);
    // Output batch 2 for first EMIT outcome
    assertTrue(sortBatch.next() == EMIT);
    assertEquals(0, sortBatch.getRecordCount());
    // Output batch for OK outcome
    assertTrue(sortBatch.next() == OK);
    assertEquals(4, sortBatch.getRecordCount());
    // verify results
    RowSet actualRowSet2 = HyperRowSetImpl.fromContainer(sortBatch.getContainer(), sortBatch.getSelectionVector4());
    new RowSetComparison(expectedRowSet2).verify(actualRowSet2);
    // Output batch for NONE outcome
    assertTrue(sortBatch.next() == NONE);
    // Release memory for row sets
    nonEmptyInputRowSet2.clear();
    expectedRowSet2.clear();
    expectedRowSet1.clear();
}

Example 14 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testUnsupportedSelectionVector.

/**
 * Test unsupported incoming batch to LATERAL with SelectionVector
 * @throws Exception
 */
@Test
public void testUnsupportedSelectionVector() throws Exception {
    final RowSet.SingleRowSet leftRowSet2 = fixture.rowSetBuilder(leftSchema).addRow(2, 20, "item20").withSv2().build();
    // Get the left container with dummy data for Lateral Join
    leftContainer.add(leftRowSet2.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());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    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 {
        ljBatch.next();
        fail();
    } catch (UserException e) {
        assertEquals(ErrorType.UNSUPPORTED_OPERATION, e.getErrorType());
    } catch (AssertionError | Exception error) {
        fail();
    } finally {
        // Close all the resources for this test case
        ljBatch.close();
        leftMockBatch.close();
        rightMockBatch.close();
        leftRowSet2.clear();
    }
}

Example 15 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testFillingUpOutputBatch_With2ExcludedColumns.

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