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

Example 71 with MockRecordBatch

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

the class TestTopNEmitOutcome method testTopNResetsAfterFirstEmitOutcome.

/**
 * Verifies that if TopNBatch receives multiple non-empty record batch with EMIT outcome in between then it produces
 * output for those input batch correctly. In this case it receives first non-empty batch with OK_NEW_SCHEMA in
 * buildSchema phase followed by an empty batch with EMIT outcome. For this combination it produces output for the
 * record received so far along with EMIT outcome. Then it receives second non-empty batch with OK outcome and
 * produces output for it differently. The test validates that for each output received the order of the records are
 * correct
 */
@Test
public void testTopNResetsAfterFirstEmitOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item2").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(3, 30, "item3").addRow(2, 20, "item2").build();
    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.EMIT);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final TopN topNConfig = new TopN(null, Lists.newArrayList(ordering("id_left", RelFieldCollation.Direction.DESCENDING, RelFieldCollation.NullDirection.FIRST)), false, 10);
    final TopNBatch topNBatch = new TopNBatch(topNConfig, operatorFixture.getFragmentContext(), mockInputBatch);
    assertTrue(topNBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(topNBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertEquals(1, topNBatch.getRecordCount());
    // verify results with baseline
    RowSet actualRowSet1 = HyperRowSetImpl.fromContainer(topNBatch.getContainer(), topNBatch.getSelectionVector4());
    new RowSetComparison(expectedRowSet1).verify(actualRowSet1);
    assertTrue(topNBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, topNBatch.getRecordCount());
    // State refresh happens and limit again works on new data batches
    assertTrue(topNBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(2, topNBatch.getRecordCount());
    // verify results with baseline
    RowSet actualRowSet2 = HyperRowSetImpl.fromContainer(topNBatch.getContainer(), topNBatch.getSelectionVector4());
    new RowSetComparison(expectedRowSet2).verify(actualRowSet2);
    // Release memory for row sets
    nonEmptyInputRowSet2.clear();
    expectedRowSet2.clear();
    expectedRowSet1.clear();
}

Example 72 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testMultiLevelLateral_MultipleOutput.

/**
 * This test generates an operator tree for multi level LATERAL by stacking 2 LATERAL and finally an UNNEST pair
 * (using MockRecord Batch) as left and right child of lower level LATERAL. Then we call next() on top level
 * LATERAL to simulate the operator tree and compare the outcome and record count generated with expected values.
 * This test also changes the MAX_BATCH_ROW_COUNT to simulate the output being produced in multiple batches.
 * @throws Exception
 */
@Test
public void testMultiLevelLateral_MultipleOutput() throws Exception {
    // ** Prepare first pair of left batch and right batch for lower level LATERAL Lateral_1 **
    final LateralJoinPOP popConfig_1 = new LateralJoinPOP(null, null, JoinRelType.INNER, DrillLateralJoinRelBase.IMPLICIT_COLUMN, Lists.newArrayList());
    TupleMetadata leftSchemaWithImplicit = new SchemaBuilder().add(popConfig_1.getImplicitRIDColumn(), TypeProtos.MinorType.INT).add("id_left", TypeProtos.MinorType.INT).add("cost_left", TypeProtos.MinorType.INT).add("name_left", TypeProtos.MinorType.VARCHAR).buildSchema();
    final RowSet.SingleRowSet emptyLeftRowSet_1 = fixture.rowSetBuilder(leftSchemaWithImplicit).build();
    final RowSet.SingleRowSet nonEmptyLeftRowSet_1 = fixture.rowSetBuilder(leftSchemaWithImplicit).addRow(1, 1, 10, "item1").build();
    final RowSet.SingleRowSet nonEmptyLeftRowSet_2 = fixture.rowSetBuilder(leftSchemaWithImplicit).addRow(1, 2, 20, "item2").build();
    leftContainer.add(emptyLeftRowSet_1.container());
    leftContainer.add(nonEmptyLeftRowSet_1.container());
    leftContainer.add(nonEmptyLeftRowSet_2.container());
    // Get the left IterOutcomes for Lateral Join
    leftOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    leftOutcomes.add(RecordBatch.IterOutcome.OK);
    leftOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final CloseableRecordBatch leftMockBatch_1 = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, leftContainer, leftOutcomes, leftContainer.get(0).getSchema());
    // Get the right container with dummy data
    final RowSet.SingleRowSet nonEmptyRightRowSet_1 = fixture.rowSetBuilder(rightSchema).addRow(1, 5, 51, "item51").addRow(1, 6, 61, "item61").addRow(1, 7, 71, "item71").build();
    rightContainer.add(emptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet_1.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.EMIT);
    rightOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final CloseableRecordBatch rightMockBatch_1 = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, rightContainer, rightOutcomes, rightContainer.get(0).getSchema());
    final LateralJoinBatch lowerLateral = new LateralJoinBatch(popConfig_1, fixture.getFragmentContext(), leftMockBatch_1, rightMockBatch_1);
    // Use below api to enforce static output batch limit
    lowerLateral.setUseMemoryManager(false);
    lowerLateral.setMaxOutputRowCount(2);
    // ** Prepare second pair of left and right batch for upper LATERAL Lateral_2 **
    // Create left input schema
    TupleMetadata leftSchema2 = new SchemaBuilder().add("id_left_1", TypeProtos.MinorType.INT).add("cost_left_1", TypeProtos.MinorType.INT).add("name_left_1", TypeProtos.MinorType.VARCHAR).buildSchema();
    final RowSet.SingleRowSet emptyLeftRowSet2 = fixture.rowSetBuilder(leftSchema2).build();
    final RowSet.SingleRowSet nonEmptyLeftRowSet2 = fixture.rowSetBuilder(leftSchema2).addRow(6, 60, "item6").build();
    // Get the left container with dummy data
    final List<VectorContainer> leftContainer2 = new ArrayList<>(5);
    leftContainer2.add(emptyLeftRowSet2.container());
    leftContainer2.add(nonEmptyLeftRowSet2.container());
    // Get the left incoming batch outcomes
    final List<RecordBatch.IterOutcome> leftOutcomes2 = new ArrayList<>(5);
    leftOutcomes2.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    leftOutcomes2.add(RecordBatch.IterOutcome.OK);
    final CloseableRecordBatch leftMockBatch_2 = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, leftContainer2, leftOutcomes2, leftContainer2.get(0).getSchema());
    final LateralJoinBatch upperLateral = new LateralJoinBatch(popConfig_1, fixture.getFragmentContext(), leftMockBatch_2, lowerLateral);
    // Use below api to enforce static output batch limit
    upperLateral.setUseMemoryManager(false);
    upperLateral.setMaxOutputRowCount(2);
    try {
        final int expectedOutputRecordCount = 6;
        int actualOutputRecordCount = 0;
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == upperLateral.next());
        assertTrue(RecordBatch.IterOutcome.OK == upperLateral.next());
        actualOutputRecordCount += upperLateral.getRecordCount();
        assertTrue(RecordBatch.IterOutcome.OK == upperLateral.next());
        actualOutputRecordCount += upperLateral.getRecordCount();
        assertTrue(RecordBatch.IterOutcome.OK == upperLateral.next());
        actualOutputRecordCount += upperLateral.getRecordCount();
        assertTrue(RecordBatch.IterOutcome.NONE == upperLateral.next());
        assertTrue(actualOutputRecordCount == expectedOutputRecordCount);
    } catch (AssertionError | Exception error) {
        fail();
    } finally {
        // Close all the resources for this test case
        upperLateral.close();
        leftMockBatch_2.close();
        lowerLateral.close();
        leftMockBatch_1.close();
        rightMockBatch_1.close();
        leftContainer2.clear();
        leftOutcomes2.clear();
    }
}

