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

Example 21 with BufferAllocator

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

the class TestLoad method testSchemaChange.

@Test
public void testSchemaChange() throws SchemaChangeException {
    final BufferAllocator allocator = RootAllocatorFactory.newRoot(drillConfig);
    final RecordBatchLoader batchLoader = new RecordBatchLoader(allocator);
    // Initial schema: a: INT, b: VARCHAR
    // Schema change: N/A
    SchemaBuilder schemaBuilder1 = new SchemaBuilder().add("a", MinorType.INT).add("b", MinorType.VARCHAR);
    BatchSchema schema1 = new BatchSchemaBuilder().withSchemaBuilder(schemaBuilder1).build();
    {
        assertTrue(loadBatch(allocator, batchLoader, schema1));
        assertTrue(schema1.isEquivalent(batchLoader.getSchema()));
        batchLoader.getContainer().zeroVectors();
    }
    // Same schema
    // Schema change: No
    {
        assertFalse(loadBatch(allocator, batchLoader, schema1));
        assertTrue(schema1.isEquivalent(batchLoader.getSchema()));
        batchLoader.getContainer().zeroVectors();
    }
    // Reverse columns: b: VARCHAR, a: INT
    // Schema change: No
    {
        SchemaBuilder schemaBuilder = new SchemaBuilder().add("b", MinorType.VARCHAR).add("a", MinorType.INT);
        BatchSchema schema = new BatchSchemaBuilder().withSchemaBuilder(schemaBuilder).build();
        assertFalse(loadBatch(allocator, batchLoader, schema));
        // Potential bug: see DRILL-5828
        assertTrue(schema.isEquivalent(batchLoader.getSchema()));
        batchLoader.getContainer().zeroVectors();
    }
    // Drop a column: a: INT
    // Schema change: Yes
    {
        SchemaBuilder schemaBuilder = new SchemaBuilder().add("a", MinorType.INT);
        BatchSchema schema = new BatchSchemaBuilder().withSchemaBuilder(schemaBuilder).build();
        assertTrue(loadBatch(allocator, batchLoader, schema));
        assertTrue(schema.isEquivalent(batchLoader.getSchema()));
        batchLoader.getContainer().zeroVectors();
    }
    // Add a column: a: INT, b: VARCHAR, c: INT
    // Schema change: Yes
    {
        assertTrue(loadBatch(allocator, batchLoader, schema1));
        assertTrue(schema1.isEquivalent(batchLoader.getSchema()));
        batchLoader.getContainer().zeroVectors();
        SchemaBuilder schemaBuilder = new SchemaBuilder().add("a", MinorType.INT).add("b", MinorType.VARCHAR).add("c", MinorType.INT);
        BatchSchema schema = new BatchSchemaBuilder().withSchemaBuilder(schemaBuilder).build();
        assertTrue(loadBatch(allocator, batchLoader, schema));
        assertTrue(schema.isEquivalent(batchLoader.getSchema()));
        batchLoader.getContainer().zeroVectors();
    }
    // Change a column type: a: INT, b: VARCHAR, c: VARCHAR
    // Schema change: Yes
    {
        SchemaBuilder schemaBuilder = new SchemaBuilder().add("a", MinorType.INT).add("b", MinorType.VARCHAR).add("c", MinorType.VARCHAR);
        BatchSchema schema = new BatchSchemaBuilder().withSchemaBuilder(schemaBuilder).build();
        assertTrue(loadBatch(allocator, batchLoader, schema));
        assertTrue(schema.isEquivalent(batchLoader.getSchema()));
        batchLoader.getContainer().zeroVectors();
    }
    // Empty schema
    // Schema change: Yes
    {
        BatchSchema schema = new BatchSchemaBuilder().withSchemaBuilder(new SchemaBuilder()).build();
        assertTrue(loadBatch(allocator, batchLoader, schema));
        assertTrue(schema.isEquivalent(batchLoader.getSchema()));
        batchLoader.getContainer().zeroVectors();
    }
    batchLoader.clear();
    allocator.close();
}

Example 22 with BufferAllocator

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

the class TestMemoryRetention method main.

public static void main(String[] args) throws Exception {
    final DrillConfig config = DrillConfig.create();
    final BufferAllocator a = RootAllocatorFactory.newRoot(config);
    for (int i = 0; i < PARALLEL_THREADS; i++) {
        Alloc alloc = new Alloc(a);
        alloc.start();
    }
}

Example 23 with BufferAllocator

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

the class DrillParquetReader method setup.

