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

Example 91 with QueryDataBatch

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

the class TestDrillbitResilience method assertDrillbitsOk.

/**
 * Check that all the drillbits are ok.
 * <p/>
 * <p>The current implementation does this by counting the number of drillbits using a query.
 */
private static void assertDrillbitsOk() {
    SingleRowListener listener = new SingleRowListener() {

        private final BufferAllocator bufferAllocator = RootAllocatorFactory.newRoot(cluster.config());

        private final RecordBatchLoader loader = new RecordBatchLoader(bufferAllocator);

        @Override
        public void rowArrived(QueryDataBatch queryResultBatch) {
            // load the single record
            final QueryData queryData = queryResultBatch.getHeader();
            loader.load(queryData.getDef(), queryResultBatch.getData());
            assertEquals(1, loader.getRecordCount());
            // there should only be one column
            final BatchSchema batchSchema = loader.getSchema();
            assertEquals(1, batchSchema.getFieldCount());
            // the column should be an integer
            final MaterializedField countField = batchSchema.getColumn(0);
            final MinorType fieldType = countField.getType().getMinorType();
            assertEquals(MinorType.BIGINT, fieldType);
            // get the column value
            final VectorWrapper<?> vw = loader.iterator().next();
            final Object obj = vw.getValueVector().getAccessor().getObject(0);
            assertTrue(obj instanceof Long);
            final Long countValue = (Long) obj;
            // assume this means all the drillbits are still ok
            assertEquals(cluster.drillbits().size(), countValue.intValue());
            loader.clear();
        }

        @Override
        public void cleanup() {
            loader.clear();
            DrillAutoCloseables.closeNoChecked(bufferAllocator);
        }
    };
    try {
        QueryTestUtil.testWithListener(client.client(), QueryType.SQL, "select count(*) from sys.memory", listener);
        listener.waitForCompletion();
        QueryState state = listener.getQueryState();
        assertSame(state, QueryState.COMPLETED, () -> String.format("QueryState should be COMPLETED (and not %s).", state));
        assertTrue(listener.getErrorList().isEmpty(), "There should not be any errors when checking if Drillbits are OK");
    } catch (final Exception e) {
        throw new RuntimeException("Couldn't query active drillbits", e);
    } finally {
        logger.debug("Cleanup listener");
        listener.cleanup();
    }
    logger.debug("Drillbits are ok.");
}

Example 92 with QueryDataBatch

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

the class TestBroadcastExchange method TestMultipleSendLocationBroadcastExchange.

@Test
public void TestMultipleSendLocationBroadcastExchange() throws Exception {
    RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
    try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet);
        Drillbit bit2 = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) {
        bit1.run();
        bit2.run();
        client.connect();
        String physicalPlan = Files.asCharSource(DrillFileUtils.getResourceAsFile("/sender/broadcast_exchange_long_run.json"), Charsets.UTF_8).read();
        List<QueryDataBatch> results = client.runQuery(QueryType.PHYSICAL, physicalPlan);
        int count = 0;
        for (QueryDataBatch b : results) {
            if (b.getHeader().getRowCount() != 0) {
                count += b.getHeader().getRowCount();
            }
            b.release();
        }
    // Nothing done with count?
    }
}

Example 93 with QueryDataBatch

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

the class TestCastVarCharToBigInt method testCastToBigInt.

// private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(TestCastVarCharToBigInt.class);
@Test
public void testCastToBigInt() throws Exception {
    try (RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
        Drillbit bit = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) {
        // run query.
        bit.run();
        client.connect();
        List<QueryDataBatch> results = client.runQuery(org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.asCharSource(DrillFileUtils.getResourceAsFile("/functions/cast/test_cast_varchar_to_bigint.json"), Charsets.UTF_8).read().replace("#{TEST_FILE}", "/scan_json_test_cast.json"));
        RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());
        QueryDataBatch batch = results.get(0);
        assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
        for (VectorWrapper<?> v : batchLoader) {
            ValueVector.Accessor accessor = v.getValueVector().getAccessor();
            assertEquals(accessor.getObject(0), 2008L);
            assertEquals(accessor.getObject(1), 2007L);
            assertEquals(accessor.getObject(2), 2006L);
        }
        for (QueryDataBatch b : results) {
            b.release();
        }
        batchLoader.clear();
    }
}

Example 94 with QueryDataBatch

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

the class TestDistributedFragmentRun method oneBitOneExchangeOneEntryRun.

@Test
public void oneBitOneExchangeOneEntryRun() throws Exception {
    RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
    try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) {
        bit1.run();
        client.connect();
        List<QueryDataBatch> results = client.runQuery(QueryType.PHYSICAL, Files.asCharSource(DrillFileUtils.getResourceAsFile("/physical_single_exchange.json"), Charsets.UTF_8).read());
        int count = 0;
        for (QueryDataBatch b : results) {
            count += b.getHeader().getRowCount();
            b.release();
        }
        assertEquals(100, count);
    }
}

Example 95 with QueryDataBatch

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

the class TestDistributedFragmentRun method oneBitOneExchangeTwoEntryRunLogical.

@Test
public void oneBitOneExchangeTwoEntryRunLogical() throws Exception {
    RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
    try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) {
        bit1.run();
        client.connect();
        List<QueryDataBatch> results = client.runQuery(QueryType.LOGICAL, Files.asCharSource(DrillFileUtils.getResourceAsFile("/scan_screen_logical.json"), Charsets.UTF_8).read());
        int count = 0;
        for (QueryDataBatch b : results) {
            count += b.getHeader().getRowCount();
            b.release();
        }
        assertEquals(100, count);
    }
}