@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 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());
}
@Override
public void run() {
// if a cancel thread has already entered this executor, we have not reason to continue.
if (!hasCloseoutThread.compareAndSet(false, true)) {
return;
}
final Thread myThread = Thread.currentThread();
myThreadRef.set(myThread);
final String originalThreadName = myThread.getName();
final FragmentHandle fragmentHandle = fragmentContext.getHandle();
final DrillbitContext drillbitContext = fragmentContext.getDrillbitContext();
final ClusterCoordinator clusterCoordinator = drillbitContext.getClusterCoordinator();
final DrillbitStatusListener drillbitStatusListener = new FragmentDrillbitStatusListener();
final String newThreadName = QueryIdHelper.getExecutorThreadName(fragmentHandle);
try {
myThread.setName(newThreadName);
// if we didn't get the root operator when the executor was created, create it now.
final FragmentRoot rootOperator =
this.rootOperator != null
? this.rootOperator
: drillbitContext.getPlanReader().readFragmentOperator(fragment.getFragmentJson());
root = ImplCreator.getExec(fragmentContext, rootOperator);
if (root == null) {
return;
}
clusterCoordinator.addDrillbitStatusListener(drillbitStatusListener);
updateState(FragmentState.RUNNING);
acceptExternalEvents.countDown();
injector.injectPause(fragmentContext.getExecutionControls(), "fragment-running", logger);
final DrillbitEndpoint endpoint = drillbitContext.getEndpoint();
logger.debug(
"Starting fragment {}:{} on {}:{}",
fragmentHandle.getMajorFragmentId(),
fragmentHandle.getMinorFragmentId(),
endpoint.getAddress(),
endpoint.getUserPort());
final UserGroupInformation queryUserUgi =
fragmentContext.isImpersonationEnabled()
? ImpersonationUtil.createProxyUgi(fragmentContext.getQueryUserName())
: ImpersonationUtil.getProcessUserUGI();
queryUserUgi.doAs(
new PrivilegedExceptionAction<Void>() {
public Void run() throws Exception {
injector.injectChecked(
fragmentContext.getExecutionControls(), "fragment-execution", IOException.class);
/*
* Run the query until root.next returns false OR we no longer need to continue.
*/
while (shouldContinue() && root.next()) {
// loop
}
return null;
}
});
} catch (OutOfMemoryError | OutOfMemoryException e) {
if (!(e instanceof OutOfMemoryError) || "Direct buffer memory".equals(e.getMessage())) {
fail(UserException.memoryError(e).build(logger));
} else {
// we have a heap out of memory error. The JVM in unstable, exit.
CatastrophicFailure.exit(
e, "Unable to handle out of memory condition in FragmentExecutor.", -2);
}
} catch (AssertionError | Exception e) {
fail(e);
} finally {
// no longer allow this thread to be interrupted. We synchronize here to make sure that cancel
// can't set an
// interruption after we have moved beyond this block.
synchronized (myThreadRef) {
myThreadRef.set(null);
Thread.interrupted();
}
// We need to sure we countDown at least once. We'll do it here to guarantee that.
acceptExternalEvents.countDown();
// here we could be in FAILED, RUNNING, or CANCELLATION_REQUESTED
cleanup(FragmentState.FINISHED);
clusterCoordinator.removeDrillbitStatusListener(drillbitStatusListener);
myThread.setName(originalThreadName);
}
}