Examples with WriteChannel org.apache.arrow.vector.ipc.WriteChannel used on opensource projects

Example 1 with WriteChannel

use of org.apache.arrow.vector.ipc.WriteChannel in project flink by apache.

the class ArrowUtils method readNextBatch.

private static byte[] readNextBatch(ReadableByteChannel channel) throws IOException {
    MessageMetadataResult metadata = MessageSerializer.readMessage(new ReadChannel(channel));
    if (metadata == null) {
        return null;
    }
    long bodyLength = metadata.getMessageBodyLength();
    // Only care about RecordBatch messages and skip the other kind of messages
    if (metadata.getMessage().headerType() == MessageHeader.RecordBatch) {
        // Buffer backed output large enough to hold 8-byte length + complete serialized message
        ByteArrayOutputStreamWithPos baos = new ByteArrayOutputStreamWithPos((int) (8 + metadata.getMessageLength() + bodyLength));
        // Write message metadata to ByteBuffer output stream
        MessageSerializer.writeMessageBuffer(new WriteChannel(Channels.newChannel(baos)), metadata.getMessageLength(), metadata.getMessageBuffer());
        baos.close();
        ByteBuffer result = ByteBuffer.wrap(baos.getBuf());
        result.position(baos.getPosition());
        result.limit(result.capacity());
        readFully(channel, result);
        return result.array();
    } else {
        if (bodyLength > 0) {
            // Skip message body if not a RecordBatch
            Channels.newInputStream(channel).skip(bodyLength);
        }
        // Proceed to next message
        return readNextBatch(channel);
    }
}

Example 2 with WriteChannel

use of org.apache.arrow.vector.ipc.WriteChannel 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();
}