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

Example 41 with DrillbitContext

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

the class TestPartitionSender method testPartitionSenderCostToThreads.

@Test
public /**
 * Main test to go over different scenarios
 * @throws Exception
 */
void testPartitionSenderCostToThreads() throws Exception {
    final VectorContainer container = new VectorContainer();
    container.buildSchema(SelectionVectorMode.FOUR_BYTE);
    final SelectionVector4 sv = Mockito.mock(SelectionVector4.class, "SelectionVector4");
    Mockito.when(sv.getCount()).thenReturn(100);
    Mockito.when(sv.getTotalCount()).thenReturn(100);
    for (int i = 0; i < 100; i++) {
        Mockito.when(sv.get(i)).thenReturn(i);
    }
    final TopNBatch.SimpleSV4RecordBatch incoming = new TopNBatch.SimpleSV4RecordBatch(container, sv, null);
    updateTestCluster(DRILLBITS_COUNT, null);
    test("ALTER SESSION SET `planner.slice_target`=1");
    String plan = getPlanInString("EXPLAIN PLAN FOR " + groupByQuery, JSON_FORMAT);
    final DrillbitContext drillbitContext = getDrillbitContext();
    final PhysicalPlanReader planReader = drillbitContext.getPlanReader();
    final PhysicalPlan physicalPlan = planReader.readPhysicalPlan(plan);
    final Fragment rootFragment = PopUnitTestBase.getRootFragmentFromPlanString(planReader, plan);
    final PlanningSet planningSet = new PlanningSet();
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(config);
    // Create a planningSet to get the assignment of major fragment ids to fragments.
    PARALLELIZER.initFragmentWrappers(rootFragment, planningSet);
    final List<PhysicalOperator> operators = physicalPlan.getSortedOperators(false);
    // get HashToRandomExchange physical operator
    HashToRandomExchange hashToRandomExchange = null;
    for (PhysicalOperator operator : operators) {
        if (operator instanceof HashToRandomExchange) {
            hashToRandomExchange = (HashToRandomExchange) operator;
            break;
        }
    }
    final OptionList options = new OptionList();
    // try multiple scenarios with different set of options
    options.add(OptionValue.create(OptionValue.AccessibleScopes.SESSION, "planner.slice_target", 1, OptionScope.SESSION));
    testThreadsHelper(hashToRandomExchange, drillbitContext, options, incoming, registry, planReader, planningSet, rootFragment, 1);
    options.clear();
    options.add(OptionValue.create(AccessibleScopes.SESSION, "planner.slice_target", 1, OptionScope.SESSION));
    options.add(OptionValue.create(OptionValue.AccessibleScopes.SESSION, "planner.partitioner_sender_max_threads", 10, OptionScope.SESSION));
    hashToRandomExchange.setCost(new PrelCostEstimates(1000, 1000));
    testThreadsHelper(hashToRandomExchange, drillbitContext, options, incoming, registry, planReader, planningSet, rootFragment, 10);
    options.clear();
    options.add(OptionValue.create(AccessibleScopes.SESSION, "planner.slice_target", 1000, OptionScope.SESSION));
    options.add(OptionValue.create(AccessibleScopes.SESSION, "planner.partitioner_sender_threads_factor", 2, OptionScope.SESSION));
    hashToRandomExchange.setCost(new PrelCostEstimates(14000, 14000));
    testThreadsHelper(hashToRandomExchange, drillbitContext, options, incoming, registry, planReader, planningSet, rootFragment, 2);
}

Example 42 with DrillbitContext

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

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);
            assertNull("No exception should be thrown.", ex.value);
            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();
    }
}

Example 43 with DrillbitContext

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

the class TestHashJoin method testHJMockScanCommon.

private void testHJMockScanCommon(String physicalPlan, int expectedRows) 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(physicalPlan), 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()));
    int totalRecordCount = 0;
    while (exec.next()) {
        totalRecordCount += exec.getRecordCount();
    }
    exec.close();
    assertEquals(expectedRows, totalRecordCount);
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 44 with DrillbitContext

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

the class TestMergeJoin method simpleEqualityJoin.

@Test
// this doesn't have a sort.  it also causes an infinite loop.  these may or may not be related.
@Ignore
public void simpleEqualityJoin() 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("/join/merge_join.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()));
    int totalRecordCount = 0;
    final StringBuilder sb = new StringBuilder();
    while (exec.next()) {
        totalRecordCount += exec.getRecordCount();
        for (final ValueVector v : exec) {
            sb.append("[" + v.getField().getName() + "]        ");
        }
        sb.append("\n\n");
        for (int valueIdx = 0; valueIdx < exec.getRecordCount(); valueIdx++) {
            final List<Object> row = new ArrayList<>();
            for (final ValueVector v : exec) {
                row.add(v.getAccessor().getObject(valueIdx));
            }
            for (final Object cell : row) {
                if (cell == null) {
                    sb.append("<null>          ");
                    continue;
                }
                final int len = cell.toString().length();
                sb.append(cell);
                for (int i = 0; i < (14 - len); ++i) {
                    sb.append(" ");
                }
            }
            sb.append("\n");
        }
        sb.append("\n");
    }
    logger.info(sb.toString());
    assertEquals(100, totalRecordCount);
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 45 with DrillbitContext

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

the class TestMergeJoin method testJoinBatchSize.

@Test
public void testJoinBatchSize() 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("/join/join_batchsize.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()));
    // skip schema batch
    exec.next();
    while (exec.next()) {
        assertEquals(100, exec.getRecordCount());
    }
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}