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

Example 91 with DrillbitContext

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

the class TestComparisonFunctions method runTest.

public void runTest(String expression, int expectedResults) throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final String planString = Resources.toString(Resources.getResource(COMPARISON_TEST_PHYSICAL_PLAN), Charsets.UTF_8).replaceAll("EXPRESSION", expression);
    if (reader == null) {
        reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    }
    if (registry == null) {
        registry = new FunctionImplementationRegistry(c);
    }
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final PhysicalPlan plan = reader.readPhysicalPlan(planString);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    while (exec.next()) {
        assertEquals(String.format("Expression: %s;", expression), expectedResults, exec.getSelectionVector2().getCount());
    }
    exec.close();
    context.close();
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 92 with DrillbitContext

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

the class TestMathFunctions method testBasicMathFunctions.

@Test
public void testBasicMathFunctions() 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.toString(DrillFileUtils.getResourceAsFile("/functions/simple_math_functions.json"), Charsets.UTF_8));
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, BitControl.PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    while (exec.next()) {
        final IntVector intMulVector = exec.getValueVectorById(new SchemaPath("INTMUL", ExpressionPosition.UNKNOWN), IntVector.class);
        final Float8Vector floatMulVector = exec.getValueVectorById(new SchemaPath("FLOATMUL", ExpressionPosition.UNKNOWN), Float8Vector.class);
        final IntVector intAddVector = exec.getValueVectorById(new SchemaPath("INTADD", ExpressionPosition.UNKNOWN), IntVector.class);
        final Float8Vector floatAddVector = exec.getValueVectorById(new SchemaPath("FLOATADD", ExpressionPosition.UNKNOWN), Float8Vector.class);
        assertEquals(exec.getRecordCount(), 1);
        assertEquals(intMulVector.getAccessor().get(0), 2);
        assertEquals(floatMulVector.getAccessor().get(0), (1.1 * 2.2), 0);
        assertEquals(intAddVector.getAccessor().get(0), 3);
        assertEquals(floatAddVector.getAccessor().get(0), (1.1 + 2.2), 0);
    }
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 93 with DrillbitContext

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

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)
 */
@SuppressWarnings("resource")
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(ClassTransformer.SCALAR_REPLACEMENT_OPTION);
    optionManager.setLocalOption(ClassTransformer.SCALAR_REPLACEMENT_OPTION, srOption.name().toLowerCase());
    // flush the code cache
    drillbitContext.getCompiler().flushCache();
    return originalOptionValue;
}

Example 94 with DrillbitContext

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

the class TestExceptionInjection method injectionOnSpecificBit.

@SuppressWarnings("static-method")
@Test
public void injectionOnSpecificBit() {
    final RemoteServiceSet remoteServiceSet = RemoteServiceSet.getLocalServiceSet();
    final ZookeeperHelper zkHelper = new ZookeeperHelper();
    zkHelper.startZookeeper(1);
    try {
        // Creating two drillbits
        final Drillbit drillbit1, drillbit2;
        final DrillConfig drillConfig = zkHelper.getConfig();
        try {
            drillbit1 = Drillbit.start(drillConfig, remoteServiceSet);
            drillbit2 = Drillbit.start(drillConfig, remoteServiceSet);
        } catch (DrillbitStartupException e) {
            throw new RuntimeException("Failed to start drillbits.", e);
        }
        final DrillbitContext drillbitContext1 = drillbit1.getContext();
        final DrillbitContext drillbitContext2 = drillbit2.getContext();
        final UserSession session = UserSession.Builder.newBuilder().withCredentials(UserBitShared.UserCredentials.newBuilder().setUserName("foo").build()).withUserProperties(UserProperties.getDefaultInstance()).withOptionManager(drillbitContext1.getOptionManager()).build();
        final String passthroughDesc = "<<injected from descPassthrough>>";
        final int nSkip = 7;
        final int nFire = 3;
        final Class<? extends Throwable> exceptionClass = RuntimeException.class;
        // only drillbit1's (address, port)
        final String controls = Controls.newBuilder().addExceptionOnBit(DummyClass.class, passthroughDesc, exceptionClass, drillbitContext1.getEndpoint(), nSkip, nFire).build();
        ControlsInjectionUtil.setControls(session, controls);
        {
            final QueryContext queryContext1 = new QueryContext(session, drillbitContext1, QueryId.getDefaultInstance());
            final DummyClass class1 = new DummyClass(queryContext1);
            // these shouldn't throw
            for (int i = 0; i < nSkip; ++i) {
                class1.descPassthroughMethod(passthroughDesc);
            }
            // these should throw
            for (int i = 0; i < nFire; ++i) {
                assertPassthroughThrows(class1, exceptionClass.getName(), passthroughDesc);
            }
            // this shouldn't throw
            class1.descPassthroughMethod(passthroughDesc);
            try {
                queryContext1.close();
            } catch (Exception e) {
                fail();
            }
        }
        {
            final QueryContext queryContext2 = new QueryContext(session, drillbitContext2, QueryId.getDefaultInstance());
            final DummyClass class2 = new DummyClass(queryContext2);
            // these shouldn't throw
            for (int i = 0; i < nSkip; ++i) {
                class2.descPassthroughMethod(passthroughDesc);
            }
            // these shouldn't throw
            for (int i = 0; i < nFire; ++i) {
                class2.descPassthroughMethod(passthroughDesc);
            }
            // this shouldn't throw
            class2.descPassthroughMethod(passthroughDesc);
            try {
                queryContext2.close();
            } catch (Exception e) {
                fail();
            }
        }
    } finally {
        zkHelper.stopZookeeper();
    }
}

