Examples with DrillbitContext org.apache.drill.exec.server.DrillbitContext used on opensource projects

Example 66 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class TestTraceOutputDump method testFilter.

@Test
public void testFilter() throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final PhysicalPlanReader reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    final PhysicalPlan plan = reader.readPhysicalPlan(Files.asCharSource(DrillFileUtils.getResourceAsFile("/trace/simple_trace.json"), Charsets.UTF_8).read());
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    while (exec.next()) {
    }
    exec.close();
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
    final FragmentHandle handle = context.getHandle();
    /* Form the file name to which the trace output will dump the record batches */
    final String qid = QueryIdHelper.getQueryId(handle.getQueryId());
    final int majorFragmentId = handle.getMajorFragmentId();
    final int minorFragmentId = handle.getMinorFragmentId();
    final String logLocation = c.getString(ExecConstants.TRACE_DUMP_DIRECTORY);
    final String filename = String.format("%s//%s_%d_%d_mock-scan", logLocation, qid, majorFragmentId, minorFragmentId);
    final Configuration conf = new Configuration();
    conf.set(FileSystem.FS_DEFAULT_NAME_KEY, c.getString(ExecConstants.TRACE_DUMP_FILESYSTEM));
    final FileSystem fs = FileSystem.get(conf);
    final Path path = new Path(filename);
    assertTrue("Trace file does not exist", fs.exists(path));
    final FSDataInputStream in = fs.open(path);
    final VectorAccessibleSerializable wrap = new VectorAccessibleSerializable(context.getAllocator());
    wrap.readFromStream(in);
    final VectorAccessible container = wrap.get();
    /* Assert there are no selection vectors */
    assertNull(wrap.getSv2());
    /* Assert there is only one record */
    assertEquals(1, container.getRecordCount());
    /* Read the Integer value and ASSERT its Integer.MIN_VALUE */
    final int value = (int) container.iterator().next().getValueVector().getAccessor().getObject(0);
    assertEquals(value, Integer.MIN_VALUE);
}

Example 67 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class QueryTestUtil method setupScalarReplacementOption.

/**
 * Set up the options to test the scalar replacement retry option (see
 * ClassTransformer.java). Scalar replacement rewrites bytecode to replace
 * value holders (essentially boxed values) with their member variables as
 * locals. There is still one pattern that doesn't work, and occasionally new
 * ones are introduced. This can be used in tests that exercise failing patterns.
 *
 * <p>This also flushes the compiled code cache.
 *
 * @param drillbit the drillbit
 * @param srOption the scalar replacement option value to use
 * @return the original scalar replacement option setting (so it can be restored)
 */
public static OptionValue setupScalarReplacementOption(final Drillbit drillbit, final ClassTransformer.ScalarReplacementOption srOption) {
    // set the system option
    final DrillbitContext drillbitContext = drillbit.getContext();
    final SystemOptionManager optionManager = drillbitContext.getOptionManager();
    final OptionValue originalOptionValue = optionManager.getOption(ExecConstants.SCALAR_REPLACEMENT_OPTION);
    optionManager.setLocalOption(ExecConstants.SCALAR_REPLACEMENT_OPTION, srOption.name().toLowerCase());
    // flush the code cache
    drillbitContext.getCompiler().flushCache();
    return originalOptionValue;
}

Example 68 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class QueryTestUtil method restoreScalarReplacementOption.

/**
 * Restore the original scalar replacement option returned from
 * setupScalarReplacementOption().
 *
 * <p>This also flushes the compiled code cache.
 *
 * @param drillbit the drillbit
 * @param srOption the scalar replacement option value to use
 */
public static void restoreScalarReplacementOption(final Drillbit drillbit, final String srOption) {
    final DrillbitContext drillbitContext = drillbit.getContext();
    final OptionManager optionManager = drillbitContext.getOptionManager();
    optionManager.setLocalOption(ExecConstants.SCALAR_REPLACEMENT_OPTION, srOption);
    // flush the code cache
    drillbitContext.getCompiler().flushCache();
}

Example 69 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class TestOptiqPlans method doLogicalTest.

private SimpleRootExec doLogicalTest(final BootStrapContext context, UserClientConnection connection, String file, ClusterCoordinator coord, DataConnectionCreator com, Controller controller, WorkEventBus workBus) throws Exception {
    new NonStrictExpectations() {

        {
            context.getMetrics();
            result = new MetricRegistry();
            context.getAllocator();
            result = RootAllocatorFactory.newRoot(config);
            context.getConfig();
            result = config;
        }
    };
    final RemoteServiceSet lss = RemoteServiceSet.getLocalServiceSet();
    final DrillbitContext bitContext = new DrillbitContext(DrillbitEndpoint.getDefaultInstance(), context, coord, controller, com, workBus, new LocalPersistentStoreProvider(config));
    final QueryContext qc = new QueryContext(UserSession.Builder.newBuilder().setSupportComplexTypes(true).build(), bitContext, QueryId.getDefaultInstance());
    final PhysicalPlanReader reader = bitContext.getPlanReader();
    final LogicalPlan plan = reader.readLogicalPlan(Files.toString(FileUtils.getResourceAsFile(file), Charsets.UTF_8));
    final PhysicalPlan pp = new BasicOptimizer(qc, connection).optimize(new BasicOptimizer.BasicOptimizationContext(qc), plan);
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(config);
    final FragmentContext fctxt = new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(fctxt, (FragmentRoot) pp.getSortedOperators(false).iterator().next()));
    return exec;
}

Example 70 with DrillbitContext

use of org.apache.drill.exec.server.DrillbitContext in project drill by apache.

the class ControlMessageHandler method startNewRemoteFragment.

private void startNewRemoteFragment(final PlanFragment fragment) throws UserRpcException {
    logger.debug("Received remote fragment start instruction", fragment);
    final DrillbitContext drillbitContext = bee.getContext();
    try {
        // we either need to start the fragment if it is a leaf fragment, or set up a fragment manager if it is non leaf.
        if (fragment.getLeafFragment()) {
            final FragmentContext context = new FragmentContext(drillbitContext, fragment, drillbitContext.getFunctionImplementationRegistry());
            final ControlTunnel tunnel = drillbitContext.getController().getTunnel(fragment.getForeman());
            final FragmentStatusReporter statusReporter = new FragmentStatusReporter(context, tunnel);
            final FragmentExecutor fr = new FragmentExecutor(context, fragment, statusReporter);
            bee.addFragmentRunner(fr);
        } else {
            // isIntermediate, store for incoming data.
            final NonRootFragmentManager manager = new NonRootFragmentManager(fragment, drillbitContext);
            drillbitContext.getWorkBus().addFragmentManager(manager);
        }
    } catch (final Exception e) {
        throw new UserRpcException(drillbitContext.getEndpoint(), "Failure while trying to start remote fragment", e);
    } catch (final OutOfMemoryError t) {
        if (t.getMessage().startsWith("Direct buffer")) {
            throw new UserRpcException(drillbitContext.getEndpoint(), "Out of direct memory while trying to start remote fragment", t);
        } else {
            throw t;
        }
    }
}