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

Example 1 with PartitionLimitRecordBatch

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

the class TestPartitionLimitBatch method testPartitionLimitCommon.

/**
 * Common method used by all the tests for {@link PartitionLimitRecordBatch} below. It creates the MockRecordBatch
 * and {@link PartitionLimitRecordBatch} with the populated containers and outcomes list in the test. It also
 * verifies the expected outcomes list and record count populated by each test against each next() call to
 * {@link PartitionLimitRecordBatch}. For cases when the expected record count is >0 it verifies the actual output
 * returned by {@link PartitionLimitRecordBatch} with expected output rows.
 * @param start - Start offset for {@link PartitionLimit} PopConfig
 * @param end - End offset for {@link PartitionLimit} PopConfig
 */
private void testPartitionLimitCommon(Integer start, Integer end) {
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, inputContainerSv2, inputContainer.get(0).getSchema());
    final PartitionLimit limitConf = new PartitionLimit(null, start, end, PARTITION_COLUMN);
    limitBatch = new PartitionLimitRecordBatch(limitConf, operatorFixture.getFragmentContext(), mockInputBatch);
    int i = 0;
    int expectedRowSetIndex = 0;
    while (i < expectedOutcomes.size()) {
        try {
            assertTrue(expectedOutcomes.get(i) == limitBatch.next());
            assertTrue(expectedRecordCounts.get(i++) == limitBatch.getRecordCount());
            if (limitBatch.getRecordCount() > 0) {
                final RowSet actualRowSet = IndirectRowSet.fromSv2(limitBatch.getContainer(), limitBatch.getSelectionVector2());
                new RowSetComparison(expectedRowSets.get(expectedRowSetIndex++)).verify(actualRowSet);
            }
        } finally {
            limitBatch.getSelectionVector2().clear();
            limitBatch.getContainer().zeroVectors();
        }
    }
}