@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); } }