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

Example 56 with MockRecordBatch

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

the class TestStreamingAggEmitOutcome method t3_testStreamingAggrEmptyBatchFollowedByNonEmptyBatchEmitOutcome.

@Test
public void t3_testStreamingAggrEmptyBatchFollowedByNonEmptyBatchEmitOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(13, 130, "item13").addRow(0, 1300, "item13").addRow(2, 20, "item2").addRow(0, 2000, "item2").addRow(4, 40, "item4").addRow(0, 4000, "item4").build();
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(resultSchema).addRow("item13", (long) 1443).addRow("item2", (long) 2022).addRow("item4", (long) 4044).build();
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, parseExprs("name_left", "name"), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(3, strAggBatch.getRecordCount());
    RowSet actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
    // Release memory for row sets
    nonEmptyInputRowSet2.clear();
    expectedRowSet.clear();
}

Example 57 with MockRecordBatch

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

the class TestStreamingAggEmitOutcome method t5_testStreamingAgrResetsAfterFirstEmitOutcome.

/**
 * Verifies that if StreamingAggr 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.
 * @throws Exception
 */
@Test
public void t5_testStreamingAgrResetsAfterFirstEmitOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item2").addRow(2, 20, "item2").addRow(3, 30, "item3").addRow(3, 30, "item3").addRow(3, 30, "item3").addRow(3, 30, "item3").addRow(3, 30, "item3").addRow(3, 30, "item3").addRow(3, 30, "item3").addRow(3, 30, "item3").addRow(3, 30, "item3").addRow(3, 30, "item3").build();
    final RowSet.SingleRowSet expectedRowSet1 = operatorFixture.rowSetBuilder(resultSchema).addRow("item1", (long) 11).build();
    final RowSet.SingleRowSet expectedRowSet2 = operatorFixture.rowSetBuilder(resultSchema).addRow("item2", (long) 44).addRow("item3", (long) 330).build();
    inputContainer.add(emptyInputRowSet.container());
    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_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 StreamingAggregate streamAggrConfig = new StreamingAggregate(null, parseExprs("name_left", "name"), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertEquals(1, strAggBatch.getRecordCount());
    RowSet actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet1).verify(actualRowSet);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(2, strAggBatch.getRecordCount());
    actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet2).verify(actualRowSet);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
    // Release memory for row sets
    nonEmptyInputRowSet2.clear();
    expectedRowSet2.clear();
    expectedRowSet1.clear();
}

Example 58 with MockRecordBatch

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

the class TestStreamingAggEmitOutcome method t11_testStreamingAggrEmptyBatchEmitOutcome.

/**
 *****************************************************
 * Tests for EMIT with empty batches and no group by
 * (Tests t1-t8 are repeated with no group by)
 ******************************************************
 */
/**
 * Repeats t1_testStreamingAggrEmptyBatchEmitOutcome with no group by
 */
@Test
public void t11_testStreamingAggrEmptyBatchEmitOutcome() {
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputOutcomes.add(OK_NEW_SCHEMA);
    inputOutcomes.add(OK_NEW_SCHEMA);
    inputOutcomes.add(EMIT);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, new ArrayList<NamedExpression>(), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
}

Example 59 with MockRecordBatch

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

the class AbstractGenericCopierTest method testCopyRecords.

@Test
public void testCopyRecords() throws Exception {
    try (OperatorFixture operatorFixture = new OperatorFixture.Builder(baseDirTestWatcher).build()) {
        final BufferAllocator allocator = operatorFixture.allocator();
        final BatchSchema batchSchema = createTestSchema(BatchSchema.SelectionVectorMode.NONE);
        final RowSet srcRowSet = createSrcRowSet(allocator);
        final VectorContainer destContainer = new VectorContainer(allocator, batchSchema);
        destContainer.setRecordCount(0);
        final RowSet expectedRowSet = createExpectedRowset(allocator);
        MockRecordBatch mockRecordBatch = null;
        try {
            mockRecordBatch = new MockRecordBatch.Builder().sendData(srcRowSet).build(operatorFixture.getFragmentContext());
            mockRecordBatch.next();
            final Copier copier = createCopier(mockRecordBatch, destContainer, null);
            copier.copyRecords(0, 3);
            new RowSetComparison(expectedRowSet).verify(DirectRowSet.fromContainer(destContainer));
        } finally {
            if (mockRecordBatch != null) {
                mockRecordBatch.close();
            }
            srcRowSet.clear();
            destContainer.clear();
            expectedRowSet.clear();
        }
    }
}

Example 60 with MockRecordBatch

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

the class TestUnnestWithLateralCorrectness method testNestedUnnest.

