Search in sources :

Example 11 with FlatBufferBuilder

use of com.google.flatbuffers.FlatBufferBuilder in project deephaven-core by deephaven.

the class BarrageStreamGenerator method getInputStream.

/**
 * Returns an InputStream of the message filtered to the viewport.
 *
 * @param view the view of the overall chunk to generate a RecordBatch for
 * @param metadata the optional flight data metadata to attach to the message
 * @param columnVisitor the helper method responsible for appending the payload columns to the RecordBatch
 * @return an InputStream ready to be drained by GRPC
 */
private InputStream getInputStream(final View view, final ByteBuffer metadata, final ColumnVisitor columnVisitor) throws IOException {
    final ArrayDeque<InputStream> streams = new ArrayDeque<>();
    final MutableInt size = new MutableInt();
    final Consumer<InputStream> addStream = (final InputStream is) -> {
        try {
            final int sz = is.available();
            if (sz == 0) {
                is.close();
                return;
            }
            streams.add(is);
            size.add(sz);
        } catch (final IOException e) {
            throw new UncheckedDeephavenException("Unexpected IOException", e);
        }
        // C++.
        if (size.intValue() % 8 != 0) {
            final int paddingBytes = (8 - (size.intValue() % 8));
            size.add(paddingBytes);
            streams.add(new DrainableByteArrayInputStream(PADDING_BUFFER, 0, paddingBytes));
        }
    };
    final FlatBufferBuilder header = new FlatBufferBuilder();
    final long numRows;
    final int nodesOffset;
    final int buffersOffset;
    try (final SizedChunk<Values> nodeOffsets = new SizedChunk<>(ChunkType.Object);
        final SizedLongChunk<Values> bufferInfos = new SizedLongChunk<>()) {
        nodeOffsets.ensureCapacity(addColumnData.length);
        nodeOffsets.get().setSize(0);
        bufferInfos.ensureCapacity(addColumnData.length * 3);
        bufferInfos.get().setSize(0);
        final MutableLong totalBufferLength = new MutableLong();
        final ChunkInputStreamGenerator.FieldNodeListener fieldNodeListener = (numElements, nullCount) -> {
            nodeOffsets.ensureCapacityPreserve(nodeOffsets.get().size() + 1);
            nodeOffsets.get().asWritableObjectChunk().add(new ChunkInputStreamGenerator.FieldNodeInfo(numElements, nullCount));
        };
        final ChunkInputStreamGenerator.BufferListener bufferListener = (length) -> {
            totalBufferLength.add(length);
            bufferInfos.ensureCapacityPreserve(bufferInfos.get().size() + 1);
            bufferInfos.get().add(length);
        };
        numRows = columnVisitor.visit(view, addStream, fieldNodeListener, bufferListener);
        final WritableChunk<Values> noChunk = nodeOffsets.get();
        RecordBatch.startNodesVector(header, noChunk.size());
        for (int i = noChunk.size() - 1; i >= 0; --i) {
            final ChunkInputStreamGenerator.FieldNodeInfo node = (ChunkInputStreamGenerator.FieldNodeInfo) noChunk.asObjectChunk().get(i);
            FieldNode.createFieldNode(header, node.numElements, node.nullCount);
        }
        nodesOffset = header.endVector();
        final WritableLongChunk<Values> biChunk = bufferInfos.get();
        RecordBatch.startBuffersVector(header, biChunk.size());
        for (int i = biChunk.size() - 1; i >= 0; --i) {
            totalBufferLength.subtract(biChunk.get(i));
            Buffer.createBuffer(header, totalBufferLength.longValue(), biChunk.get(i));
        }
        buffersOffset = header.endVector();
    }
    RecordBatch.startRecordBatch(header);
    RecordBatch.addNodes(header, nodesOffset);
    RecordBatch.addBuffers(header, buffersOffset);
    RecordBatch.addLength(header, numRows);
    final int headerOffset = RecordBatch.endRecordBatch(header);
    header.finish(MessageHelper.wrapInMessage(header, headerOffset, org.apache.arrow.flatbuf.MessageHeader.RecordBatch, size.intValue()));
    // now create the proto header
    try (final ExposedByteArrayOutputStream baos = new ExposedByteArrayOutputStream()) {
        final CodedOutputStream cos = CodedOutputStream.newInstance(baos);
        cos.writeByteBuffer(Flight.FlightData.DATA_HEADER_FIELD_NUMBER, header.dataBuffer().slice());
        if (metadata != null) {
            cos.writeByteBuffer(Flight.FlightData.APP_METADATA_FIELD_NUMBER, metadata);
        }
        cos.writeTag(Flight.FlightData.DATA_BODY_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED);
        cos.writeUInt32NoTag(size.intValue());
        cos.flush();
        streams.addFirst(new DrainableByteArrayInputStream(baos.peekBuffer(), 0, baos.size()));
        return new ConsecutiveDrainableStreams(streams.toArray(new InputStream[0]));
    } catch (final IOException ex) {
        throw new UncheckedDeephavenException("Unexpected IOException", ex);
    }
}
Also used : MutableInt(org.apache.commons.lang3.mutable.MutableInt) TIntArrayList(gnu.trove.list.array.TIntArrayList) ExternalizableRowSetUtils(io.deephaven.engine.rowset.impl.ExternalizableRowSetUtils) PADDING_BUFFER(io.deephaven.extensions.barrage.chunk.BaseChunkInputStreamGenerator.PADDING_BUFFER) ByteStringAccess(com.google.protobuf.ByteStringAccess) SizedChunk(io.deephaven.chunk.sized.SizedChunk) ByteBuffer(java.nio.ByteBuffer) SafeCloseable(io.deephaven.util.SafeCloseable) MutableLong(org.apache.commons.lang3.mutable.MutableLong) Map(java.util.Map) RowSetBuilderSequential(io.deephaven.engine.rowset.RowSetBuilderSequential) Values(io.deephaven.chunk.attributes.Values) LoggerFactory(io.deephaven.internal.log.LoggerFactory) ExposedByteArrayOutputStream(io.deephaven.extensions.barrage.util.BarrageProtoUtil.ExposedByteArrayOutputStream) Objects(java.util.Objects) Nullable(org.jetbrains.annotations.Nullable) RecordBatch(org.apache.arrow.flatbuf.RecordBatch) BarrageSubscriptionOptions(io.deephaven.extensions.barrage.BarrageSubscriptionOptions) LittleEndianDataOutputStream(com.google.common.io.LittleEndianDataOutputStream) BarrageUtil(io.deephaven.extensions.barrage.util.BarrageUtil) MessageHelper(io.deephaven.proto.flight.util.MessageHelper) ChunkInputStreamGenerator(io.deephaven.extensions.barrage.chunk.ChunkInputStreamGenerator) FlatBufferBuilder(com.google.flatbuffers.FlatBufferBuilder) ChunkType(io.deephaven.chunk.ChunkType) Flight(org.apache.arrow.flight.impl.Flight) io.deephaven.barrage.flatbuf(io.deephaven.barrage.flatbuf) WritableRowSet(io.deephaven.engine.rowset.WritableRowSet) Drainable(io.grpc.Drainable) Singleton(javax.inject.Singleton) RowSet(io.deephaven.engine.rowset.RowSet) UncheckedDeephavenException(io.deephaven.UncheckedDeephavenException) RowSetFactory(io.deephaven.engine.rowset.RowSetFactory) Inject(javax.inject.Inject) BarrageMessage(io.deephaven.engine.table.impl.util.BarrageMessage) SizedLongChunk(io.deephaven.chunk.sized.SizedLongChunk) WritableLongChunk(io.deephaven.chunk.WritableLongChunk) StreamReaderOptions(io.deephaven.extensions.barrage.util.StreamReaderOptions) WireFormat(com.google.protobuf.WireFormat) CodedOutputStream(com.google.protobuf.CodedOutputStream) FieldNode(org.apache.arrow.flatbuf.FieldNode) OutputStream(java.io.OutputStream) RowSetShiftData(io.deephaven.engine.rowset.RowSetShiftData) IOException(java.io.IOException) WritableChunk(io.deephaven.chunk.WritableChunk) DefensiveDrainable(io.deephaven.extensions.barrage.util.DefensiveDrainable) Consumer(java.util.function.Consumer) BarrageSnapshotOptions(io.deephaven.extensions.barrage.BarrageSnapshotOptions) Buffer(org.apache.arrow.flatbuf.Buffer) Logger(io.deephaven.io.logger.Logger) TableDefinition(io.deephaven.engine.table.TableDefinition) BitSet(java.util.BitSet) ArrayDeque(java.util.ArrayDeque) InputStream(java.io.InputStream) SizedChunk(io.deephaven.chunk.sized.SizedChunk) ExposedByteArrayOutputStream(io.deephaven.extensions.barrage.util.BarrageProtoUtil.ExposedByteArrayOutputStream) Values(io.deephaven.chunk.attributes.Values) SizedLongChunk(io.deephaven.chunk.sized.SizedLongChunk) FlatBufferBuilder(com.google.flatbuffers.FlatBufferBuilder) InputStream(java.io.InputStream) CodedOutputStream(com.google.protobuf.CodedOutputStream) IOException(java.io.IOException) ArrayDeque(java.util.ArrayDeque) MutableLong(org.apache.commons.lang3.mutable.MutableLong) UncheckedDeephavenException(io.deephaven.UncheckedDeephavenException) MutableInt(org.apache.commons.lang3.mutable.MutableInt) ChunkInputStreamGenerator(io.deephaven.extensions.barrage.chunk.ChunkInputStreamGenerator)

