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

Example 11 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.

the class TestDecimal method testSimpleDecimalMathFunc.

@Test
public void testSimpleDecimalMathFunc() 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.toString(FileUtils.getResourceAsFile("/decimal/simple_decimal_math.json"), Charsets.UTF_8).replace("#{TEST_FILE}", "/input_simple_decimal.json"));
        RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());
        QueryDataBatch batch = results.get(0);
        assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
        Iterator<VectorWrapper<?>> itr = batchLoader.iterator();
        // Check the output of decimal18
        ValueVector.Accessor dec18Accessor = itr.next().getValueVector().getAccessor();
        assertEquals(6, dec18Accessor.getValueCount());
        batchLoader.clear();
        for (QueryDataBatch result : results) {
            result.release();
        }
    }
}

Example 12 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader 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() {
    final SingleRowListener listener = new SingleRowListener() {

        private final BufferAllocator bufferAllocator = RootAllocatorFactory.newRoot(zkHelper.getConfig());

        private final RecordBatchLoader loader = new RecordBatchLoader(bufferAllocator);

        @Override
        public void rowArrived(final QueryDataBatch queryResultBatch) {
            // load the single record
            final QueryData queryData = queryResultBatch.getHeader();
            try {
                loader.load(queryData.getDef(), queryResultBatch.getData());
            // TODO:  Clean:  DRILL-2933:  That load(...) no longer throws
            // SchemaChangeException, so check/clean catch clause below.
            } catch (final SchemaChangeException e) {
                fail(e.toString());
            }
            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(drillbits.size(), countValue.intValue());
            loader.clear();
        }

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

Example 13 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.

the class TestQueriesOnLargeFile method testRead.

@Test
public void testRead() throws Exception {
    List<QueryDataBatch> results = testSqlWithResults(String.format("SELECT count(*) FROM dfs_test.`default`.`%s`", dataFile.getPath()));
    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();
    }
}

Example 14 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.

the class TestDateTypes method testDateAdd.

@Test
public void testDateAdd() 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.toString(FileUtils.getResourceAsFile("/record/vector/test_date_add.json"), Charsets.UTF_8).replace("#{TEST_FILE}", "/test_simple_date.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).toString()), ("2008-03-27 00:00:00.000"));
        }
        batchLoader.clear();
        for (QueryDataBatch b : results) {
            b.release();
        }
    }
}

Example 15 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.

the class TestDateTypes method testLiterals.

@Test
public void testLiterals() 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.toString(FileUtils.getResourceAsFile("/record/vector/test_all_date_literals.json"), Charsets.UTF_8).replace("#{TEST_FILE}", "/test_simple_date.json"));
        RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());
        QueryDataBatch batch = results.get(0);
        assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
        String[] result = { "2008-02-27", "2008-02-27 01:02:03.000", "10:11:13.999", "2 years 2 months 3 days 0:1:3.89" };
        int idx = 0;
        for (VectorWrapper<?> v : batchLoader) {
            ValueVector.Accessor accessor = v.getValueVector().getAccessor();
            assertEquals((accessor.getObject(0).toString()), (result[idx++]));
        }
        batchLoader.clear();
        for (QueryDataBatch b : results) {
            b.release();
        }
    }
}