/**
 *     Run a plan like the following for various input batches :
 *             Lateral1
 *               /    \
 *              /    Lateral2
 *            Scan      / \
 *                     /   \
 *                Project1 Project2
 *                   /       \
 *                  /         \
 *              Unnest1      Unnest2
 *
 * @param incomingSchemas
 * @param iterOutcomes
 * @param execKill
 * @param data
 * @param baseline
 * @param <T>
 * @throws Exception
 */
private <T> void testNestedUnnest(TupleMetadata[] incomingSchemas, RecordBatch.IterOutcome[] iterOutcomes, // number of batches after which to kill the execution (!)
int execKill, T[][] data, T[][][] baseline) 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 unnestPopConfig1 = new UnnestPOP(null, SchemaPath.getSimplePath("unnestColumn"), DrillUnnestRelBase.IMPLICIT_COLUMN);
    final UnnestPOP unnestPopConfig2 = new UnnestPOP(null, SchemaPath.getSimplePath("colB"), 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 unnestBatch1 = new UnnestRecordBatch(unnestPopConfig1, fixture.getFragmentContext());
    final UnnestRecordBatch unnestBatch2 = new UnnestRecordBatch(unnestPopConfig2, fixture.getFragmentContext());
    // Create intermediate Project
    final Project projectPopConfig1 = new Project(DrillLogicalTestUtils.parseExprs("unnestColumn.colB", "colB", unnestPopConfig1.getImplicitColumn(), unnestPopConfig1.getImplicitColumn()), unnestPopConfig1);
    final ProjectRecordBatch projectBatch1 = new ProjectRecordBatch(projectPopConfig1, unnestBatch1, fixture.getFragmentContext());
    final Project projectPopConfig2 = new Project(DrillLogicalTestUtils.parseExprs("colB", "unnestColumn2", unnestPopConfig2.getImplicitColumn(), unnestPopConfig2.getImplicitColumn()), unnestPopConfig2);
    final ProjectRecordBatch projectBatch2 = new ProjectRecordBatch(projectPopConfig2, unnestBatch2, fixture.getFragmentContext());
    final LateralJoinPOP ljPopConfig2 = new LateralJoinPOP(projectPopConfig1, projectPopConfig2, JoinRelType.INNER, DrillLateralJoinRelBase.IMPLICIT_COLUMN, Lists.newArrayList());
    final LateralJoinPOP ljPopConfig1 = new LateralJoinPOP(mockPopConfig, ljPopConfig2, JoinRelType.INNER, DrillLateralJoinRelBase.IMPLICIT_COLUMN, Lists.newArrayList());
    final LateralJoinBatch lateralJoinBatch2 = new LateralJoinBatch(ljPopConfig2, fixture.getFragmentContext(), projectBatch1, projectBatch2);
    final LateralJoinBatch lateralJoinBatch1 = new LateralJoinBatch(ljPopConfig1, fixture.getFragmentContext(), incomingMockBatch, lateralJoinBatch2);
    // set pointer to Lateral in unnest
    unnestBatch1.setIncoming((LateralContract) lateralJoinBatch1);
    unnestBatch2.setIncoming((LateralContract) lateralJoinBatch2);
    // Simulate the pipeline by calling next on the incoming
    // results is an array ot 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(lateralJoinBatch1.next())) {
                if (lateralJoinBatch1.getRecordCount() > 0) {
                    addBatchToResults(resultList, lateralJoinBatch1);
                }
                batchesProcessed++;
                if (batchesProcessed == execKill) {
                    lateralJoinBatch1.getContext().getExecutorState().fail(new DrillException("Testing failure of execution."));
                    lateralJoinBatch1.cancel();
                }
            // else nothing to do
            }
        } catch (UserException e) {
            throw e;
        } catch (Exception e) {
            throw new Exception("Test failed to execute lateralJoinBatch.next() because: " + 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();
            if (vectorCount != baseline[batchIndex].length + 2) {
                // baseline does not include the original unnest column(s)
                fail("Test failed in validating unnest output. Batch column count mismatch.");
            }
            for (ValueVector vv : batch) {
                if (vv.getField().getName().equals("unnestColumn") || vv.getField().getName().equals("colB")) {
                    // 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
        unnestBatch1.close();
        lateralJoinBatch1.close();
        unnestBatch2.close();
        lateralJoinBatch2.close();
        incomingMockBatch.close();
        if (results != null) {
            for (List<ValueVector> batch : results) {
                for (ValueVector vv : batch) {
                    vv.clear();
                }
            }
        }
        for (RowSet.SingleRowSet rowSet : rowSets) {
            rowSet.clear();
        }
    }
}