Example 12 with FlatBufferBuilder

use of com.google.flatbuffers.FlatBufferBuilder in project deephaven-core by deephaven.

the class BarrageStreamGenerator method getSubscriptionMetadata.

private ByteBuffer getSubscriptionMetadata(final SubView view) throws IOException {
    final FlatBufferBuilder metadata = new FlatBufferBuilder();
    int effectiveViewportOffset = 0;
    if (isSnapshot && view.isViewport()) {
        try (final RowSetGenerator viewportGen = new RowSetGenerator(view.viewport)) {
            effectiveViewportOffset = viewportGen.addToFlatBuffer(metadata);
        }
    }
    int effectiveColumnSetOffset = 0;
    if (isSnapshot && view.subscribedColumns != null) {
        effectiveColumnSetOffset = new BitSetGenerator(view.subscribedColumns).addToFlatBuffer(metadata);
    }
    final int rowsAddedOffset;
    if (isSnapshot && !view.isInitialSnapshot) {
        // client's don't need/want to receive the full RowSet on every snapshot
        rowsAddedOffset = EmptyRowSetGenerator.INSTANCE.addToFlatBuffer(metadata);
    } else {
        rowsAddedOffset = rowsAdded.addToFlatBuffer(metadata);
    }
    final int rowsRemovedOffset = rowsRemoved.addToFlatBuffer(metadata);
    final int shiftDataOffset = shifted.addToFlatBuffer(metadata);
    // Added Chunk Data:
    int addedRowsIncludedOffset = 0;
    if (view.isViewport()) {
        addedRowsIncludedOffset = rowsIncluded.addToFlatBuffer(view.keyspaceViewport, metadata);
    }
    // now add mod-column streams, and write the mod column indexes
    TIntArrayList modOffsets = new TIntArrayList(modColumnData.length);
    for (final ModColumnData mcd : modColumnData) {
        final int myModRowOffset;
        if (view.isViewport()) {
            myModRowOffset = mcd.rowsModified.addToFlatBuffer(view.keyspaceViewport, metadata);
        } else {
            myModRowOffset = mcd.rowsModified.addToFlatBuffer(metadata);
        }
        modOffsets.add(BarrageModColumnMetadata.createBarrageModColumnMetadata(metadata, myModRowOffset));
    }
    BarrageUpdateMetadata.startModColumnNodesVector(metadata, modOffsets.size());
    modOffsets.forEachDescending(offset -> {
        metadata.addOffset(offset);
        return true;
    });
    final int nodesOffset = metadata.endVector();
    BarrageUpdateMetadata.startBarrageUpdateMetadata(metadata);
    BarrageUpdateMetadata.addNumAddBatches(metadata, view.hasAddBatch ? 1 : 0);
    BarrageUpdateMetadata.addNumModBatches(metadata, view.hasModBatch ? 1 : 0);
    BarrageUpdateMetadata.addIsSnapshot(metadata, isSnapshot);
    BarrageUpdateMetadata.addFirstSeq(metadata, firstSeq);
    BarrageUpdateMetadata.addLastSeq(metadata, lastSeq);
    BarrageUpdateMetadata.addEffectiveViewport(metadata, effectiveViewportOffset);
    BarrageUpdateMetadata.addEffectiveColumnSet(metadata, effectiveColumnSetOffset);
    BarrageUpdateMetadata.addEffectiveReverseViewport(metadata, view.reverseViewport);
    BarrageUpdateMetadata.addAddedRows(metadata, rowsAddedOffset);
    BarrageUpdateMetadata.addRemovedRows(metadata, rowsRemovedOffset);
    BarrageUpdateMetadata.addShiftData(metadata, shiftDataOffset);
    BarrageUpdateMetadata.addAddedRowsIncluded(metadata, addedRowsIncludedOffset);
    BarrageUpdateMetadata.addModColumnNodes(metadata, nodesOffset);
    metadata.finish(BarrageUpdateMetadata.endBarrageUpdateMetadata(metadata));
    final FlatBufferBuilder header = new FlatBufferBuilder();
    final int payloadOffset = BarrageMessageWrapper.createMsgPayloadVector(header, metadata.dataBuffer());
    BarrageMessageWrapper.startBarrageMessageWrapper(header);
    BarrageMessageWrapper.addMagic(header, BarrageUtil.FLATBUFFER_MAGIC);
    BarrageMessageWrapper.addMsgType(header, BarrageMessageType.BarrageUpdateMetadata);
    BarrageMessageWrapper.addMsgPayload(header, payloadOffset);
    header.finish(BarrageMessageWrapper.endBarrageMessageWrapper(header));
    return header.dataBuffer().slice();
}
Also used : FlatBufferBuilder(com.google.flatbuffers.FlatBufferBuilder) TIntArrayList(gnu.trove.list.array.TIntArrayList)

