Search in sources :

Example 1 with ArrowSchema

use of com.google.cloud.bigquery.storage.v1.ArrowSchema in project beam by apache.

the class BigQueryIOStorageReadTest method createResponseArrow.

private ReadRowsResponse createResponseArrow(org.apache.arrow.vector.types.pojo.Schema arrowSchema, List<String> name, List<Long> number, double progressAtResponseStart, double progressAtResponseEnd) {
    ArrowRecordBatch serializedRecord;
    try (VectorSchemaRoot schemaRoot = VectorSchemaRoot.create(arrowSchema, allocator)) {
        schemaRoot.allocateNew();
        schemaRoot.setRowCount(name.size());
        VarCharVector strVector = (VarCharVector) schemaRoot.getFieldVectors().get(0);
        BigIntVector bigIntVector = (BigIntVector) schemaRoot.getFieldVectors().get(1);
        for (int i = 0; i < name.size(); i++) {
            bigIntVector.set(i, number.get(i));
            strVector.set(i, new Text(name.get(i)));
        }
        VectorUnloader unLoader = new VectorUnloader(schemaRoot);
        try (org.apache.arrow.vector.ipc.message.ArrowRecordBatch records = unLoader.getRecordBatch()) {
            try (ByteArrayOutputStream os = new ByteArrayOutputStream()) {
                MessageSerializer.serialize(new WriteChannel(Channels.newChannel(os)), records);
                serializedRecord = ArrowRecordBatch.newBuilder().setRowCount(records.getLength()).setSerializedRecordBatch(ByteString.copyFrom(os.toByteArray())).build();
            } catch (IOException e) {
                throw new RuntimeException("Error writing to byte array output stream", e);
            }
        }
    }
    return ReadRowsResponse.newBuilder().setArrowRecordBatch(serializedRecord).setRowCount(name.size()).setStats(StreamStats.newBuilder().setProgress(Progress.newBuilder().setAtResponseStart(progressAtResponseStart).setAtResponseEnd(progressAtResponseEnd))).build();
}
Also used : VectorSchemaRoot(org.apache.arrow.vector.VectorSchemaRoot) VarCharVector(org.apache.arrow.vector.VarCharVector) Text(org.apache.arrow.vector.util.Text) ByteArrayOutputStream(java.io.ByteArrayOutputStream) IOException(java.io.IOException) BigIntVector(org.apache.arrow.vector.BigIntVector) VectorUnloader(org.apache.arrow.vector.VectorUnloader) StatusRuntimeException(io.grpc.StatusRuntimeException) ArrowRecordBatch(com.google.cloud.bigquery.storage.v1.ArrowRecordBatch) WriteChannel(org.apache.arrow.vector.ipc.WriteChannel)

Example 2 with ArrowSchema

use of com.google.cloud.bigquery.storage.v1.ArrowSchema in project beam by apache.

the class BigQueryStorageArrowReader method processReadRowsResponse.

@Override
public void processReadRowsResponse(ReadRowsResponse readRowsResponse) throws IOException {
    com.google.cloud.bigquery.storage.v1.ArrowRecordBatch recordBatch = readRowsResponse.getArrowRecordBatch();
    rowCount = recordBatch.getRowCount();
    this.alloc = new RootAllocator(Long.MAX_VALUE);
    InputStream input = protoSchema.getSerializedSchema().newInput();
    Schema arrowSchema = ArrowConversion.arrowSchemaFromInput(input);
    this.recordBatchIterator = ArrowConversion.rowsFromSerializedRecordBatch(arrowSchema, recordBatch.getSerializedRecordBatch().newInput(), this.alloc);
}
Also used : RootAllocator(org.apache.arrow.memory.RootAllocator) InputStream(java.io.InputStream) Schema(org.apache.arrow.vector.types.pojo.Schema) ArrowSchema(com.google.cloud.bigquery.storage.v1.ArrowSchema)

Aggregations

ArrowRecordBatch (com.google.cloud.bigquery.storage.v1.ArrowRecordBatch)1 ArrowSchema (com.google.cloud.bigquery.storage.v1.ArrowSchema)1 StatusRuntimeException (io.grpc.StatusRuntimeException)1 ByteArrayOutputStream (java.io.ByteArrayOutputStream)1 IOException (java.io.IOException)1 InputStream (java.io.InputStream)1 RootAllocator (org.apache.arrow.memory.RootAllocator)1 BigIntVector (org.apache.arrow.vector.BigIntVector)1 VarCharVector (org.apache.arrow.vector.VarCharVector)1 VectorSchemaRoot (org.apache.arrow.vector.VectorSchemaRoot)1 VectorUnloader (org.apache.arrow.vector.VectorUnloader)1 WriteChannel (org.apache.arrow.vector.ipc.WriteChannel)1 Schema (org.apache.arrow.vector.types.pojo.Schema)1 Text (org.apache.arrow.vector.util.Text)1