Search in sources :

Example 26 with VectorContainer

use of org.apache.drill.exec.record.VectorContainer in project drill by apache.

the class TestWriteToDisk method test.

@Test
@SuppressWarnings("static-method")
public void test() throws Exception {
    final List<ValueVector> vectorList = Lists.newArrayList();
    final DrillConfig config = DrillConfig.create();
    try (final RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
        final Drillbit bit = new Drillbit(config, serviceSet)) {
        bit.run();
        final DrillbitContext context = bit.getContext();
        final MaterializedField intField = MaterializedField.create("int", Types.required(TypeProtos.MinorType.INT));
        final MaterializedField binField = MaterializedField.create("binary", Types.required(TypeProtos.MinorType.VARBINARY));
        try (final IntVector intVector = (IntVector) TypeHelper.getNewVector(intField, context.getAllocator());
            final VarBinaryVector binVector = (VarBinaryVector) TypeHelper.getNewVector(binField, context.getAllocator())) {
            AllocationHelper.allocate(intVector, 4, 4);
            AllocationHelper.allocate(binVector, 4, 5);
            vectorList.add(intVector);
            vectorList.add(binVector);
            intVector.getMutator().setSafe(0, 0);
            binVector.getMutator().setSafe(0, "ZERO".getBytes());
            intVector.getMutator().setSafe(1, 1);
            binVector.getMutator().setSafe(1, "ONE".getBytes());
            intVector.getMutator().setSafe(2, 2);
            binVector.getMutator().setSafe(2, "TWO".getBytes());
            intVector.getMutator().setSafe(3, 3);
            binVector.getMutator().setSafe(3, "THREE".getBytes());
            intVector.getMutator().setValueCount(4);
            binVector.getMutator().setValueCount(4);
            VectorContainer container = new VectorContainer();
            container.addCollection(vectorList);
            container.setRecordCount(4);
            WritableBatch batch = WritableBatch.getBatchNoHVWrap(container.getRecordCount(), container, false);
            VectorAccessibleSerializable wrap = new VectorAccessibleSerializable(batch, context.getAllocator());
            Configuration conf = new Configuration();
            conf.set(FileSystem.FS_DEFAULT_NAME_KEY, FileSystem.DEFAULT_FS);
            final VectorAccessibleSerializable newWrap = new VectorAccessibleSerializable(context.getAllocator());
            try (final FileSystem fs = FileSystem.get(conf)) {
                final File tempDir = Files.createTempDir();
                tempDir.deleteOnExit();
                final Path path = new Path(tempDir.getAbsolutePath(), "drillSerializable");
                try (final FSDataOutputStream out = fs.create(path)) {
                    wrap.writeToStream(out);
                    out.close();
                }
                try (final FSDataInputStream in = fs.open(path)) {
                    newWrap.readFromStream(in);
                }
            }
            final VectorAccessible newContainer = newWrap.get();
            for (VectorWrapper<?> w : newContainer) {
                try (ValueVector vv = w.getValueVector()) {
                    int values = vv.getAccessor().getValueCount();
                    for (int i = 0; i < values; i++) {
                        final Object o = vv.getAccessor().getObject(i);
                        if (o instanceof byte[]) {
                            System.out.println(new String((byte[]) o));
                        } else {
                            System.out.println(o);
                        }
                    }
                }
            }
        }
    }
}
Also used : DrillbitContext(org.apache.drill.exec.server.DrillbitContext) Path(org.apache.hadoop.fs.Path) SchemaPath(org.apache.drill.common.expression.SchemaPath) IntVector(org.apache.drill.exec.vector.IntVector) Configuration(org.apache.hadoop.conf.Configuration) VectorAccessible(org.apache.drill.exec.record.VectorAccessible) MaterializedField(org.apache.drill.exec.record.MaterializedField) VarBinaryVector(org.apache.drill.exec.vector.VarBinaryVector) VectorContainer(org.apache.drill.exec.record.VectorContainer) ValueVector(org.apache.drill.exec.vector.ValueVector) DrillConfig(org.apache.drill.common.config.DrillConfig) Drillbit(org.apache.drill.exec.server.Drillbit) RemoteServiceSet(org.apache.drill.exec.server.RemoteServiceSet) FileSystem(org.apache.hadoop.fs.FileSystem) WritableBatch(org.apache.drill.exec.record.WritableBatch) FSDataInputStream(org.apache.hadoop.fs.FSDataInputStream) FSDataOutputStream(org.apache.hadoop.fs.FSDataOutputStream) File(java.io.File) ExecTest(org.apache.drill.exec.ExecTest) Test(org.junit.Test)