Example 13 with FlatBufferBuilder

use of com.google.flatbuffers.FlatBufferBuilder in project deephaven-core by deephaven.

the class FlightServiceGrpcImpl method doGetCustom.

public void doGetCustom(final Flight.Ticket request, final StreamObserver<InputStream> responseObserver) {
    GrpcUtil.rpcWrapper(log, responseObserver, () -> {
        final SessionState session = sessionService.getCurrentSession();
        final SessionState.ExportObject<BaseTable> export = ticketRouter.resolve(session, request, "request");
        session.nonExport().require(export).onError(responseObserver).submit(() -> {
            final BaseTable table = export.get();
            // create an adapter for the response observer
            final StreamObserver<BarrageStreamGenerator.View> listener = ArrowModule.provideListenerAdapter().adapt(responseObserver);
            // Send Schema wrapped in Message
            final FlatBufferBuilder builder = new FlatBufferBuilder();
            final int schemaOffset = BarrageUtil.makeSchemaPayload(builder, table.getDefinition(), table.getAttributes());
            builder.finish(MessageHelper.wrapInMessage(builder, schemaOffset, org.apache.arrow.flatbuf.MessageHeader.Schema));
            final ByteBuffer serializedMessage = builder.dataBuffer();
            // leverage the stream generator SchemaView constructor
            final BarrageStreamGenerator.SchemaView schemaView = new BarrageStreamGenerator.SchemaView(serializedMessage);
            // push the schema to the listener
            listener.onNext(schemaView);
            // get ourselves some data!
            final BarrageMessage msg = ConstructSnapshot.constructBackplaneSnapshot(this, table);
            // actually no mod column data for DoGet
            msg.modColumnData = ZERO_MOD_COLUMNS;
            // translate the viewport to keyspace and make the call
            try (final BarrageStreamGenerator bsg = new BarrageStreamGenerator(msg)) {
                listener.onNext(bsg.getSnapshotView(DEFAULT_SNAPSHOT_DESER_OPTIONS));
            }
            listener.onCompleted();
        });
    });
}
Also used : SessionState(io.deephaven.server.session.SessionState) BaseTable(io.deephaven.engine.table.impl.BaseTable) BarrageStreamGenerator(io.deephaven.server.barrage.BarrageStreamGenerator) FlatBufferBuilder(com.google.flatbuffers.FlatBufferBuilder) BarrageMessage(io.deephaven.engine.table.impl.util.BarrageMessage) ByteBuffer(java.nio.ByteBuffer)

