[GitHub] drill pull request #788: DRILL-5318: Sub-operator test fixture

[paul-rogers <gi...@git.apache.org>]

Github user paul-rogers commented on a diff in the pull request:

    https://github.com/apache/drill/pull/788#discussion_r109541450
  
    --- Diff: exec/java-exec/src/test/java/org/apache/drill/test/QueryBuilder.java ---
    @@ -271,6 +276,91 @@ public QuerySummary run() throws Exception {
       }
     
       /**
    +   * 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();
    +      }
    +    }
    +
    +    // No results?
    +
    +    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);
    +    }
    +  }
    +
    +  /**
    +   * Run the query that is expected to return (at least) one row
    +   * with the only (or first) column returning a long value.
    +   *
    +   * @return the value of the first column of the first row
    +   * @throws RpcException if anything goes wrong
    +   */
    +
    +  public long singletonLong() throws RpcException {
    +    RowSet rowSet = rowSet();
    +    if (rowSet == null) {
    +      throw new IllegalStateException("No rows returned");
    +    }
    +    RowSetReader reader = rowSet.reader();
    +    reader.next();
    +    long value = reader.column(0).getLong();
    --- End diff --
    
    Refactored some. To reuse all code, I'd need to create a class so we can grab the typed value in the middle of the other stuff. The few lines of repetition are the price of avoiding a trivial class. Since this is test code, I'm inclined to live with the (reduced) repetition. Take a look and let me know what you think. Also added a `singletonString()` method (because, why not?) and a new `reader()` method with some of the common code.


---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastructure@apache.org or file a JIRA ticket
with INFRA.
---