Examples with BufferAllocator org.apache.drill.exec.memory.BufferAllocator used on opensource projects

Example 1 with BufferAllocator

use of org.apache.drill.exec.memory.BufferAllocator in project drill by apache.

the class TestBitRpc method testConnectionBackpressure.

@Test
public void testConnectionBackpressure(@Injectable WorkerBee bee, @Injectable final WorkEventBus workBus) throws Exception {
    DrillConfig config1 = DrillConfig.create();
    final BootStrapContext c = new BootStrapContext(config1, ClassPathScanner.fromPrescan(config1));
    DrillConfig config2 = DrillConfig.create();
    BootStrapContext c2 = new BootStrapContext(config2, ClassPathScanner.fromPrescan(config2));
    final FragmentContext fcon = new MockUp<FragmentContext>() {

        BufferAllocator getAllocator() {
            return c.getAllocator();
        }
    }.getMockInstance();
    final FragmentManager fman = new MockUp<FragmentManager>() {

        int v = 0;

        @Mock
        boolean handle(IncomingDataBatch batch) throws FragmentSetupException, IOException {
            try {
                v++;
                if (v % 10 == 0) {
                    System.out.println("sleeping.");
                    Thread.sleep(3000);
                }
            } catch (InterruptedException e) {
            }
            RawFragmentBatch rfb = batch.newRawFragmentBatch(c.getAllocator());
            rfb.sendOk();
            rfb.release();
            return true;
        }

        public FragmentContext getFragmentContext() {
            return fcon;
        }
    }.getMockInstance();
    new NonStrictExpectations() {

        {
            workBus.getFragmentManagerIfExists((FragmentHandle) any);
            result = fman;
            workBus.getFragmentManager((FragmentHandle) any);
            result = fman;
        }
    };
    int port = 1234;
    DataConnectionConfig config = new DataConnectionConfig(c.getAllocator(), c, new DataServerRequestHandler(workBus, bee));
    DataServer server = new DataServer(config);
    port = server.bind(port, true);
    DrillbitEndpoint ep = DrillbitEndpoint.newBuilder().setAddress("localhost").setDataPort(port).build();
    DataConnectionManager manager = new DataConnectionManager(ep, config);
    DataTunnel tunnel = new DataTunnel(manager);
    AtomicLong max = new AtomicLong(0);
    for (int i = 0; i < 40; i++) {
        long t1 = System.currentTimeMillis();
        tunnel.sendRecordBatch(new TimingOutcome(max), new FragmentWritableBatch(false, QueryId.getDefaultInstance(), 1, 1, 1, 1, getRandomBatch(c.getAllocator(), 5000)));
        System.out.println(System.currentTimeMillis() - t1);
    // System.out.println("sent.");
    }
    System.out.println(String.format("Max time: %d", max.get()));
    assertTrue(max.get() > 2700);
    Thread.sleep(5000);
}

Example 2 with BufferAllocator

use of org.apache.drill.exec.memory.BufferAllocator in project drill by apache.

the class WebUserConnection method sendData.