Example 14 with FlatBufferBuilder

use of com.google.flatbuffers.FlatBufferBuilder in project apisix-java-plugin-runner by apache.

the class A6ConfigHandlerTest method testAddFilter3.

@Test
@DisplayName("test skip the same name filter")
void testAddFilter3() {
    FlatBufferBuilder builder = new FlatBufferBuilder();
    int foo1 = builder.createString("FooFilter");
    int bar1 = builder.createString("Bar1");
    int filter1 = TextEntry.createTextEntry(builder, foo1, bar1);
    int foo2 = builder.createString("FooFilter");
    int bar2 = builder.createString("Bar2");
    int filter2 = TextEntry.createTextEntry(builder, foo2, bar2);
    int confVector = Req.createConfVector(builder, new int[] { filter1, filter2 });
    Req.startReq(builder);
    Req.addConf(builder, confVector);
    builder.finish(Req.endReq(builder));
    Req req = Req.getRootAsReq(builder.dataBuffer());
    A6ConfigRequest request = new A6ConfigRequest(req);
    EmbeddedChannel channel = new EmbeddedChannel(new BinaryProtocolDecoder(), prepareConfHandler);
    channel.writeInbound(request);
    channel.finish();
    A6ConfigResponse response = channel.readOutbound();
    A6Conf config = cache.getIfPresent(response.getConfToken());
    Assertions.assertEquals(config.getChain().getFilters().size(), 1);
}
Also used : A6ConfigRequest(org.apache.apisix.plugin.runner.A6ConfigRequest) A6ConfigResponse(org.apache.apisix.plugin.runner.A6ConfigResponse) FlatBufferBuilder(com.google.flatbuffers.FlatBufferBuilder) EmbeddedChannel(io.netty.channel.embedded.EmbeddedChannel) A6Conf(org.apache.apisix.plugin.runner.A6Conf) Req(io.github.api7.A6.PrepareConf.Req) Test(org.junit.jupiter.api.Test) DisplayName(org.junit.jupiter.api.DisplayName)

