use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.
the class TestAggregateFunction method runTest.
public void runTest(Object[] values, String planPath, String dataPath) throws Throwable {
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(QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile(planPath), Charsets.UTF_8).replace("#{TEST_FILE}", dataPath));
RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());
QueryDataBatch batch = results.get(1);
assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
int i = 0;
for (VectorWrapper<?> v : batchLoader) {
ValueVector.Accessor accessor = v.getValueVector().getAccessor();
assertEquals(values[i++], (accessor.getObject(0)));
}
batchLoader.clear();
for (QueryDataBatch b : results) {
b.release();
}
}
}
use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.
the class TestDateFunctions method testCommon.
public void testCommon(String[] expectedResults, String physicalPlan, String resourceFile) 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(physicalPlan), Charsets.UTF_8).replace("#{TEST_FILE}", resourceFile));
RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());
QueryDataBatch batch = results.get(0);
assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
int i = 0;
for (VectorWrapper<?> v : batchLoader) {
ValueVector.Accessor accessor = v.getValueVector().getAccessor();
System.out.println(accessor.getObject(0));
assertEquals(expectedResults[i++], accessor.getObject(0).toString());
}
batchLoader.clear();
for (QueryDataBatch b : results) {
b.release();
}
}
}
use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.
the class TextRecordReaderTest method testFullExecution.
@Test
public void testFullExecution() 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(org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/store/text/test.json"), Charsets.UTF_8).replace("#{DATA_FILE}", FileUtils.getResourceAsFile("/store/text/data/regions.csv").toURI().toString()));
int count = 0;
RecordBatchLoader loader = new RecordBatchLoader(bit1.getContext().getAllocator());
for (QueryDataBatch b : results) {
if (b.getHeader().getRowCount() != 0) {
count += b.getHeader().getRowCount();
}
loader.load(b.getHeader().getDef(), b.getData());
VectorUtil.showVectorAccessibleContent(loader);
loader.clear();
b.release();
}
assertEquals(5, count);
}
}
use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.
the class TestJsonReader method readComplexWithStar.
@Test
public void readComplexWithStar() throws Exception {
List<QueryDataBatch> results = testSqlWithResults("select * from cp.`/store/json/test_complex_read_with_star.json`");
assertEquals(1, results.size());
RecordBatchLoader batchLoader = new RecordBatchLoader(getAllocator());
QueryDataBatch batch = results.get(0);
assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
assertEquals(3, batchLoader.getSchema().getFieldCount());
testExistentColumns(batchLoader);
batch.release();
batchLoader.clear();
}
use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.
the class TestComplexToJson method test.
@Test
public void test() throws Exception {
DrillClient parent_client = client;
List<QueryDataBatch> results;
RecordBatchLoader loader = new RecordBatchLoader(getAllocator());
client = new DrillClient(config, serviceSet.getCoordinator());
client.setSupportComplexTypes(false);
client.connect();
results = testSqlWithResults("select * from dfs_test.`[WORKING_PATH]/src/test/resources/store/text/data/regions.csv`");
loader.load(results.get(0).getHeader().getDef(), results.get(0).getData());
RecordBatchDef def = results.get(0).getHeader().getDef();
// the entire row is returned as a single column
assertEquals(1, def.getFieldCount());
// with setSupportComplexTypes == false, the column mode should be REQUIRED
assertTrue(def.getField(0).getMajorType().getMode() == DataMode.REQUIRED);
loader.clear();
for (QueryDataBatch result : results) {
result.release();
}
client.close();
client = new DrillClient(config, serviceSet.getCoordinator());
client.setSupportComplexTypes(true);
client.connect();
results = testSqlWithResults("select * from dfs_test.`[WORKING_PATH]/src/test/resources/store/text/data/regions.csv`");
loader.load(results.get(0).getHeader().getDef(), results.get(0).getData());
def = results.get(0).getHeader().getDef();
// the entire row is returned as a single column
assertEquals(1, def.getFieldCount());
// with setSupportComplexTypes == true, the column mode should be REPEATED
assertTrue(def.getField(0).getMajorType().getMode() == DataMode.REPEATED);
loader.clear();
for (QueryDataBatch result : results) {
result.release();
}
client.close();
client = parent_client;
}
Aggregations