Example 73 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testHandlingSchemaChangeForNonUnnestField.

/**
 * When multiple left 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 in left batch.
 * Moreover in this case the schema change was only for columns which are not produced by the UNNEST or right branch.
 *
 * @throws Exception
 */
@Test
public void testHandlingSchemaChangeForNonUnnestField() 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();
    final RowSet.SingleRowSet leftRowSet2 = fixture.rowSetBuilder(leftSchema2).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(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());
        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(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 74 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testHandlingEmptyEMITAfterOK.

/**
 * 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 empty batch. Then in this case LATERAL should produce the previous output batch
 * with EMIT outcome.
 *
 * @throws Exception
 */
@Test
public void testHandlingEmptyEMITAfterOK() throws Exception {
    // Get the left container with dummy data for Lateral Join
    leftContainer.add(nonEmptyLeftRowSet.container());
    leftContainer.add(emptyLeftRowSet.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());
    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.EMIT == 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 75 with MockRecordBatch

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

the class TestLateralJoinCorrectness method testLeftLateralJoin_WithAndWithoutMatching.

/**
 * Test to see if there are multiple rows in left batch and for some rows right side produces batch with records
 * and for other rows right side produces empty batches then based on left join type we are populating the output
 * batch correctly. Expectation is that for left rows if we find corresponding right rows then we will output both
 * using cross-join but for left rows for which there is empty right side we will produce only left row in output
 * batch. In this case all the output will be produces only in 1 record batch.
 *
 * @throws Exception
 */
@Test
public void testLeftLateralJoin_WithAndWithoutMatching() throws Exception {
    // Get the left container with dummy data for Lateral Join
    final RowSet.SingleRowSet leftRowSet2 = fixture.rowSetBuilder(leftSchema).addRow(1, 10, "item10").addRow(2, 20, "item20").addRow(3, 30, "item30").build();
    final RowSet.SingleRowSet nonEmptyRightRowSet2 = fixture.rowSetBuilder(rightSchema).addRow(3, 6, 60, "item61").addRow(3, 7, 70, "item71").addRow(3, 8, 80, "item81").build();
    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());
    rightContainer.add(nonEmptyRightRowSet.container());
    rightContainer.add(emptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet2.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.OK);
    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 LateralJoinPOP popConfig = new LateralJoinPOP(null, null, JoinRelType.LEFT, DrillLateralJoinRelBase.IMPLICIT_COLUMN, Lists.newArrayList());
    final LateralJoinBatch ljBatch = new LateralJoinBatch(popConfig, fixture.getFragmentContext(), leftMockBatch, rightMockBatch);
    try {
        // 3 for first left row and 1 for second left row
        final int expectedOutputRecordCount = 7;
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        assertTrue(ljBatch.getRecordCount() == expectedOutputRecordCount);
        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();
        nonEmptyRightRowSet2.clear();
    }
}