Example 27 with VectorContainer

use of org.apache.drill.exec.record.VectorContainer in project drill by apache.

the class QueryBuilder method rowSet.

/**
   * Run the query and return the first result set as a
   * {@link DirectRowSet} object that can be inspected directly
   * by the code using a {@link RowSetReader}.
   * <p>
   * An enhancement is to provide a way to read a series of result
   * batches as row sets.
   * @return a row set that represents the first batch returned from
   * the query
   * @throws RpcException if anything goes wrong
   */
public DirectRowSet rowSet() throws RpcException {
    // Ignore all but the first non-empty batch.
    QueryDataBatch dataBatch = null;
    for (QueryDataBatch batch : results()) {
        if (dataBatch == null && batch.getHeader().getRowCount() != 0) {
            dataBatch = batch;
        } else {
            batch.release();
        }
    }
    if (dataBatch == null) {
        return null;
    }
    // Unload the batch and convert to a row set.
    final RecordBatchLoader loader = new RecordBatchLoader(client.allocator());
    try {
        loader.load(dataBatch.getHeader().getDef(), dataBatch.getData());
        dataBatch.release();
        VectorContainer container = loader.getContainer();
        container.setRecordCount(loader.getRecordCount());
        return new DirectRowSet(client.allocator(), container);
    } catch (SchemaChangeException e) {
        throw new IllegalStateException(e);
    }
}
Also used : SchemaChangeException(org.apache.drill.exec.exception.SchemaChangeException) QueryDataBatch(org.apache.drill.exec.rpc.user.QueryDataBatch) RecordBatchLoader(org.apache.drill.exec.record.RecordBatchLoader) DirectRowSet(org.apache.drill.test.rowSet.DirectRowSet) VectorContainer(org.apache.drill.exec.record.VectorContainer)

Aggregations

VectorContainer (org.apache.drill.exec.record.VectorContainer)27 ValueVector (org.apache.drill.exec.vector.ValueVector)11 MaterializedField (org.apache.drill.exec.record.MaterializedField)8 SchemaChangeException (org.apache.drill.exec.exception.SchemaChangeException)6 SelectionVector4 (org.apache.drill.exec.record.selection.SelectionVector4)6 Stopwatch (com.google.common.base.Stopwatch)5 SortRecordBatchBuilder (org.apache.drill.exec.physical.impl.sort.SortRecordBatchBuilder)5 IOException (java.io.IOException)4 SchemaPath (org.apache.drill.common.expression.SchemaPath)4 BatchSchema (org.apache.drill.exec.record.BatchSchema)4 CachedVectorContainer (org.apache.drill.exec.cache.CachedVectorContainer)3 VectorAccessibleSerializable (org.apache.drill.exec.cache.VectorAccessibleSerializable)3 VectorWrapper (org.apache.drill.exec.record.VectorWrapper)3 WritableBatch (org.apache.drill.exec.record.WritableBatch)3 DrillBuf (io.netty.buffer.DrillBuf)2 LogicalExpression (org.apache.drill.common.expression.LogicalExpression)2 MajorType (org.apache.drill.common.types.TypeProtos.MajorType)2 ClassTransformationException (org.apache.drill.exec.exception.ClassTransformationException)2 OutOfMemoryException (org.apache.drill.exec.exception.OutOfMemoryException)2 RecordBatchData (org.apache.drill.exec.physical.impl.sort.RecordBatchData)2