@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.toString(
FileUtils.getResourceAsFile("/sender/broadcast_exchange_long_run.json"),
Charsets.UTF_8);
List<QueryResultBatch> results = client.runQuery(QueryType.PHYSICAL, physicalPlan);
int count = 0;
for (QueryResultBatch b : results) {
if (b.getHeader().getRowCount() != 0) count += b.getHeader().getRowCount();
b.release();
}
System.out.println(count);
}
}
@Test
public void twoBitTwoExchangeTwoEntryRun() 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();
List<QueryResultBatch> results =
client.runQuery(
org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
Files.toString(
FileUtils.getResourceAsFile("/sender/union_exchange.json"), Charsets.UTF_8));
int count = 0;
for (QueryResultBatch b : results) {
if (b.getHeader().getRowCount() != 0) {
count += b.getHeader().getRowCount();
}
b.release();
}
assertEquals(150, count);
}
}
@Test
public void TestSingleBroadcastExchangeWithTwoScans() 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.toString(
FileUtils.getResourceAsFile("/sender/broadcast_exchange.json"), Charsets.UTF_8)
.replace(
"#{LEFT_FILE}",
FileUtils.getResourceAsFile("/join/merge_single_batch.left.json")
.toURI()
.toString())
.replace(
"#{RIGHT_FILE}",
FileUtils.getResourceAsFile("/join/merge_single_batch.right.json")
.toURI()
.toString());
List<QueryResultBatch> results = client.runQuery(QueryType.PHYSICAL, physicalPlan);
int count = 0;
for (QueryResultBatch b : results) {
if (b.getHeader().getRowCount() != 0) count += b.getHeader().getRowCount();
b.release();
}
assertEquals(25, count);
}
}
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();
}
}
}
@Test
public void testMultipleProvidersMixedSizes() 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();
List<QueryResultBatch> results =
client.runQuery(
org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
Files.toString(
FileUtils.getResourceAsFile("/mergerecv/multiple_providers.json"),
Charsets.UTF_8));
int count = 0;
RecordBatchLoader batchLoader = new RecordBatchLoader(client.getAllocator());
// print the results
Long lastBlueValue = null;
for (QueryResultBatch b : results) {
count += b.getHeader().getRowCount();
for (int valueIdx = 0; valueIdx < b.getHeader().getRowCount(); valueIdx++) {
List<Object> row = Lists.newArrayList();
batchLoader.load(b.getHeader().getDef(), b.getData());
for (VectorWrapper vw : batchLoader) {
row.add(
vw.getValueVector().getField().toExpr()
+ ":"
+ vw.getValueVector().getAccessor().getObject(valueIdx));
if (vw.getValueVector()
.getField()
.getAsSchemaPath()
.getRootSegment()
.getPath()
.equals("blue")) {
// assert order is ascending
if (((Long) vw.getValueVector().getAccessor().getObject(valueIdx)).longValue() == 0)
continue; // ignore initial 0's from sort
if (lastBlueValue != null)
assertTrue(
((Long) vw.getValueVector().getAccessor().getObject(valueIdx)).longValue()
>= ((Long) lastBlueValue).longValue());
lastBlueValue = (Long) vw.getValueVector().getAccessor().getObject(valueIdx);
}
}
for (Object cell : row) {
int len = cell.toString().length();
System.out.print(cell + " ");
for (int i = 0; i < (30 - len); ++i) System.out.print(" ");
}
System.out.println();
}
b.release();
batchLoader.clear();
}
assertEquals(400, count);
}
}
@Test
public void twoBitTwoExchange() 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();
List<QueryResultBatch> results =
client.runQuery(
org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
Files.toString(
FileUtils.getResourceAsFile("/mergerecv/merging_receiver.json"), Charsets.UTF_8));
int count = 0;
RecordBatchLoader batchLoader = new RecordBatchLoader(client.getAllocator());
// print the results
for (QueryResultBatch b : results) {
count += b.getHeader().getRowCount();
for (int valueIdx = 0; valueIdx < b.getHeader().getRowCount(); valueIdx++) {
List<Object> row = Lists.newArrayList();
batchLoader.load(b.getHeader().getDef(), b.getData());
for (VectorWrapper<?> vw : batchLoader)
row.add(
vw.getValueVector().getField().toExpr()
+ ":"
+ vw.getValueVector().getAccessor().getObject(valueIdx));
for (Object cell : row) {
if (cell == null) {
System.out.print("<null> ");
continue;
}
int len = cell.toString().length();
System.out.print(cell + " ");
for (int i = 0; i < (30 - len); ++i) System.out.print(" ");
}
System.out.println();
}
b.release();
batchLoader.clear();
}
assertEquals(200, count);
}
}
@Test
@Ignore
public void testParseParquetPhysicalPlan() throws Exception {
RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
DrillConfig config = DrillConfig.create();
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,
Resources.toString(Resources.getResource(fileName), Charsets.UTF_8));
RecordBatchLoader loader = new RecordBatchLoader(bit1.getContext().getAllocator());
int count = 0;
for (QueryDataBatch b : results) {
System.out.println(String.format("Got %d results", b.getHeader().getRowCount()));
count += b.getHeader().getRowCount();
loader.load(b.getHeader().getDef(), b.getData());
for (VectorWrapper vw : loader) {
System.out.print(vw.getValueVector().getField().toExpr() + ": ");
ValueVector vv = vw.getValueVector();
for (int i = 0; i < vv.getAccessor().getValueCount(); i++) {
Object o = vv.getAccessor().getObject(i);
if (o instanceof byte[]) {
System.out.print(" [" + new String((byte[]) o) + "]");
} else {
System.out.print(" [" + vv.getAccessor().getObject(i) + "]");
}
// break;
}
System.out.println();
}
loader.clear();
b.release();
}
client.close();
System.out.println(String.format("Got %d total results", count));
}
}
@Test
@Ignore
public void testParseParquetPhysicalPlanRemote() throws Exception {
DrillConfig config = DrillConfig.create();
try (DrillClient client = new DrillClient(config); ) {
client.connect();
ParquetResultsListener listener = new ParquetResultsListener();
Stopwatch watch = new Stopwatch();
watch.start();
client.runQuery(
org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
Resources.toString(Resources.getResource(fileName), Charsets.UTF_8),
listener);
System.out.println(
String.format(
"Got %d total records in %d seconds",
listener.await(), watch.elapsed(TimeUnit.SECONDS)));
client.close();
}
}
public static void setSessionOption(
final DrillClient drillClient, final String option, final String value) {
try {
final List<QueryDataBatch> results =
drillClient.runQuery(
UserBitShared.QueryType.SQL,
String.format("ALTER session SET `%s` = %s", option, value));
for (final QueryDataBatch data : results) {
data.release();
}
} catch (final RpcException e) {
fail("Could not set option: " + e.toString());
}
}
public QueryResult run(final DrillClient client, final BufferAllocator allocator)
throws Exception {
Listener listener = new Listener(allocator);
client.runQuery(getType(), query, listener);
listener.waitForCompletion();
if (listener.results.isEmpty()) {
listener.results.add(Maps.<String, String>newHashMap());
}
final Map<String, String> first = listener.results.get(0);
for (String columnName : listener.columns) {
if (!first.containsKey(columnName)) {
first.put(columnName, null);
}
}
return new QueryResult(listener.columns, listener.results);
}
