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

Example 11 with BufferAllocator

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

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;
    DataServer server = new DataServer(c, c.getAllocator(), workBus, null);
    port = server.bind(port, true);
    DrillbitEndpoint ep = DrillbitEndpoint.newBuilder().setAddress("localhost").setDataPort(port).build();
    DataConnectionManager manager = new DataConnectionManager(ep, c2);
    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 12 with BufferAllocator

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

the class TestBsonRecordReader method setUp.

@BeforeClass
public static void setUp() {
    BufferAllocator bufferAllocator = getDrillbitContext().getAllocator();
    mutator = new TestOutputMutator(bufferAllocator);
    writer = new VectorContainerWriter(mutator);
    bsonReader = new BsonRecordReader(bufferAllocator.buffer(1024), false, false);
}

Example 13 with BufferAllocator

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

the class TestPromotableWriter method list.

@Test
public void list() throws Exception {
    BufferAllocator allocator = RootAllocatorFactory.newRoot(DrillConfig.create());
    TestOutputMutator output = new TestOutputMutator(allocator);
    ComplexWriter rootWriter = new VectorContainerWriter(output, true);
    MapWriter writer = rootWriter.rootAsMap();
    rootWriter.setPosition(0);
    {
        writer.map("map").bigInt("a").writeBigInt(1);
    }
    rootWriter.setPosition(1);
    {
        writer.map("map").float4("a").writeFloat4(2.0f);
    }
    rootWriter.setPosition(2);
    {
        writer.map("map").list("a").startList();
        writer.map("map").list("a").endList();
    }
    rootWriter.setPosition(3);
    {
        writer.map("map").list("a").startList();
        writer.map("map").list("a").bigInt().writeBigInt(3);
        writer.map("map").list("a").float4().writeFloat4(4);
        writer.map("map").list("a").endList();
    }
    rootWriter.setValueCount(4);
    BatchPrinter.printBatch(output.getContainer());
}

Example 14 with BufferAllocator

use of org.apache.drill.exec.memory.BufferAllocator 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 15 with BufferAllocator

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

the class TestSplitAndTransfer method test.

@Test
public void test() throws Exception {
    final DrillConfig drillConfig = DrillConfig.create();
    final BufferAllocator allocator = RootAllocatorFactory.newRoot(drillConfig);
    final MaterializedField field = MaterializedField.create("field", Types.optional(MinorType.VARCHAR));
    final NullableVarCharVector varCharVector = new NullableVarCharVector(field, allocator);
    varCharVector.allocateNew(10000, 1000);
    final int valueCount = 500;
    final String[] compareArray = new String[valueCount];
    final NullableVarCharVector.Mutator mutator = varCharVector.getMutator();
    for (int i = 0; i < valueCount; i += 3) {
        final String s = String.format("%010d", i);
        mutator.set(i, s.getBytes());
        compareArray[i] = s;
    }
    mutator.setValueCount(valueCount);
    final TransferPair tp = varCharVector.getTransferPair(allocator);
    final NullableVarCharVector newVarCharVector = (NullableVarCharVector) tp.getTo();
    final Accessor accessor = newVarCharVector.getAccessor();
    final int[][] startLengths = { { 0, 201 }, { 201, 200 }, { 401, 99 } };
    for (final int[] startLength : startLengths) {
        final int start = startLength[0];
        final int length = startLength[1];
        tp.splitAndTransfer(start, length);
        newVarCharVector.getMutator().setValueCount(length);
        for (int i = 0; i < length; i++) {
            final boolean expectedSet = ((start + i) % 3) == 0;
            if (expectedSet) {
                final byte[] expectedValue = compareArray[start + i].getBytes();
                assertFalse(accessor.isNull(i));
                assertArrayEquals(expectedValue, accessor.get(i));
            } else {
                assertTrue(accessor.isNull(i));
            }
        }
        newVarCharVector.clear();
    }
    varCharVector.close();
    allocator.close();
}