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

Example 61 with MockRecordBatch

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

the class TestUnnestWithLateralCorrectness method testUnnest.

// test unnest for various input conditions optionally invoking kill. if the kill or killBatch
// parameter is greater than 0 then the record batch is sent a kill after that many batches have been processed
private <T> void testUnnest(TupleMetadata[] incomingSchemas, RecordBatch.IterOutcome[] iterOutcomes, // kill unnest after every 'unnestLimit' number of values in every record
int unnestLimit, // number of batches after which to kill the execution (!)
int execKill, T[][] data, T[][][] baseline, boolean excludeUnnestColumn) throws Exception {
    // Get the incoming container with dummy data for LJ
    final List<VectorContainer> incomingContainer = new ArrayList<>(data.length);
    // Create data
    ArrayList<RowSet.SingleRowSet> rowSets = new ArrayList<>();
    int rowNumber = 0;
    int batchNum = 0;
    for (Object[] recordBatch : data) {
        RowSetBuilder rowSetBuilder = fixture.rowSetBuilder(incomingSchemas[batchNum]);
        for (Object rowData : recordBatch) {
            rowSetBuilder.addRow(++rowNumber, rowData);
        }
        RowSet.SingleRowSet rowSet = rowSetBuilder.build();
        rowSets.add(rowSet);
        incomingContainer.add(rowSet.container());
        batchNum++;
    }
    // Get the unnest POPConfig
    final UnnestPOP unnestPopConfig = new UnnestPOP(null, SchemaPath.getCompoundPath("unnestColumn"), DrillUnnestRelBase.IMPLICIT_COLUMN);
    // Get the IterOutcomes for LJ
    final List<RecordBatch.IterOutcome> outcomes = new ArrayList<>(iterOutcomes.length);
    for (RecordBatch.IterOutcome o : iterOutcomes) {
        outcomes.add(o);
    }
    // Create incoming MockRecordBatch
    final MockRecordBatch incomingMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, incomingContainer, outcomes, incomingContainer.get(0).getSchema());
    // setup Unnest record batch
    final UnnestRecordBatch unnestBatch = new UnnestRecordBatch(unnestPopConfig, fixture.getFragmentContext());
    // project is required to rename the columns so as to disambiguate the same column name from
    // unnest operator and the regular scan.
    final Project projectPopConfig = new Project(DrillLogicalTestUtils.parseExprs("unnestColumn", "unnestColumn1", unnestPopConfig.getImplicitColumn(), unnestPopConfig.getImplicitColumn()), null);
    final ProjectRecordBatch projectBatch = new ProjectRecordBatch(projectPopConfig, unnestBatch, fixture.getFragmentContext());
    final LateralJoinBatch lateralJoinBatch = new LateralJoinBatch(ljPopConfig, fixture.getFragmentContext(), incomingMockBatch, projectBatch);
    // set pointer to Lateral in unnest
    unnestBatch.setIncoming((LateralContract) lateralJoinBatch);
    // Simulate the pipeline by calling next on the incoming
    // results is an array of batches, each batch being an array of output vectors.
    List<List<ValueVector>> resultList = new ArrayList<>();
    List<List<ValueVector>> results = null;
    int batchesProcessed = 0;
    try {
        try {
            while (!isTerminal(lateralJoinBatch.next())) {
                if (lateralJoinBatch.getRecordCount() > 0) {
                    addBatchToResults(resultList, lateralJoinBatch);
                }
                batchesProcessed++;
                if (batchesProcessed == execKill) {
                    // Simulate by skipping out of the loop
                    break;
                }
            // else nothing to do
            }
        } catch (UserException e) {
            throw e;
        } catch (Exception e) {
            fail(e.getMessage());
        }
        // Check results against baseline
        results = resultList;
        int batchIndex = 0;
        int vectorIndex = 0;
        // int valueIndex = 0;
        for (List<ValueVector> batch : results) {
            int vectorCount = batch.size();
            int expectedVectorCount = (excludeUnnestColumn) ? 0 : 1;
            expectedVectorCount += baseline[batchIndex].length;
            if (vectorCount != expectedVectorCount) {
                // baseline does not include the original unnest column
                fail("Test failed in validating unnest output. Batch column count mismatch.");
            }
            for (ValueVector vv : batch) {
                if (vv.getField().getName().equals("unnestColumn")) {
                    // skip the original input column
                    continue;
                }
                int valueCount = vv.getAccessor().getValueCount();
                if (valueCount != baseline[batchIndex][vectorIndex].length) {
                    fail("Test failed in validating unnest output. Value count mismatch in batch number " + (batchIndex + 1) + "" + ".");
                }
                for (int valueIndex = 0; valueIndex < valueCount; valueIndex++) {
                    if (vv instanceof MapVector) {
                        if (!compareMapBaseline(baseline[batchIndex][vectorIndex][valueIndex], vv.getAccessor().getObject(valueIndex))) {
                            fail("Test failed in validating unnest(Map) output. Value mismatch");
                        }
                    } else if (vv instanceof VarCharVector) {
                        Object val = vv.getAccessor().getObject(valueIndex);
                        if (((String) baseline[batchIndex][vectorIndex][valueIndex]).compareTo(val.toString()) != 0) {
                            fail("Test failed in validating unnest output. Value mismatch. Baseline value[]" + vectorIndex + "][" + valueIndex + "]" + ": " + baseline[vectorIndex][valueIndex] + "   VV.getObject(valueIndex): " + val);
                        }
                    } else {
                        Object val = vv.getAccessor().getObject(valueIndex);
                        if (!baseline[batchIndex][vectorIndex][valueIndex].equals(val)) {
                            fail("Test failed in validating unnest output. Value mismatch. Baseline value[" + vectorIndex + "][" + valueIndex + "]" + ": " + baseline[batchIndex][vectorIndex][valueIndex] + "   VV.getObject(valueIndex): " + val);
                        }
                    }
                }
                vectorIndex++;
            }
            vectorIndex = 0;
            batchIndex++;
        }
    } catch (UserException e) {
        // Valid exception
        throw e;
    } catch (Exception e) {
        fail("Test failed. Exception : " + e.getMessage());
    } finally {
        // Close all the resources for this test case
        unnestBatch.close();
        lateralJoinBatch.close();
        incomingMockBatch.close();
        if (results != null) {
            for (List<ValueVector> batch : results) {
                for (ValueVector vv : batch) {
                    vv.clear();
                }
            }
        }
        for (RowSet.SingleRowSet rowSet : rowSets) {
            rowSet.clear();
        }
    }
}

