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

Example 36 with DrillbitContext

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

the class TestCustomTunnel method ensureRoundTrip.

@Test
public void ensureRoundTrip() throws Exception {
    final DrillbitContext context = getDrillbitContext();
    final TestCustomMessageHandler handler = new TestCustomMessageHandler(context.getEndpoint(), false);
    context.getController().registerCustomHandler(1001, handler, DrillbitEndpoint.PARSER);
    final ControlTunnel loopbackTunnel = context.getController().getTunnel(context.getEndpoint());
    final CustomTunnel<DrillbitEndpoint, QueryId> tunnel = loopbackTunnel.getCustomTunnel(1001, DrillbitEndpoint.class, QueryId.PARSER);
    CustomFuture<QueryId> future = tunnel.send(context.getEndpoint());
    assertEquals(expectedId, future.get());
}

Example 37 with DrillbitContext

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

the class TestRecordIterator method testSimpleIterator.

@Test
public void testSimpleIterator() throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final PhysicalPlanReader reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    final String planStr = Files.toString(DrillFileUtils.getResourceAsFile("/record/test_recorditerator.json"), Charsets.UTF_8);
    final PhysicalPlan plan = reader.readPhysicalPlan(planStr);
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, BitControl.PlanFragment.getDefaultInstance(), connection, registry);
    final List<PhysicalOperator> operatorList = plan.getSortedOperators(false);
    SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) operatorList.iterator().next()));
    RecordBatch singleBatch = exec.getIncoming();
    PhysicalOperator dummyPop = operatorList.iterator().next();
    OpProfileDef def = new OpProfileDef(dummyPop.getOperatorId(), UserBitShared.CoreOperatorType.MOCK_SUB_SCAN_VALUE, OperatorUtilities.getChildCount(dummyPop));
    OperatorStats stats = exec.getContext().getStats().newOperatorStats(def, exec.getContext().getAllocator());
    RecordIterator iter = new RecordIterator(singleBatch, null, exec.getContext().newOperatorContext(dummyPop, stats), 0, false, null);
    int totalRecords = 0;
    List<ValueVector> vectors = null;
    while (true) {
        iter.next();
        if (iter.finished()) {
            break;
        } else {
            // First time save vectors.
            if (vectors == null) {
                vectors = Lists.newArrayList();
                for (VectorWrapper vw : iter) {
                    vectors.add(vw.getValueVector());
                }
            }
            final int position = iter.getCurrentPosition();
            if (position % 2 == 0) {
                assertTrue(checkValues(vectors, position));
            } else {
                assertTrue(checkValues(vectors, position));
            }
            totalRecords++;
        }
        assertEquals(0, iter.cachedBatches().size());
    }
    assertEquals(11112, totalRecords);
    try {
        iter.mark();
        assertTrue(false);
    } catch (UnsupportedOperationException e) {
    }
    try {
        iter.reset();
        assertTrue(false);
    } catch (UnsupportedOperationException e) {
    }
}

Example 38 with DrillbitContext

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

the class ParquetRecordReaderTest method testPerformance.

@Test
@Ignore
public void testPerformance() throws Exception {
    final DrillbitContext bitContext = mock(DrillbitContext.class);
    final UserClientConnection connection = mock(UserClientConnection.class);
    final DrillConfig c = DrillConfig.create();
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, BitControl.PlanFragment.getDefaultInstance(), connection, registry);
    final String fileName = "/tmp/parquet_test_performance.parquet";
    final HashMap<String, FieldInfo> fields = new HashMap<>();
    final ParquetTestProperties props = new ParquetTestProperties(1, 20 * 1000 * 1000, DEFAULT_BYTES_PER_PAGE, fields);
    populateFieldInfoMap(props);
    final Configuration dfsConfig = new Configuration();
    final List<Footer> footers = ParquetFileReader.readFooters(dfsConfig, new Path(fileName));
    final Footer f = footers.iterator().next();
    final List<SchemaPath> columns = Lists.newArrayList();
    columns.add(new SchemaPath("_MAP.integer", ExpressionPosition.UNKNOWN));
    columns.add(new SchemaPath("_MAP.bigInt", ExpressionPosition.UNKNOWN));
    columns.add(new SchemaPath("_MAP.f", ExpressionPosition.UNKNOWN));
    columns.add(new SchemaPath("_MAP.d", ExpressionPosition.UNKNOWN));
    columns.add(new SchemaPath("_MAP.b", ExpressionPosition.UNKNOWN));
    columns.add(new SchemaPath("_MAP.bin", ExpressionPosition.UNKNOWN));
    columns.add(new SchemaPath("_MAP.bin2", ExpressionPosition.UNKNOWN));
    int totalRowCount = 0;
    final FileSystem fs = new CachedSingleFileSystem(fileName);
    final BufferAllocator allocator = RootAllocatorFactory.newRoot(c);
    for (int i = 0; i < 25; i++) {
        final ParquetRecordReader rr = new ParquetRecordReader(context, fileName, 0, fs, CodecFactory.createDirectCodecFactory(dfsConfig, new ParquetDirectByteBufferAllocator(allocator), 0), f.getParquetMetadata(), columns, ParquetReaderUtility.DateCorruptionStatus.META_SHOWS_CORRUPTION);
        final TestOutputMutator mutator = new TestOutputMutator(allocator);
        rr.setup(null, mutator);
        final Stopwatch watch = Stopwatch.createStarted();
        int rowCount = 0;
        while ((rowCount = rr.next()) > 0) {
            totalRowCount += rowCount;
        }
        System.out.println(String.format("Time completed: %s. ", watch.elapsed(TimeUnit.MILLISECONDS)));
        rr.close();
    }
    allocator.close();
    System.out.println(String.format("Total row count %s", totalRowCount));
}

Example 39 with DrillbitContext

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

the class TestPauseInjection method pauseOnSpecificBit.

@Test
public void pauseOnSpecificBit() {
    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 (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().addPauseOnBit(DummyClass.class, DummyClass.PAUSES, drillbitEndpoint1).build();
        ControlsInjectionUtil.setControls(session, controls);
        {
            final long expectedDuration = 1000L;
            final ExtendedLatch trigger = new ExtendedLatch(1);
            final Pointer<Exception> ex = new Pointer<>();
            final QueryContext queryContext = new QueryContext(session, drillbitContext1, QueryId.getDefaultInstance());
            (new ResumingThread(queryContext, trigger, ex, expectedDuration)).start();
            // 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();
    }
}

Example 40 with DrillbitContext

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

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) {
    @SuppressWarnings("resource") final DrillbitContext drillbitContext = drillbit.getContext();
    @SuppressWarnings("resource") final OptionManager optionManager = drillbitContext.getOptionManager();
    optionManager.setLocalOption(ClassTransformer.SCALAR_REPLACEMENT_OPTION, srOption);
    // flush the code cache
    drillbitContext.getCompiler().flushCache();
}