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

Example 41 with RecordBatchLoader

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

the class TestSimpleExternalSort method validateResults.

private void validateResults(BufferAllocator allocator, List<QueryDataBatch> results) throws SchemaChangeException {
    long previousBigInt = Long.MAX_VALUE;
    int recordCount = 0;
    int batchCount = 0;
    for (QueryDataBatch b : results) {
        RecordBatchLoader loader = new RecordBatchLoader(allocator);
        if (b.getHeader().getRowCount() > 0) {
            batchCount++;
            loader.load(b.getHeader().getDef(), b.getData());
            @SuppressWarnings("resource") BigIntVector c1 = (BigIntVector) loader.getValueAccessorById(BigIntVector.class, loader.getValueVectorId(new SchemaPath("blue", ExpressionPosition.UNKNOWN)).getFieldIds()).getValueVector();
            BigIntVector.Accessor a1 = c1.getAccessor();
            for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
                recordCount++;
                assertTrue(String.format("%d > %d", previousBigInt, a1.get(i)), previousBigInt >= a1.get(i));
                previousBigInt = a1.get(i);
            }
        }
        loader.clear();
        b.release();
    }
    System.out.println(String.format("Sorted %,d records in %d batches.", recordCount, batchCount));
}

Example 42 with RecordBatchLoader

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

the class TestWriter method simpleCsv.

@Test
public void simpleCsv() throws Exception {
    // before executing the test deleting the existing CSV files in /tmp/csvtest
    Path path = new Path("/tmp/csvtest");
    if (fs.exists(path)) {
        fs.delete(path, true);
    }
    String plan = Files.toString(FileUtils.getResourceAsFile("/writer/simple_csv_writer.json"), Charsets.UTF_8);
    List<QueryDataBatch> results = testPhysicalWithResults(plan);
    RecordBatchLoader batchLoader = new RecordBatchLoader(getAllocator());
    QueryDataBatch batch = results.get(0);
    assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
    VarCharVector fragmentIdV = (VarCharVector) batchLoader.getValueAccessorById(VarCharVector.class, 0).getValueVector();
    BigIntVector recordWrittenV = (BigIntVector) batchLoader.getValueAccessorById(BigIntVector.class, 1).getValueVector();
    // expected only one row in output
    assertEquals(1, batchLoader.getRecordCount());
    assertEquals("0_0", fragmentIdV.getAccessor().getObject(0).toString());
    assertEquals(132000, recordWrittenV.getAccessor().get(0));
    // now verify csv files are written to disk
    assertTrue(fs.exists(path));
    // expect two files
    FileStatus[] fileStatuses = fs.globStatus(new Path(path.toString(), "*.csv"));
    assertTrue(2 == fileStatuses.length);
    for (QueryDataBatch b : results) {
        b.release();
    }
    batchLoader.clear();
}

Example 43 with RecordBatchLoader

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

the class TestWriter method testCTASQueryHelper.

private void testCTASQueryHelper(String tableName, String testQuery, int expectedOutputCount) throws Exception {
    try {
        List<QueryDataBatch> results = testSqlWithResults(testQuery);
        RecordBatchLoader batchLoader = new RecordBatchLoader(getAllocator());
        int recordsWritten = 0;
        for (QueryDataBatch batch : results) {
            batchLoader.load(batch.getHeader().getDef(), batch.getData());
            if (batchLoader.getRecordCount() <= 0) {
                continue;
            }
            BigIntVector recordWrittenV = (BigIntVector) batchLoader.getValueAccessorById(BigIntVector.class, 1).getValueVector();
            for (int i = 0; i < batchLoader.getRecordCount(); i++) {
                recordsWritten += recordWrittenV.getAccessor().get(i);
            }
            batchLoader.clear();
            batch.release();
        }
        assertEquals(expectedOutputCount, recordsWritten);
    } finally {
        try {
            Path path = new Path(getDfsTestTmpSchemaLocation(), tableName);
            if (fs.exists(path)) {
                fs.delete(path, true);
            }
        } catch (Exception e) {
            // ignore exceptions.
            logger.warn("Failed to delete the table [{}, {}] created as part of the test", getDfsTestTmpSchemaLocation(), tableName);
        }
    }
}

Example 44 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader 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);
    }
}

Example 45 with RecordBatchLoader

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

the class TestDateTypes method testDate.

//    private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(TestDateTypes.class);
@Test
public void testDate() 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.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()), ("1970-01-02"));
            assertEquals((accessor.getObject(1).toString()), ("2008-12-28"));
            assertEquals((accessor.getObject(2).toString()), ("2000-02-27"));
        }
        batchLoader.clear();
        for (QueryDataBatch b : results) {
            b.release();
        }
    }
}