Example 62 with MockRecordBatch

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

the class OperatorTestBuilderTest method executeTest.

private void executeTest(final RowSet expectedRowSet, boolean unordered, boolean combine) throws Exception {
    final MockRecordBatch inputRowSetBatch = new MockRecordBatch.Builder().sendData(buildInputData()).build(fragContext);
    final OperatorTestBuilder testBuilder = opTestBuilder().physicalOperator(createProjectPhysicalOperator());
    if (combine) {
        testBuilder.combineOutputBatches();
    }
    if (unordered) {
        testBuilder.unordered();
    }
    testBuilder.addUpstreamBatch(inputRowSetBatch).addExpectedResult(expectedRowSet).go();
}

Example 63 with MockRecordBatch

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

the class TestHashAggBatch method batchSumTest.

private void batchSumTest(int totalCount, int maxInputBatchSize, AggPrelBase.OperatorPhase phase) throws Exception {
    final HashAggregate hashAggregate = createHashAggPhysicalOperator(phase);
    final List<RowSet> inputRowSets = buildInputRowSets(TypeProtos.MinorType.INT, TypeProtos.DataMode.REQUIRED, totalCount, maxInputBatchSize);
    final MockRecordBatch.Builder rowSetBatchBuilder = new MockRecordBatch.Builder();
    inputRowSets.forEach(rowSet -> rowSetBatchBuilder.sendData(rowSet));
    final MockRecordBatch inputRowSetBatch = rowSetBatchBuilder.build(fragContext);
    final RowSet expectedRowSet = buildIntExpectedRowSet(totalCount);
    opTestBuilder().physicalOperator(hashAggregate).combineOutputBatches().unordered().addUpstreamBatch(inputRowSetBatch).addExpectedResult(expectedRowSet).go();
}

Example 64 with MockRecordBatch

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

the class TestTopNEmitOutcome method testTopN_WithEmptyNonEmptyBatchesAndOKOutcome.

/**
 ***************************************************************************************
 *   Tests for validating regular TopN behavior with no EMIT outcome
 *****************************************************************************************
 */
@Test
public void testTopN_WithEmptyNonEmptyBatchesAndOKOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(7, 70, "item7").addRow(3, 30, "item3").addRow(13, 130, "item13").build();
    final RowSet.SingleRowSet nonEmptyInputRowSet3 = operatorFixture.rowSetBuilder(inputSchema).addRow(17, 170, "item17").addRow(23, 230, "item23").addRow(130, 1300, "item130").build();
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(inputSchema).addRow(130, 1300, "item130").addRow(23, 230, "item23").addRow(17, 170, "item17").addRow(13, 130, "item13").build();
    inputContainer.add(nonEmptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet3.container());
    inputContainer.add(emptyInputRowSet.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    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, 4);
    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(4, topNBatch.getRecordCount());
    RowSet actualRowSet = HyperRowSetImpl.fromContainer(topNBatch.getContainer(), topNBatch.getSelectionVector4());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    assertTrue(topNBatch.next() == RecordBatch.IterOutcome.NONE);
    nonEmptyInputRowSet2.clear();
    nonEmptyInputRowSet3.clear();
    expectedRowSet.clear();
}

Example 65 with MockRecordBatch

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

the class TestTopNEmitOutcome method testTopNMultipleInputToSingleOutputBatch.

@Test
public void testTopNMultipleInputToSingleOutputBatch() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item2").build();
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item2").addRow(1, 10, "item1").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.OK);
    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(2, topNBatch.getRecordCount());
    // Verify results
    RowSet actualRowSet = HyperRowSetImpl.fromContainer(topNBatch.getContainer(), topNBatch.getSelectionVector4());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    assertTrue(topNBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, topNBatch.getRecordCount());
    nonEmptyInputRowSet2.clear();
}