Search in sources :

Example 1 with AbstractSubScan

use of org.apache.drill.exec.physical.base.AbstractSubScan in project drill by apache.

the class TestScanBatchWriters method sanityTest.

@Test
public void sanityTest() throws Exception {
    Scan scanConfig = new AbstractSubScan("bob") {

        @Override
        public String getOperatorType() {
            return "";
        }
    };
    OperatorContext opContext = fixture.newOperatorContext(scanConfig);
    // Setup: normally done by ScanBatch
    VectorContainer container = new VectorContainer(fixture.allocator());
    OutputMutator output = new ScanBatch.Mutator(opContext, fixture.allocator(), container);
    DrillBuf buffer = opContext.getManagedBuffer();
    try (VectorContainerWriter writer = new VectorContainerWriter(output)) {
        // Per-batch
        writer.allocate();
        writer.reset();
        BaseWriter.MapWriter map = writer.rootAsMap();
        // Write one record (10, "Fred", [100, 110, 120] )
        map.integer("a").writeInt(10);
        byte[] bytes = "Fred".getBytes("UTF-8");
        buffer.setBytes(0, bytes, 0, bytes.length);
        map.varChar("b").writeVarChar(0, bytes.length, buffer);
        try (ListWriter list = map.list("c")) {
            list.startList();
            list.integer().writeInt(100);
            list.integer().writeInt(110);
            list.integer().writeInt(120);
            list.endList();
            // Write another record: (20, "Wilma", [])
            writer.setPosition(1);
            map.integer("a").writeInt(20);
            bytes = "Wilma".getBytes("UTF-8");
            buffer.setBytes(0, bytes, 0, bytes.length);
            map.varChar("b").writeVarChar(0, bytes.length, buffer);
            writer.setValueCount(2);
            // Wrap-up done by ScanBatch
            container.setRecordCount(2);
            container.buildSchema(SelectionVectorMode.NONE);
            RowSet rowSet = fixture.wrap(container);
            // Expected
            TupleMetadata schema = new SchemaBuilder().addNullable("a", MinorType.INT).addNullable("b", MinorType.VARCHAR).addArray("c", MinorType.INT).buildSchema();
            RowSet expected = fixture.rowSetBuilder(schema).addRow(10, "Fred", new int[] { 100, 110, 120 }).addRow(20, "Wilma", null).build();
            new RowSetComparison(expected).verifyAndClearAll(rowSet);
        }
    } finally {
        opContext.close();
    }
}
Also used : AbstractSubScan(org.apache.drill.exec.physical.base.AbstractSubScan) VectorContainerWriter(org.apache.drill.exec.vector.complex.impl.VectorContainerWriter) BaseWriter(org.apache.drill.exec.vector.complex.writer.BaseWriter) RowSet(org.apache.drill.exec.physical.rowSet.RowSet) OutputMutator(org.apache.drill.exec.physical.impl.OutputMutator) VectorContainer(org.apache.drill.exec.record.VectorContainer) RowSetComparison(org.apache.drill.test.rowSet.RowSetComparison) OutputMutator(org.apache.drill.exec.physical.impl.OutputMutator) OperatorContext(org.apache.drill.exec.ops.OperatorContext) ListWriter(org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter) TupleMetadata(org.apache.drill.exec.record.metadata.TupleMetadata) SchemaBuilder(org.apache.drill.exec.record.metadata.SchemaBuilder) Scan(org.apache.drill.exec.physical.base.Scan) AbstractSubScan(org.apache.drill.exec.physical.base.AbstractSubScan) DrillBuf(io.netty.buffer.DrillBuf) SubOperatorTest(org.apache.drill.test.SubOperatorTest) Test(org.junit.Test)

Aggregations

DrillBuf (io.netty.buffer.DrillBuf)1 OperatorContext (org.apache.drill.exec.ops.OperatorContext)1 AbstractSubScan (org.apache.drill.exec.physical.base.AbstractSubScan)1 Scan (org.apache.drill.exec.physical.base.Scan)1 OutputMutator (org.apache.drill.exec.physical.impl.OutputMutator)1 RowSet (org.apache.drill.exec.physical.rowSet.RowSet)1 VectorContainer (org.apache.drill.exec.record.VectorContainer)1 SchemaBuilder (org.apache.drill.exec.record.metadata.SchemaBuilder)1 TupleMetadata (org.apache.drill.exec.record.metadata.TupleMetadata)1 VectorContainerWriter (org.apache.drill.exec.vector.complex.impl.VectorContainerWriter)1 BaseWriter (org.apache.drill.exec.vector.complex.writer.BaseWriter)1 ListWriter (org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter)1 SubOperatorTest (org.apache.drill.test.SubOperatorTest)1 RowSetComparison (org.apache.drill.test.rowSet.RowSetComparison)1 Test (org.junit.Test)1