Example 15 with FlatBufferBuilder

use of com.google.flatbuffers.FlatBufferBuilder in project apisix-java-plugin-runner by apache.

the class A6ConfigHandlerTest method testAddFilter2.

@Test
@DisplayName("test filter sort by it's order")
void testAddFilter2() {
    FlatBufferBuilder builder = new FlatBufferBuilder();
    int cat = builder.createString("CatFilter");
    int dog = builder.createString("Dog");
    int filter2 = TextEntry.createTextEntry(builder, cat, dog);
    int foo = builder.createString("FooFilter");
    int bar = builder.createString("Bar");
    int filter1 = TextEntry.createTextEntry(builder, foo, bar);
    int confVector = Req.createConfVector(builder, new int[] { filter1, filter2 });
    Req.startReq(builder);
    Req.addConf(builder, confVector);
    builder.finish(Req.endReq(builder));
    Req req = Req.getRootAsReq(builder.dataBuffer());
    A6ConfigRequest request = new A6ConfigRequest(req);
    EmbeddedChannel channel = new EmbeddedChannel(new BinaryProtocolDecoder(), prepareConfHandler);
    channel.writeInbound(request);
    channel.finish();
    A6ConfigResponse response = channel.readOutbound();
    A6Conf config = cache.getIfPresent(response.getConfToken());
    Assertions.assertEquals(config.getChain().getFilters().size(), 2);
}
Also used : A6ConfigRequest(org.apache.apisix.plugin.runner.A6ConfigRequest) A6ConfigResponse(org.apache.apisix.plugin.runner.A6ConfigResponse) FlatBufferBuilder(com.google.flatbuffers.FlatBufferBuilder) EmbeddedChannel(io.netty.channel.embedded.EmbeddedChannel) A6Conf(org.apache.apisix.plugin.runner.A6Conf) Req(io.github.api7.A6.PrepareConf.Req) Test(org.junit.jupiter.api.Test) DisplayName(org.junit.jupiter.api.DisplayName)

Aggregations

FlatBufferBuilder (com.google.flatbuffers.FlatBufferBuilder)47 DisplayName (org.junit.jupiter.api.DisplayName)14 Test (org.junit.jupiter.api.Test)14 BigInteger (java.math.BigInteger)10 A6ConfigRequest (org.apache.apisix.plugin.runner.A6ConfigRequest)9 HttpRequest (org.apache.apisix.plugin.runner.HttpRequest)9 EmbeddedChannel (io.netty.channel.embedded.EmbeddedChannel)7 A6ConfigResponse (org.apache.apisix.plugin.runner.A6ConfigResponse)7 Test (org.junit.Test)7 ByteBuffer (java.nio.ByteBuffer)6 Req (io.github.api7.A6.PrepareConf.Req)5 A6Conf (org.apache.apisix.plugin.runner.A6Conf)5 HashMap (java.util.HashMap)3 List (java.util.List)3 Map (java.util.Map)3 lombok.val (lombok.val)3 ExtraInfoRequest (org.apache.apisix.plugin.runner.ExtraInfoRequest)3 HttpResponse (org.apache.apisix.plugin.runner.HttpResponse)3 BaseNd4jTest (org.nd4j.linalg.BaseNd4jTest)3 ByteStringAccess (com.google.protobuf.ByteStringAccess)2