@Override
public void sendData(RpcOutcomeListener<Ack> listener, QueryWritableBatch result) {
    // Check if there is any data or not. There can be overflow here but DrillBuf doesn't support allocating with
    // bytes in long. Hence we are just preserving the earlier behavior and logging debug log for the case.
    final int dataByteCount = (int) result.getByteCount();
    if (dataByteCount <= 0) {
        if (logger.isDebugEnabled()) {
            logger.debug("Either no data received in this batch or there is BufferOverflow in dataByteCount: {}", dataByteCount);
        }
        listener.success(Acks.OK, null);
        return;
    }
    // If here that means there is some data for sure. Create a ByteBuf with all the data in it.
    final int rows = result.getHeader().getRowCount();
    final BufferAllocator allocator = webSessionResources.getAllocator();
    final DrillBuf bufferWithData = allocator.buffer(dataByteCount);
    try {
        final ByteBuf[] resultDataBuffers = result.getBuffers();
        for (final ByteBuf buffer : resultDataBuffers) {
            bufferWithData.writeBytes(buffer);
            buffer.release();
        }
        final RecordBatchLoader loader = new RecordBatchLoader(allocator);
        try {
            loader.load(result.getHeader().getDef(), bufferWithData);
            // SchemaChangeException, so check/clean catch clause below.
            for (int i = 0; i < loader.getSchema().getFieldCount(); ++i) {
                columns.add(loader.getSchema().getColumn(i).getPath());
            }
            for (int i = 0; i < rows; ++i) {
                final Map<String, String> record = Maps.newHashMap();
                for (VectorWrapper<?> vw : loader) {
                    final String field = vw.getValueVector().getMetadata().getNamePart().getName();
                    final Accessor accessor = vw.getValueVector().getAccessor();
                    final Object value = i < accessor.getValueCount() ? accessor.getObject(i) : null;
                    final String display = value == null ? null : value.toString();
                    record.put(field, display);
                }
                results.add(record);
            }
        } finally {
            loader.clear();
        }
    } catch (Exception e) {
        exception = UserException.systemError(e).build(logger);
    } finally {
        // Notify the listener with ACK.OK both in error/success case because data was send successfully from Drillbit.
        bufferWithData.release();
        listener.success(Acks.OK, null);
    }
}

Example 3 with BufferAllocator

use of org.apache.drill.exec.memory.BufferAllocator in project drill by apache.

the class PhysicalOpUnitTestBase method mockOpContext.

protected void mockOpContext(long initReservation, long maxAllocation) throws Exception {
    final BufferAllocator allocator = this.allocator.newChildAllocator("allocator_for_operator_test", initReservation, maxAllocation);
    new NonStrictExpectations() {

        {
            opContext.getStats();
            result = opStats;
            opContext.getAllocator();
            result = allocator;
            fragContext.newOperatorContext(withAny(popConf));
            result = opContext;
        }
    };
}

Example 4 with BufferAllocator

use of org.apache.drill.exec.memory.BufferAllocator in project drill by apache.

the class ParquetRecordReaderTest method testPerformance.

@Test
@Ignore
public void testPerformance(@Injectable final DrillbitContext bitContext, @Injectable UserClientConnection connection) throws Exception {
    final DrillConfig c = DrillConfig.create();
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContext context = new FragmentContext(bitContext, BitControl.PlanFragment.getDefaultInstance(), connection, registry);
    //    new NonStrictExpectations() {
    //      {
    //        context.getAllocator(); result = BufferAllocator.getAllocator(DrillConfig.create());
    //      }
    //    };
    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);
    //generateParquetFile(fileName, 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 5 with BufferAllocator

use of org.apache.drill.exec.memory.BufferAllocator in project drill by apache.

the class ExternalSortBatch method createCopier.

private void createCopier(VectorAccessible batch, List<BatchGroup> batchGroupList, VectorContainer outputContainer, boolean spilling) throws SchemaChangeException {
    try {
        if (copier == null) {
            CodeGenerator<PriorityQueueCopier> cg = CodeGenerator.get(PriorityQueueCopier.TEMPLATE_DEFINITION, context.getFunctionRegistry(), context.getOptions());
            cg.plainJavaCapable(true);
            // Uncomment out this line to debug the generated code.
            //        cg.saveCodeForDebugging(true);
            ClassGenerator<PriorityQueueCopier> g = cg.getRoot();
            generateComparisons(g, batch);
            g.setMappingSet(COPIER_MAPPING_SET);
            CopyUtil.generateCopies(g, batch, true);
            g.setMappingSet(MAIN_MAPPING);
            copier = context.getImplementationClass(cg);
        } else {
            copier.close();
        }
        @SuppressWarnings("resource") BufferAllocator allocator = spilling ? copierAllocator : oAllocator;
        for (VectorWrapper<?> i : batch) {
            @SuppressWarnings("resource") ValueVector v = TypeHelper.getNewVector(i.getField(), allocator);
            outputContainer.add(v);
        }
        copier.setup(context, allocator, batch, batchGroupList, outputContainer);
    } catch (ClassTransformationException | IOException e) {
        throw new RuntimeException(e);
    }
}