Examples with DirectRowSet org.apache.drill.test.rowSet.DirectRowSet used on opensource projects

Example 1 with DirectRowSet

use of org.apache.drill.test.rowSet.DirectRowSet in project drill by apache.

the class QueryBuilder method rowSet.

/**
   * Run the query and return the first result set as a
   * {@link DirectRowSet} object that can be inspected directly
   * by the code using a {@link RowSetReader}.
   * <p>
   * An enhancement is to provide a way to read a series of result
   * batches as row sets.
   * @return a row set that represents the first 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 new DirectRowSet(client.allocator(), container);
    } catch (SchemaChangeException e) {
        throw new IllegalStateException(e);
    }
}