@Override
public void setup(OperatorContext context, OutputMutator output) throws ExecutionSetupException {
    try {
        this.operatorContext = context;
        schema = footer.getFileMetaData().getSchema();
        MessageType projection;
        final List<SchemaPath> columnsNotFound = new ArrayList<>(getColumns().size());
        if (isStarQuery()) {
            projection = schema;
        } else {
            projection = getProjection(schema, getColumns(), columnsNotFound);
            if (projection == null) {
                projection = schema;
            }
            if (!columnsNotFound.isEmpty()) {
                nullFilledVectors = new ArrayList<>(columnsNotFound.size());
                for (SchemaPath col : columnsNotFound) {
                    // col.toExpr() is used here as field name since we don't want to see these fields in the existing maps
                    nullFilledVectors.add(output.addField(MaterializedField.create(col.toExpr(), OPTIONAL_INT), NullableIntVector.class));
                }
                noColumnsFound = columnsNotFound.size() == getColumns().size();
            }
        }
        logger.debug("Requesting schema {}", projection);
        if (!noColumnsFound) {
            // Discard the columns not found in the schema when create DrillParquetRecordMaterializer, since they have been added to output already.
            @SuppressWarnings("unchecked") Collection<SchemaPath> columns = columnsNotFound.isEmpty() ? getColumns() : CollectionUtils.subtract(getColumns(), columnsNotFound);
            recordMaterializer = new DrillParquetRecordMaterializer(output, projection, columns, fragmentContext.getOptions(), containsCorruptedDates);
        }
        if (numRecordsToRead == 0 || noColumnsFound) {
            // no need to init readers
            return;
        }
        ColumnIOFactory factory = new ColumnIOFactory(false);
        MessageColumnIO columnIO = factory.getColumnIO(projection, schema);
        BlockMetaData blockMetaData = footer.getBlocks().get(entry.getRowGroupIndex());
        Map<ColumnPath, ColumnChunkMetaData> paths = blockMetaData.getColumns().stream().collect(Collectors.toMap(ColumnChunkMetaData::getPath, Function.identity(), (o, n) -> n));
        BufferAllocator allocator = operatorContext.getAllocator();
        CompressionCodecFactory ccf = DrillCompressionCodecFactory.createDirectCodecFactory(drillFileSystem.getConf(), new ParquetDirectByteBufferAllocator(allocator), 0);
        pageReadStore = new ColumnChunkIncReadStore(numRecordsToRead, ccf, allocator, drillFileSystem, entry.getPath());
        for (String[] path : schema.getPaths()) {
            Type type = schema.getType(path);
            if (type.isPrimitive()) {
                ColumnChunkMetaData md = paths.get(ColumnPath.get(path));
                pageReadStore.addColumn(schema.getColumnDescription(path), md);
            }
        }
        recordReader = columnIO.getRecordReader(pageReadStore, recordMaterializer);
    } catch (Exception e) {
        throw handleAndRaise("Failure in setting up reader", e);
    }
}

Example 24 with BufferAllocator

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

the class AbstractGenericCopierTest method testCopyRecords.

@Test
public void testCopyRecords() throws Exception {
    try (OperatorFixture operatorFixture = new OperatorFixture.Builder(baseDirTestWatcher).build()) {
        final BufferAllocator allocator = operatorFixture.allocator();
        final BatchSchema batchSchema = createTestSchema(BatchSchema.SelectionVectorMode.NONE);
        final RowSet srcRowSet = createSrcRowSet(allocator);
        final VectorContainer destContainer = new VectorContainer(allocator, batchSchema);
        destContainer.setRecordCount(0);
        final RowSet expectedRowSet = createExpectedRowset(allocator);
        MockRecordBatch mockRecordBatch = null;
        try {
            mockRecordBatch = new MockRecordBatch.Builder().sendData(srcRowSet).build(operatorFixture.getFragmentContext());
            mockRecordBatch.next();
            final Copier copier = createCopier(mockRecordBatch, destContainer, null);
            copier.copyRecords(0, 3);
            new RowSetComparison(expectedRowSet).verify(DirectRowSet.fromContainer(destContainer));
        } finally {
            if (mockRecordBatch != null) {
                mockRecordBatch.close();
            }
            srcRowSet.clear();
            destContainer.clear();
            expectedRowSet.clear();
        }
    }
}

Example 25 with BufferAllocator

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

the class TestBitBitKerberos method setupFragmentContextAndManager.

private static void setupFragmentContextAndManager() {
    final FragmentContext fcontext = new MockUp<FragmentContext>() {

        @SuppressWarnings("unused")
        BufferAllocator getAllocator() {
            return c1.getAllocator();
        }
    }.getMockInstance();
    manager = 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(c1.getAllocator());
            rfb.sendOk();
            rfb.release();
            return true;
        }

        @SuppressWarnings("unused")
        public FragmentContext getFragmentContext() {
            return fcontext;
        }
    }.getMockInstance();
}