Examples with QueryDataBatch org.apache.drill.exec.rpc.user.QueryDataBatch used on opensource projects

Example 16 with QueryDataBatch

use of org.apache.drill.exec.rpc.user.QueryDataBatch in project drill by axbaretto.

the class DrillTestWrapper method addToMaterializedResults.

public static void addToMaterializedResults(List<Map<String, Object>> materializedRecords, List<QueryDataBatch> records, RecordBatchLoader loader) throws SchemaChangeException, UnsupportedEncodingException {
    long totalRecords = 0;
    QueryDataBatch batch;
    int size = records.size();
    for (int i = 0; i < size; i++) {
        batch = records.get(0);
        loader.load(batch.getHeader().getDef(), batch.getData());
        // TODO:  Clean:  DRILL-2933:  That load(...) no longer throws
        // SchemaChangeException, so check/clean throws clause above.
        logger.debug("reading batch with " + loader.getRecordCount() + " rows, total read so far " + totalRecords);
        totalRecords += loader.getRecordCount();
        for (int j = 0; j < loader.getRecordCount(); j++) {
            Map<String, Object> record = new TreeMap<>();
            for (VectorWrapper<?> w : loader) {
                Object obj = w.getValueVector().getAccessor().getObject(j);
                if (obj != null) {
                    if (obj instanceof Text) {
                        obj = obj.toString();
                    }
                    record.put(SchemaPath.getSimplePath(w.getField().getName()).toExpr(), obj);
                }
                record.put(SchemaPath.getSimplePath(w.getField().getName()).toExpr(), obj);
            }
            materializedRecords.add(record);
        }
        records.remove(0);
        batch.release();
        loader.clear();
    }
}

Example 17 with QueryDataBatch

use of org.apache.drill.exec.rpc.user.QueryDataBatch in project drill by axbaretto.

the class DrillTestWrapper method compareUnorderedResults.

/**
 * Use this method only if necessary to validate one query against another. If you are just validating against a
 * baseline file use one of the simpler interfaces that will write the validation query for you.
 *
 * @throws Exception
 */
protected void compareUnorderedResults() throws Exception {
    RecordBatchLoader loader = new RecordBatchLoader(getAllocator());
    List<QueryDataBatch> actual = Collections.emptyList();
    List<QueryDataBatch> expected = Collections.emptyList();
    List<Map<String, Object>> expectedRecords = new ArrayList<>();
    List<Map<String, Object>> actualRecords = new ArrayList<>();
    try {
        test(testOptionSettingQueries);
        actual = testRunAndReturn(queryType, query);
        checkNumBatches(actual);
        addTypeInfoIfMissing(actual.get(0), testBuilder);
        addToMaterializedResults(actualRecords, actual, loader);
        // the cases where the baseline is stored in a file.
        if (baselineRecords == null) {
            test(baselineOptionSettingQueries);
            expected = testRunAndReturn(baselineQueryType, testBuilder.getValidationQuery());
            addToMaterializedResults(expectedRecords, expected, loader);
        } else {
            expectedRecords = baselineRecords;
        }
        compareResults(expectedRecords, actualRecords);
    } finally {
        cleanupBatches(actual, expected);
    }
}

Example 18 with QueryDataBatch

use of org.apache.drill.exec.rpc.user.QueryDataBatch in project drill by axbaretto.

the class PlanTestBase method getPlanInString.

/*
   * This will submit an "EXPLAIN" statement, and return the column value which
   * contains the plan's string.
   */
protected static String getPlanInString(String sql, String columnName) throws Exception {
    final List<QueryDataBatch> results = testSqlWithResults(sql);
    final RecordBatchLoader loader = new RecordBatchLoader(getDrillbitContext().getAllocator());
    final StringBuilder builder = new StringBuilder();
    final boolean silent = config != null && config.getBoolean(QueryTestUtil.TEST_QUERY_PRINTING_SILENT);
    for (final QueryDataBatch b : results) {
        if (!b.hasData()) {
            continue;
        }
        loader.load(b.getHeader().getDef(), b.getData());
        final VectorWrapper<?> vw;
        try {
            vw = loader.getValueAccessorById(NullableVarCharVector.class, loader.getValueVectorId(SchemaPath.getSimplePath(columnName)).getFieldIds());
        } catch (Throwable t) {
            throw new Exception("Looks like you did not provide an explain plan query, please add EXPLAIN PLAN FOR to the beginning of your query.");
        }
        if (!silent) {
            System.out.println(vw.getValueVector().getField().getName());
        }
        final ValueVector vv = vw.getValueVector();
        for (int i = 0; i < vv.getAccessor().getValueCount(); i++) {
            final Object o = vv.getAccessor().getObject(i);
            builder.append(o);
            if (!silent) {
                System.out.println(o);
            }
        }
        loader.clear();
        b.release();
    }
    return builder.toString();
}

Example 19 with QueryDataBatch

use of org.apache.drill.exec.rpc.user.QueryDataBatch in project drill by axbaretto.

the class QueryTestUtil method createClient.

/**
 * Create a DrillClient that can be used to query a drill cluster.
 *
 * @param drillConfig
 * @param remoteServiceSet remote service set
 * @param maxWidth maximum width per node
 * @param props Connection properties contains properties such as "user", "password", "schema" etc
 * @return the newly created client
 * @throws RpcException if there is a problem setting up the client
 */
public static DrillClient createClient(final DrillConfig drillConfig, final RemoteServiceSet remoteServiceSet, final int maxWidth, final Properties props) throws RpcException, OutOfMemoryException {
    final DrillClient drillClient = new DrillClient(drillConfig, remoteServiceSet.getCoordinator());
    drillClient.connect(props);
    final List<QueryDataBatch> results = drillClient.runQuery(QueryType.SQL, String.format("alter session set `%s` = %d", ExecConstants.MAX_WIDTH_PER_NODE_KEY, maxWidth));
    for (QueryDataBatch queryDataBatch : results) {
        queryDataBatch.release();
    }
    return drillClient;
}

Example 20 with QueryDataBatch

use of org.apache.drill.exec.rpc.user.QueryDataBatch in project drill by axbaretto.

the class TestQueriesOnLargeFile method testRead.

@Test
public void testRead() throws Exception {
    List<QueryDataBatch> results = testSqlWithResults(String.format("SELECT count(*) FROM dfs.`default`.`%s`", dataFile.getName()));
    RecordBatchLoader batchLoader = new RecordBatchLoader(getAllocator());
    for (QueryDataBatch batch : results) {
        batchLoader.load(batch.getHeader().getDef(), batch.getData());
        if (batchLoader.getRecordCount() <= 0) {
            continue;
        }
        BigIntVector countV = (BigIntVector) batchLoader.getValueAccessorById(BigIntVector.class, 0).getValueVector();
        assertTrue("Total of " + NUM_RECORDS + " records expected in count", countV.getAccessor().get(0) == NUM_RECORDS);
        batchLoader.clear();
        batch.release();
    }
}