Examples with BarrageUpdateMetadata io.deephaven.barrage.flatbuf.BarrageUpdateMetadata used on opensource projects

Example 1 with BarrageUpdateMetadata

use of io.deephaven.barrage.flatbuf.BarrageUpdateMetadata in project deephaven-core by deephaven.

the class BarrageStreamReader method safelyParseFrom.

@Override
public BarrageMessage safelyParseFrom(final StreamReaderOptions options, final BitSet expectedColumns, final ChunkType[] columnChunkTypes, final Class<?>[] columnTypes, final Class<?>[] componentTypes, final InputStream stream) {
    final long startDeserTm = System.nanoTime();
    Message header = null;
    try {
        boolean bodyParsed = false;
        final CodedInputStream decoder = CodedInputStream.newInstance(stream);
        for (int tag = decoder.readTag(); tag != 0; tag = decoder.readTag()) {
            if (tag == BarrageProtoUtil.DATA_HEADER_TAG) {
                final int size = decoder.readRawVarint32();
                header = Message.getRootAsMessage(ByteBuffer.wrap(decoder.readRawBytes(size)));
                continue;
            } else if (tag == BarrageProtoUtil.APP_METADATA_TAG) {
                final int size = decoder.readRawVarint32();
                final ByteBuffer msgAsBB = ByteBuffer.wrap(decoder.readRawBytes(size));
                final BarrageMessageWrapper wrapper = BarrageMessageWrapper.getRootAsBarrageMessageWrapper(msgAsBB);
                if (wrapper.magic() != BarrageUtil.FLATBUFFER_MAGIC) {
                    log.warn().append("BarrageStreamReader: skipping app_metadata that does not look like BarrageMessageWrapper").endl();
                } else if (wrapper.msgType() == BarrageMessageType.BarrageUpdateMetadata) {
                    if (msg != null) {
                        throw new IllegalStateException("Previous message was not complete; pending " + (numAddRowsTotal - numAddRowsRead) + " add rows and " + (numModRowsTotal - numModRowsRead) + " mod rows");
                    }
                    final BarrageUpdateMetadata metadata = BarrageUpdateMetadata.getRootAsBarrageUpdateMetadata(wrapper.msgPayloadAsByteBuffer());
                    msg = new BarrageMessage();
                    msg.isSnapshot = metadata.isSnapshot();
                    msg.snapshotRowSetIsReversed = metadata.effectiveReverseViewport();
                    numAddRowsRead = 0;
                    numModRowsRead = 0;
                    lastAddStartIndex = 0;
                    lastModStartIndex = 0;
                    if (msg.isSnapshot) {
                        final ByteBuffer effectiveViewport = metadata.effectiveViewportAsByteBuffer();
                        if (effectiveViewport != null) {
                            msg.snapshotRowSet = extractIndex(effectiveViewport);
                        }
                        final ByteBuffer effectiveSnapshotColumns = metadata.effectiveColumnSetAsByteBuffer();
                        if (effectiveSnapshotColumns != null) {
                            msg.snapshotColumns = extractBitSet(effectiveSnapshotColumns);
                        }
                    }
                    msg.firstSeq = metadata.firstSeq();
                    msg.lastSeq = metadata.lastSeq();
                    msg.rowsAdded = extractIndex(metadata.addedRowsAsByteBuffer());
                    msg.rowsRemoved = extractIndex(metadata.removedRowsAsByteBuffer());
                    msg.shifted = extractIndexShiftData(metadata.shiftDataAsByteBuffer());
                    final ByteBuffer rowsIncluded = metadata.addedRowsIncludedAsByteBuffer();
                    msg.rowsIncluded = rowsIncluded != null ? extractIndex(rowsIncluded) : msg.rowsAdded.copy();
                    msg.addColumnData = new BarrageMessage.AddColumnData[columnTypes.length];
                    for (int ci = 0; ci < msg.addColumnData.length; ++ci) {
                        msg.addColumnData[ci] = new BarrageMessage.AddColumnData();
                        msg.addColumnData[ci].type = columnTypes[ci];
                        msg.addColumnData[ci].componentType = componentTypes[ci];
                        msg.addColumnData[ci].data = new ArrayList<>();
                        // create an initial chunk of the correct size
                        final int chunkSize = (int) (Math.min(msg.rowsIncluded.size(), MAX_CHUNK_SIZE));
                        msg.addColumnData[ci].data.add(columnChunkTypes[ci].makeWritableChunk(chunkSize));
                    }
                    numAddRowsTotal = msg.rowsIncluded.size();
                    // if this message is a snapshot response (vs. subscription) then mod columns may be empty
                    numModRowsTotal = 0;
                    msg.modColumnData = new BarrageMessage.ModColumnData[metadata.modColumnNodesLength()];
                    for (int ci = 0; ci < msg.modColumnData.length; ++ci) {
                        msg.modColumnData[ci] = new BarrageMessage.ModColumnData();
                        msg.modColumnData[ci].type = columnTypes[ci];
                        msg.modColumnData[ci].componentType = componentTypes[ci];
                        msg.modColumnData[ci].data = new ArrayList<>();
                        final BarrageModColumnMetadata mcd = metadata.modColumnNodes(ci);
                        msg.modColumnData[ci].rowsModified = extractIndex(mcd.modifiedRowsAsByteBuffer());
                        // create an initial chunk of the correct size
                        final int chunkSize = (int) (Math.min(msg.modColumnData[ci].rowsModified.size(), MAX_CHUNK_SIZE));
                        msg.modColumnData[ci].data.add(columnChunkTypes[ci].makeWritableChunk(chunkSize));
                        numModRowsTotal = Math.max(numModRowsTotal, msg.modColumnData[ci].rowsModified.size());
                    }
                }
                continue;
            } else if (tag != BarrageProtoUtil.BODY_TAG) {
                decoder.skipField(tag);
                continue;
            }
            if (bodyParsed) {
                // 2) it's plain ol' inefficient!
                throw new IllegalStateException("Unexpected duplicate body tag");
            }
            if (header == null) {
                throw new IllegalStateException("Missing metadata header; cannot decode body");
            }
            if (header.headerType() != org.apache.arrow.flatbuf.MessageHeader.RecordBatch) {
                throw new IllegalStateException("Only know how to decode Schema/BarrageRecordBatch messages");
            }
            // throw an error when no app metadata (snapshots now provide by default)
            if (msg == null) {
                throw new IllegalStateException("Missing app metadata tag; cannot decode using BarrageStreamReader");
            }
            bodyParsed = true;
            final int size = decoder.readRawVarint32();
            final RecordBatch batch = (RecordBatch) header.header(new RecordBatch());
            msg.length = batch.length();
            // noinspection UnstableApiUsage
            try (final LittleEndianDataInputStream ois = new LittleEndianDataInputStream(new BarrageProtoUtil.ObjectInputStreamAdapter(decoder, size))) {
                final MutableInt bufferOffset = new MutableInt();
                final Iterator<ChunkInputStreamGenerator.FieldNodeInfo> fieldNodeIter = new FlatBufferIteratorAdapter<>(batch.nodesLength(), i -> new ChunkInputStreamGenerator.FieldNodeInfo(batch.nodes(i)));
                final TLongArrayList bufferInfo = new TLongArrayList(batch.buffersLength());
                for (int bi = 0; bi < batch.buffersLength(); ++bi) {
                    int offset = LongSizedDataStructure.intSize("BufferInfo", batch.buffers(bi).offset());
                    int length = LongSizedDataStructure.intSize("BufferInfo", batch.buffers(bi).length());
                    if (bi < batch.buffersLength() - 1) {
                        final int nextOffset = LongSizedDataStructure.intSize("BufferInfo", batch.buffers(bi + 1).offset());
                        // our parsers handle overhanging buffers
                        length += Math.max(0, nextOffset - offset - length);
                    }
                    bufferOffset.setValue(offset + length);
                    bufferInfo.add(length);
                }
                final TLongIterator bufferInfoIter = bufferInfo.iterator();
                // mod rows will be received.
                if (numAddRowsRead < numAddRowsTotal) {
                    for (int ci = 0; ci < msg.addColumnData.length; ++ci) {
                        final BarrageMessage.AddColumnData acd = msg.addColumnData[ci];
                        final long remaining = numAddRowsTotal - numAddRowsRead;
                        if (batch.length() > remaining) {
                            throw new IllegalStateException("Batch length exceeded the expected number of rows from app metadata");
                        }
                        // select the current chunk size and read the size
                        int lastChunkIndex = acd.data.size() - 1;
                        WritableChunk<Values> chunk = (WritableChunk<Values>) acd.data.get(lastChunkIndex);
                        int chunkSize = acd.data.get(lastChunkIndex).size();
                        final int chunkOffset;
                        long rowOffset = numAddRowsRead - lastAddStartIndex;
                        // reading the rows from this batch might overflow the existing chunk
                        if (rowOffset + batch.length() > chunkSize) {
                            lastAddStartIndex += chunkSize;
                            // create a new chunk before trying to write again
                            chunkSize = (int) (Math.min(remaining, MAX_CHUNK_SIZE));
                            chunk = columnChunkTypes[ci].makeWritableChunk(chunkSize);
                            acd.data.add(chunk);
                            chunkOffset = 0;
                            ++lastChunkIndex;
                        } else {
                            chunkOffset = (int) rowOffset;
                        }
                        // fill the chunk with data and assign back into the array
                        acd.data.set(lastChunkIndex, ChunkInputStreamGenerator.extractChunkFromInputStream(options, columnChunkTypes[ci], columnTypes[ci], componentTypes[ci], fieldNodeIter, bufferInfoIter, ois, chunk, chunkOffset, chunkSize));
                    }
                    numAddRowsRead += batch.length();
                } else {
                    for (int ci = 0; ci < msg.modColumnData.length; ++ci) {
                        final BarrageMessage.ModColumnData mcd = msg.modColumnData[ci];
                        long remaining = mcd.rowsModified.size() - numModRowsRead;
                        // need to add the batch row data to the column chunks
                        int lastChunkIndex = mcd.data.size() - 1;
                        WritableChunk<Values> chunk = (WritableChunk<Values>) mcd.data.get(lastChunkIndex);
                        int chunkSize = chunk.size();
                        final int chunkOffset;
                        long rowOffset = numModRowsRead - lastModStartIndex;
                        // this batch might overflow the chunk
                        if (rowOffset + Math.min(remaining, batch.length()) > chunkSize) {
                            lastModStartIndex += chunkSize;
                            // create a new chunk before trying to write again
                            chunkSize = (int) (Math.min(remaining, MAX_CHUNK_SIZE));
                            chunk = columnChunkTypes[ci].makeWritableChunk(chunkSize);
                            mcd.data.add(chunk);
                            chunkOffset = 0;
                            ++lastChunkIndex;
                        } else {
                            chunkOffset = (int) rowOffset;
                        }
                        // fill the chunk with data and assign back into the array
                        mcd.data.set(lastChunkIndex, ChunkInputStreamGenerator.extractChunkFromInputStream(options, columnChunkTypes[ci], columnTypes[ci], componentTypes[ci], fieldNodeIter, bufferInfoIter, ois, chunk, chunkOffset, chunkSize));
                    }
                    numModRowsRead += batch.length();
                }
            }
        }
        if (header != null && header.headerType() == MessageHeader.Schema) {
            // there is no body and our clients do not want to see schema messages
            return null;
        }
        if (!bodyParsed) {
            throw new IllegalStateException("Missing body tag");
        }
        deserializeTmConsumer.accept(System.nanoTime() - startDeserTm);
        if (numAddRowsRead == numAddRowsTotal && numModRowsRead == numModRowsTotal) {
            final BarrageMessage retval = msg;
            msg = null;
            return retval;
        }
        // otherwise, must wait for more data
        return null;
    } catch (final Exception e) {
        log.error().append("Unable to parse a received BarrageMessage: ").append(e).endl();
        throw new GrpcMarshallingException("Unable to parse BarrageMessage object", e);
    }
}