Examples with VectorContainer org.apache.drill.exec.record.VectorContainer used on opensource projects

Example 81 with VectorContainer

use of org.apache.drill.exec.record.VectorContainer in project drill by axbaretto.

the class AbstractGenericCopierTest method testCopyRecords.

@Test
public void testCopyRecords() throws SchemaChangeException {
    try (RootAllocator allocator = new RootAllocator(10_000_000)) {
        final BatchSchema batchSchema = createTestSchema(BatchSchema.SelectionVectorMode.NONE);
        final RowSet srcRowSet = createSrcRowSet(allocator);
        final RowSet destRowSet = new RowSetBuilder(allocator, batchSchema).build();
        final VectorContainer destContainer = destRowSet.container();
        final Copier copier = createCopier();
        final RowSet expectedRowSet = createExpectedRowset(allocator);
        copier.setup(new RowSetBatch(srcRowSet), destContainer);
        copier.copyRecords(0, 3);
        try {
            new RowSetComparison(expectedRowSet).verify(destRowSet);
        } finally {
            srcRowSet.clear();
            if (srcRowSet instanceof RowSet.HyperRowSet) {
                ((RowSet.HyperRowSet) srcRowSet).getSv4().clear();
            }
            destRowSet.clear();
            expectedRowSet.clear();
        }
    }
}

Example 82 with VectorContainer

use of org.apache.drill.exec.record.VectorContainer in project drill by axbaretto.

the class TestCopier method testEmptyInput.

@Test
public void testEmptyInput() throws Exception {
    BatchSchema schema = SortTestUtilities.nonNullSchema();
    List<BatchGroup> batches = new ArrayList<>();
    Sort popConfig = SortTestUtilities.makeCopierConfig(Ordering.ORDER_ASC, Ordering.NULLS_UNSPECIFIED);
    OperatorContext opContext = fixture.newOperatorContext(popConfig);
    PriorityQueueCopierWrapper copier = new PriorityQueueCopierWrapper(opContext);
    VectorContainer dest = new VectorContainer();
    try {
        // TODO: Create a vector allocator to pass as last parameter so
        // that the test uses the same vector allocator as the production
        // code. Only nuisance is that we don't have the required metadata
        // readily at hand here...
        @SuppressWarnings({ "resource", "unused" }) BatchMerger merger = copier.startMerge(schema, batches, dest, 10, null);
        fail();
    } catch (AssertionError e) {
    // Expected
    } finally {
        opContext.close();
    }
}

Example 83 with VectorContainer

use of org.apache.drill.exec.record.VectorContainer in project drill by axbaretto.

the class TestSortImpl method runLargeSortTest.

/**
 * Run a full-blown sort test with multiple input batches. Because we want to
 * generate multiple inputs, we don't create them statically. Instead, we generate
 * them on the fly using a data generator. A matching data validator verifies the
 * output. Here, we are focusing on overall test flow. Separate, detailed, unit
 * tests have already probed the details of each sort component and data type,
 * so we don't need to repeat that whole exercise here; using integer keys is
 * sufficient.
 *
 * @param fixture the operator test fixture
 * @param dataGen input batch generator
 * @param validator validates output batches
 */
public void runLargeSortTest(OperatorFixture fixture, DataGenerator dataGen, DataValidator validator) {
    VectorContainer dest = new VectorContainer();
    SortImpl sort = makeSortImpl(fixture, Ordering.ORDER_ASC, Ordering.NULLS_UNSPECIFIED, dest);
    int batchCount = 0;
    RowSet input;
    while ((input = dataGen.nextRowSet()) != null) {
        batchCount++;
        if (batchCount == 1) {
            // Simulates a NEW_SCHEMA event
            sort.setSchema(input.container().getSchema());
        }
        // Simulates an OK event
        sort.addBatch(input.vectorAccessible());
    }
    // Simulate returning results
    SortResults results = sort.startMerge();
    if (results.getContainer() != dest) {
        dest.clear();
        dest = results.getContainer();
    }
    while (results.next()) {
        RowSet output = toRowSet(results, dest);
        validator.validate(output);
    }
    validator.validateDone();
    results.close();
    dest.clear();
    sort.close();
    sort.opContext().close();
}

Example 84 with VectorContainer

use of org.apache.drill.exec.record.VectorContainer in project drill by axbaretto.

the class QueryBuilder method rowSet.

/**
 * Run the query and return the first non-empty batch as a
 * {@link DirectRowSet} object that can be inspected directly
 * by the code using a {@link RowSetReader}.
 * <p>
 *
 * @see {@link #rowSetIterator()} for a version that reads a series of
 * batches as row sets.
 * @return a row set that represents the first non-empty batch returned from
 * the query
 * @throws RpcException if anything goes wrong
 */
public DirectRowSet rowSet() throws RpcException {
    // Ignore all but the first non-empty batch.
    QueryDataBatch dataBatch = null;
    for (QueryDataBatch batch : results()) {
        if (dataBatch == null && batch.getHeader().getRowCount() != 0) {
            dataBatch = batch;
        } else {
            batch.release();
        }
    }
    if (dataBatch == null) {
        return null;
    }
    // Unload the batch and convert to a row set.
    final RecordBatchLoader loader = new RecordBatchLoader(client.allocator());
    try {
        loader.load(dataBatch.getHeader().getDef(), dataBatch.getData());
        dataBatch.release();
        VectorContainer container = loader.getContainer();
        container.setRecordCount(loader.getRecordCount());
        return DirectRowSet.fromContainer(container);
    } catch (SchemaChangeException e) {
        throw new IllegalStateException(e);
    }
}

Example 85 with VectorContainer

use of org.apache.drill.exec.record.VectorContainer in project drill by axbaretto.

the class QueryRowSetIterator method next.

@Override
public DirectRowSet next() {
    if (batch == null) {
        throw new IllegalStateException();
    }
    // Unload the batch and convert to a row set.
    final RecordBatchLoader loader = new RecordBatchLoader(allocator);
    try {
        loader.load(batch.getHeader().getDef(), batch.getData());
        batch.release();
        batch = null;
        VectorContainer container = loader.getContainer();
        container.setRecordCount(loader.getRecordCount());
        return DirectRowSet.fromContainer(container);
    } catch (SchemaChangeException e) {
        throw new IllegalStateException(e);
    }
}