Example 95 with DrillbitContext

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

the class TestPauseInjection method timedPauseOnSpecificBit.

@Test
public void timedPauseOnSpecificBit() {
    final RemoteServiceSet remoteServiceSet = RemoteServiceSet.getLocalServiceSet();
    final ZookeeperHelper zkHelper = new ZookeeperHelper();
    zkHelper.startZookeeper(1);
    final long pauseDuration = 2000L;
    final long expectedDuration = pauseDuration;
    try {
        // Creating two drillbits
        final Drillbit drillbit1, drillbit2;
        final DrillConfig drillConfig = zkHelper.getConfig();
        try {
            drillbit1 = Drillbit.start(drillConfig, remoteServiceSet);
            drillbit2 = Drillbit.start(drillConfig, remoteServiceSet);
        } catch (final DrillbitStartupException e) {
            throw new RuntimeException("Failed to start two drillbits.", e);
        }
        final DrillbitContext drillbitContext1 = drillbit1.getContext();
        final DrillbitContext drillbitContext2 = drillbit2.getContext();
        final UserSession session = UserSession.Builder.newBuilder().withCredentials(UserCredentials.newBuilder().setUserName("foo").build()).withUserProperties(UserProperties.getDefaultInstance()).withOptionManager(drillbitContext1.getOptionManager()).build();
        final DrillbitEndpoint drillbitEndpoint1 = drillbitContext1.getEndpoint();
        final String controls = Controls.newBuilder().addTimedPauseOnBit(DummyClass.class, DummyClass.PAUSES, drillbitEndpoint1, 0, pauseDuration).build();
        ControlsInjectionUtil.setControls(session, controls);
        {
            final ExtendedLatch trigger = new ExtendedLatch(1);
            final Pointer<Exception> ex = new Pointer<>();
            final QueryContext queryContext = new QueryContext(session, drillbitContext1, QueryId.getDefaultInstance());
            // test that the pause happens
            final DummyClass dummyClass = new DummyClass(queryContext, trigger);
            final long actualDuration = dummyClass.pauses();
            assertTrue(String.format("Test should stop for at least %d milliseconds.", expectedDuration), expectedDuration <= actualDuration);
            assertTrue("No exception should be thrown.", ex.value == null);
            try {
                queryContext.close();
            } catch (final Exception e) {
                fail("Failed to close query context: " + e);
            }
        }
        {
            final ExtendedLatch trigger = new ExtendedLatch(1);
            final QueryContext queryContext = new QueryContext(session, drillbitContext2, QueryId.getDefaultInstance());
            // if the resume did not happen, the test would hang
            final DummyClass dummyClass = new DummyClass(queryContext, trigger);
            dummyClass.pauses();
            try {
                queryContext.close();
            } catch (final Exception e) {
                fail("Failed to close query context: " + e);
            }
        }
    } finally {
        zkHelper.stopZookeeper();
    }
}