@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);
}
}
@Test // DRILL-3580
public void testExpressionInWindowFunction() throws Exception {
String root = FileUtils.getResourceAsFile("/store/text/data/t.json").toURI().toString();
String query =
String.format(
"select a1, b1, sum(b1) over (partition by a1) as c1, sum(a1 + b1) over (partition by a1) as c2\n"
+ "from dfs_test.`%s`",
root);
// Validate the plan
final String[] expectedPlan = {
"Window\\(window#0=\\[window\\(partition \\{0\\} order by \\[\\].*\\[SUM\\(\\$1\\), SUM\\(\\$2\\)\\]"
};
PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, new String[] {});
testBuilder()
.sqlQuery(query)
.unOrdered()
.baselineColumns("a1", "b1", "c1", "c2")
.baselineValues(0l, 1l, 8l, 8l)
.baselineValues(0l, 1l, 8l, 8l)
.baselineValues(0l, 2l, 8l, 8l)
.baselineValues(0l, 2l, 8l, 8l)
.baselineValues(0l, 2l, 8l, 8l)
.baselineValues(10l, 3l, 21l, 71l)
.baselineValues(10l, 3l, 21l, 71l)
.baselineValues(10l, 5l, 21l, 71l)
.baselineValues(10l, 5l, 21l, 71l)
.baselineValues(10l, 5l, 21l, 71l)
.build()
.run();
}
// specific tests should call this method, but it is not marked as a test itself intentionally
public void testParquetFullEngineEventBased(
boolean testValues,
boolean generateNew,
String plan,
String readEntries,
String filename,
int numberOfTimesRead /* specified in json plan */,
ParquetTestProperties props,
QueryType queryType)
throws Exception {
if (generateNew) TestFileGenerator.generateParquetFile(filename, props);
ParquetResultListener resultListener =
new ParquetResultListener(getAllocator(), props, numberOfTimesRead, testValues);
long C = System.nanoTime();
String planText = Files.toString(FileUtils.getResourceAsFile(plan), Charsets.UTF_8);
// substitute in the string for the read entries, allows reuse of the plan file for several
// tests
if (readEntries != null) {
planText = planText.replaceFirst("&REPLACED_IN_PARQUET_TEST&", readEntries);
}
this.testWithListener(queryType, planText, resultListener);
resultListener.getResults();
long D = System.nanoTime();
System.out.println(String.format("Took %f s to run query", (float) (D - C) / 1E9));
}
@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 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);
}
}
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 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 // see DRILL-3574
public void testWithAndWithoutPartitions() throws Exception {
String root = FileUtils.getResourceAsFile("/store/text/data/t.json").toURI().toString();
String query =
String.format(
"select sum(a1) over(partition by b1, c1) as s1, sum(a1) over() as s2 \n"
+ "from dfs_test.`%s` \n"
+ "order by a1",
root);
test("alter session set `planner.slice_target` = 1");
// Validate the plan
final String[] expectedPlan = {
"Window\\(window#0=\\[window\\(partition \\{\\}.*\n" + ".*UnionExchange"
};
PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, new String[] {});
testBuilder()
.sqlQuery(query)
.unOrdered()
.baselineColumns("s1", "s2")
.baselineValues(0l, 50l)
.baselineValues(0l, 50l)
.baselineValues(0l, 50l)
.baselineValues(0l, 50l)
.baselineValues(0l, 50l)
.baselineValues(10l, 50l)
.baselineValues(10l, 50l)
.baselineValues(10l, 50l)
.baselineValues(20l, 50l)
.baselineValues(20l, 50l)
.build()
.run();
}
@Test
public void testSubstringNegative(
@Injectable final DrillbitContext bitContext,
@Injectable UserServer.UserClientConnection connection)
throws Throwable {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry();
bitContext.getAllocator();
result = new TopLevelAllocator();
bitContext.getOperatorCreatorRegistry();
result = new OperatorCreatorRegistry(c);
bitContext.getConfig();
result = c;
}
};
PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(
FileUtils.getResourceAsFile("/functions/testSubstringNegative.json"),
Charsets.UTF_8));
FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
FragmentContext context =
new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
while (exec.next()) {
NullableVarCharVector c1 =
exec.getValueVectorById(
new SchemaPath("col3", ExpressionPosition.UNKNOWN), NullableVarCharVector.class);
NullableVarCharVector.Accessor a1;
a1 = c1.getAccessor();
int count = 0;
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
if (!a1.isNull(i)) {
NullableVarCharHolder holder = new NullableVarCharHolder();
a1.get(i, holder);
// when offset is negative, substring return empty string.
assertEquals("", holder.toString());
++count;
}
}
assertEquals(50, count);
}
if (context.getFailureCause() != null) {
throw context.getFailureCause();
}
assertTrue(!context.isFailed());
}
public Fragment getRootFragment(PhysicalPlanReader reader, String file)
throws FragmentSetupException, IOException {
MakeFragmentsVisitor f = new MakeFragmentsVisitor();
PhysicalPlan plan =
reader.readPhysicalPlan(Files.toString(FileUtils.getResourceAsFile(file), Charsets.UTF_8));
PhysicalOperator o = plan.getSortedOperators(false).iterator().next();
return o.accept(f, null);
}
public void testParquetFullEngineLocalPath(
String planFileName,
String filename,
int numberOfTimesRead /* specified in json plan */,
int numberOfRowGroups,
int recordsPerRowGroup)
throws Exception {
testParquetFullEngineLocalText(
Files.toString(FileUtils.getResourceAsFile(planFileName), Charsets.UTF_8),
filename,
numberOfTimesRead,
numberOfRowGroups,
recordsPerRowGroup,
true);
}
@Test
public void project(
@Injectable final DrillbitContext bitContext, @Injectable UserClientConnection connection)
throws Exception {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry("test");
bitContext.getAllocator();
result = BufferAllocator.getAllocator(c);
}
};
PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(FileUtils.getResourceAsFile("/project/test1.json"), Charsets.UTF_8));
FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
FragmentContext context =
new FragmentContext(
bitContext, FragmentHandle.getDefaultInstance(), connection, null, registry);
SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
while (exec.next()) {
BigIntVector c1 =
exec.getValueVectorById(
new SchemaPath("col1", ExpressionPosition.UNKNOWN), BigIntVector.class);
BigIntVector c2 =
exec.getValueVectorById(
new SchemaPath("col2", ExpressionPosition.UNKNOWN), BigIntVector.class);
int x = 0;
BigIntVector.Accessor a1, a2;
a1 = c1.getAccessor();
a2 = c2.getAccessor();
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
assertEquals(a1.get(i) + 1, a2.get(i));
x += a1.get(i);
}
System.out.println(x);
}
}
@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);
}
}
// use this method to submit physical plan
public void testParquetFullEngineLocalTextDistributed(
String planName,
String filename,
int numberOfTimesRead /* specified in json plan */,
int numberOfRowGroups,
int recordsPerRowGroup)
throws Exception {
String planText = Files.toString(FileUtils.getResourceAsFile(planName), Charsets.UTF_8);
testFull(
QueryType.PHYSICAL,
planText,
filename,
numberOfTimesRead,
numberOfRowGroups,
recordsPerRowGroup,
true);
}
@Test
public void testMultipleRowGroupsAndReads2() throws Exception {
String readEntries;
readEntries = "";
// number of times to read the file
int i = 3;
for (int j = 0; j < i; j++) {
readEntries += "\"" + fileName + "\"";
if (j < i - 1) readEntries += ",";
}
String planText =
Files.toString(
FileUtils.getResourceAsFile("/parquet/parquet_scan_screen_read_entry_replace.json"),
Charsets.UTF_8)
.replaceFirst("&REPLACED_IN_PARQUET_TEST&", readEntries);
testParquetFullEngineLocalText(
planText, fileName, i, numberRowGroups, recordsPerRowGroup, true);
}
public void testParquetFullEngineRemote(
String plan,
String filename,
int numberOfTimesRead /* specified in json plan */,
int numberOfRowGroups,
int recordsPerRowGroup)
throws Exception {
HashMap<String, FieldInfo> fields = new HashMap<>();
ParquetTestProperties props =
new ParquetTestProperties(
numberRowGroups, recordsPerRowGroup, DEFAULT_BYTES_PER_PAGE, fields);
TestFileGenerator.populateFieldInfoMap(props);
ParquetResultListener resultListener =
new ParquetResultListener(getAllocator(), props, numberOfTimesRead, true);
testWithListener(
QueryType.PHYSICAL,
Files.toString(FileUtils.getResourceAsFile(plan), Charsets.UTF_8),
resultListener);
resultListener.getResults();
}
@Test
@Ignore
public void testDictionaryError() throws Exception {
String readEntries;
readEntries = "\"/tmp/lineitem_null_dict.parquet\"";
String planText =
Files.toString(
FileUtils.getResourceAsFile("/parquet/parquet_scan_screen_read_entry_replace.json"),
Charsets.UTF_8)
.replaceFirst("&REPLACED_IN_PARQUET_TEST&", readEntries);
// testParquetFullEngineLocalText(planText, fileName, 1, 1, 100000, false);
testFull(
QueryType.SQL,
"select L_RECEIPTDATE from dfs.`/tmp/lineitem_null_dict.parquet`",
"",
1,
1,
100000,
false);
}
@Test
public void testCompoundIdentifierInWindowDefinition() throws Exception {
String root =
FileUtils.getResourceAsFile("/multilevel/csv/1994/Q1/orders_94_q1.csv").toURI().toString();
String query =
String.format(
"SELECT count(*) OVER w as col1, count(*) OVER w as col2 \n"
+ "FROM dfs_test.`%s` \n"
+ "WINDOW w AS (PARTITION BY columns[1] ORDER BY columns[0] DESC)",
root);
// Validate the plan
final String[] expectedPlan = {
"Window.*partition \\{1\\} order by \\[0 DESC\\].*COUNT\\(\\)",
"Scan.*columns=\\[`columns`\\[0\\], `columns`\\[1\\]\\]"
};
final String[] excludedPatterns = {"Scan.*columns=\\[`\\*`\\]"};
PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, excludedPatterns);
// Validate the result
testBuilder()
.sqlQuery(query)
.ordered()
.baselineColumns("col1", "col2")
.baselineValues((long) 1, (long) 1)
.baselineValues((long) 1, (long) 1)
.baselineValues((long) 1, (long) 1)
.baselineValues((long) 1, (long) 1)
.baselineValues((long) 1, (long) 1)
.baselineValues((long) 1, (long) 1)
.baselineValues((long) 1, (long) 1)
.baselineValues((long) 1, (long) 1)
.baselineValues((long) 1, (long) 1)
.baselineValues((long) 1, (long) 1)
.build()
.run();
}
@Test // DRILL-3567
public void testMultiplePartitions2() throws Exception {
String root = FileUtils.getResourceAsFile("/store/text/data/t.json").toURI().toString();
String query =
String.format(
"select count(*) over(partition by b1 order by c1) as count1, \n"
+ "count(*) over(partition by a1 order by c1) as count2, \n"
+ "sum(a1) over(partition by b1 order by c1) as sum1 \n"
+ "from dfs_test.`%s`",
root);
// Validate the plan
final String[] expectedPlan = {
"Window.*partition \\{2\\} order by \\[1\\].*COUNT\\(\\)",
"Window.*partition \\{0\\} order by \\[1\\].*COUNT\\(\\), SUM\\(\\$2\\)",
"Scan.*columns=\\[`b1`, `c1`, `a1`\\]"
};
final String[] excludedPatterns = {"Scan.*columns=\\[`\\*`\\]"};
PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, excludedPatterns);
testBuilder()
.sqlQuery(query)
.unOrdered()
.baselineColumns("count1", "count2", "sum1")
.baselineValues(1l, 2l, 0l)
.baselineValues(1l, 2l, 0l)
.baselineValues(2l, 5l, 0l)
.baselineValues(3l, 5l, 0l)
.baselineValues(3l, 5l, 0l)
.baselineValues(1l, 2l, 10l)
.baselineValues(1l, 2l, 10l)
.baselineValues(2l, 5l, 20l)
.baselineValues(3l, 5l, 30l)
.baselineValues(3l, 5l, 30l)
.build()
.run();
}
@Test // see DRILL-3657
public void testConstantsInMultiplePartitions() throws Exception {
String root = FileUtils.getResourceAsFile("/store/text/data/t.json").toURI().toString();
String query =
String.format(
"select sum(1) over(partition by b1 order by a1) as sum1, sum(1) over(partition by a1) as sum2, rank() over(order by b1) as rank1, rank() over(order by 1) as rank2 \n"
+ "from dfs_test.`%s` \n"
+ "order by 1, 2, 3, 4",
root);
// Validate the plan
final String[] expectedPlan = {
"Window.*SUM\\(\\$3\\).*\n"
+ ".*SelectionVectorRemover.*\n"
+ ".*Sort.*\n"
+ ".*Window.*SUM\\(\\$2\\).*"
};
PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, new String[] {});
testBuilder()
.sqlQuery(query)
.unOrdered()
.baselineColumns("sum1", "sum2", "rank1", "rank2")
.baselineValues(2l, 5l, 1l, 1l)
.baselineValues(2l, 5l, 1l, 1l)
.baselineValues(2l, 5l, 6l, 1l)
.baselineValues(2l, 5l, 6l, 1l)
.baselineValues(3l, 5l, 3l, 1l)
.baselineValues(3l, 5l, 3l, 1l)
.baselineValues(3l, 5l, 3l, 1l)
.baselineValues(3l, 5l, 8l, 1l)
.baselineValues(3l, 5l, 8l, 1l)
.baselineValues(3l, 5l, 8l, 1l)
.build()
.run();
}
public String getPlanForFile(String pathFileName, String parquetFileName) throws IOException {
return Files.toString(FileUtils.getResourceAsFile(pathFileName), Charsets.UTF_8)
.replaceFirst("&REPLACED_IN_PARQUET_TEST&", parquetFileName);
}
@Test
public void testAllocators() throws Exception {
// Setup a drillbit (initializes a root allocator)
final DrillConfig config = DrillConfig.create(TEST_CONFIGURATIONS);
final RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
final Drillbit bit = new Drillbit(config, serviceSet);
bit.run();
final DrillbitContext bitContext = bit.getContext();
FunctionImplementationRegistry functionRegistry =
bitContext.getFunctionImplementationRegistry();
StoragePluginRegistry storageRegistry = new StoragePluginRegistry(bitContext);
// Create a few Fragment Contexts
BitControl.PlanFragment.Builder pfBuilder1 = BitControl.PlanFragment.newBuilder();
pfBuilder1.setMemInitial(1500000);
BitControl.PlanFragment pf1 = pfBuilder1.build();
BitControl.PlanFragment.Builder pfBuilder2 = BitControl.PlanFragment.newBuilder();
pfBuilder2.setMemInitial(500000);
BitControl.PlanFragment pf2 = pfBuilder1.build();
FragmentContext fragmentContext1 = new FragmentContext(bitContext, pf1, null, functionRegistry);
FragmentContext fragmentContext2 = new FragmentContext(bitContext, pf2, null, functionRegistry);
// Get a few physical operators. Easiest way is to read a physical plan.
PhysicalPlanReader planReader =
new PhysicalPlanReader(
config,
config.getMapper(),
CoordinationProtos.DrillbitEndpoint.getDefaultInstance(),
storageRegistry);
PhysicalPlan plan =
planReader.readPhysicalPlan(
Files.toString(FileUtils.getResourceAsFile(planFile), Charsets.UTF_8));
List<PhysicalOperator> physicalOperators = plan.getSortedOperators();
Iterator<PhysicalOperator> physicalOperatorIterator = physicalOperators.iterator();
PhysicalOperator physicalOperator1 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator2 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator3 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator4 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator5 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator6 = physicalOperatorIterator.next();
// Create some bogus Operator profile defs and stats to create operator contexts
OpProfileDef def;
OperatorStats stats;
// Use some bogus operator type to create a new operator context.
def =
new OpProfileDef(
physicalOperator1.getOperatorId(),
UserBitShared.CoreOperatorType.MOCK_SUB_SCAN_VALUE,
OperatorContext.getChildCount(physicalOperator1));
stats = fragmentContext1.getStats().getOperatorStats(def, fragmentContext1.getAllocator());
// Add a couple of Operator Contexts
// Initial allocation = 1000000 bytes for all operators
OperatorContext oContext11 = fragmentContext1.newOperatorContext(physicalOperator1, true);
DrillBuf b11 = oContext11.getAllocator().buffer(1000000);
OperatorContext oContext12 =
fragmentContext1.newOperatorContext(physicalOperator2, stats, true);
DrillBuf b12 = oContext12.getAllocator().buffer(500000);
OperatorContext oContext21 = fragmentContext1.newOperatorContext(physicalOperator3, true);
def =
new OpProfileDef(
physicalOperator4.getOperatorId(),
UserBitShared.CoreOperatorType.TEXT_WRITER_VALUE,
OperatorContext.getChildCount(physicalOperator4));
stats = fragmentContext2.getStats().getOperatorStats(def, fragmentContext2.getAllocator());
OperatorContext oContext22 =
fragmentContext2.newOperatorContext(physicalOperator4, stats, true);
DrillBuf b22 = oContext22.getAllocator().buffer(2000000);
// New Fragment begins
BitControl.PlanFragment.Builder pfBuilder3 = BitControl.PlanFragment.newBuilder();
pfBuilder3.setMemInitial(1000000);
BitControl.PlanFragment pf3 = pfBuilder3.build();
FragmentContext fragmentContext3 = new FragmentContext(bitContext, pf3, null, functionRegistry);
// New fragment starts an operator that allocates an amount within the limit
def =
new OpProfileDef(
physicalOperator5.getOperatorId(),
UserBitShared.CoreOperatorType.UNION_VALUE,
OperatorContext.getChildCount(physicalOperator5));
stats = fragmentContext3.getStats().getOperatorStats(def, fragmentContext3.getAllocator());
OperatorContext oContext31 =
fragmentContext3.newOperatorContext(physicalOperator5, stats, true);
DrillBuf b31a = oContext31.getAllocator().buffer(200000);
// Previously running operator completes
b22.release();
((AutoCloseable) oContext22).close();
// Fragment 3 asks for more and fails
boolean outOfMem = false;
try {
DrillBuf b31b = oContext31.getAllocator().buffer(4400000);
if (b31b != null) {
b31b.release();
} else {
outOfMem = true;
}
} catch (Exception e) {
outOfMem = true;
}
assertEquals(true, (boolean) outOfMem);
// Operator is Exempt from Fragment limits. Fragment 3 asks for more and succeeds
outOfMem = false;
OperatorContext oContext32 = fragmentContext3.newOperatorContext(physicalOperator6, false);
DrillBuf b32 = null;
try {
b32 = oContext32.getAllocator().buffer(4400000);
} catch (Exception e) {
outOfMem = true;
} finally {
if (b32 != null) {
b32.release();
} else {
outOfMem = true;
}
closeOp(oContext32);
}
assertEquals(false, (boolean) outOfMem);
b11.release();
closeOp(oContext11);
b12.release();
closeOp(oContext12);
closeOp(oContext21);
b31a.release();
closeOp(oContext31);
fragmentContext1.close();
fragmentContext2.close();
fragmentContext3.close();
bit.close();
serviceSet.close();
}
@Test
public void sortOneKeyAscending(
@Injectable final DrillbitContext bitContext, @Injectable UserClientConnection connection)
throws Throwable {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry();
bitContext.getAllocator();
result = new TopLevelAllocator();
bitContext.getOperatorCreatorRegistry();
result = new OperatorCreatorRegistry(c);
bitContext.getConfig();
result = c;
bitContext.getCompiler();
result = CodeCompiler.getTestCompiler(c);
}
};
final PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
final PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(FileUtils.getResourceAsFile("/sort/one_key_sort.json"), Charsets.UTF_8));
final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
final FragmentContext context =
new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
final SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
int previousInt = Integer.MIN_VALUE;
int recordCount = 0;
int batchCount = 0;
while (exec.next()) {
batchCount++;
final IntVector c1 =
exec.getValueVectorById(
new SchemaPath("blue", ExpressionPosition.UNKNOWN), IntVector.class);
final IntVector c2 =
exec.getValueVectorById(
new SchemaPath("green", ExpressionPosition.UNKNOWN), IntVector.class);
final IntVector.Accessor a1 = c1.getAccessor();
final IntVector.Accessor a2 = c2.getAccessor();
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
recordCount++;
assertTrue(previousInt <= a1.get(i));
previousInt = a1.get(i);
assertEquals(previousInt, a2.get(i));
}
}
System.out.println(String.format("Sorted %,d records in %d batches.", recordCount, batchCount));
if (context.getFailureCause() != null) {
throw context.getFailureCause();
}
assertTrue(!context.isFailed());
}