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();
}
}
}
@Override
public void dataArrived(QueryDataBatch result, ConnectionThrottle throttle) {
try {
final int rows = result.getHeader().getRowCount();
if (result.hasData()) {
RecordBatchLoader loader = null;
try {
loader = new RecordBatchLoader(allocator);
loader.load(result.getHeader().getDef(), result.getData());
// TODO: Clean: DRILL-2933: That load(...) no longer throws
// SchemaChangeException, so check/clean catch clause below.
for (int i = 0; i < loader.getSchema().getFieldCount(); ++i) {
columns.add(loader.getSchema().getColumn(i).getPath());
}
for (int i = 0; i < rows; ++i) {
final Map<String, String> record = Maps.newHashMap();
for (VectorWrapper<?> vw : loader) {
final String field = vw.getValueVector().getMetadata().getNamePart().getName();
final ValueVector.Accessor accessor = vw.getValueVector().getAccessor();
final Object value = i < accessor.getValueCount() ? accessor.getObject(i) : null;
final String display = value == null ? null : value.toString();
record.put(field, display);
}
results.add(record);
}
} finally {
if (loader != null) {
loader.clear();
}
}
}
} catch (SchemaChangeException e) {
throw new RuntimeException(e);
} finally {
result.release();
}
}