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);
}
}
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();
}
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();
});
});
}
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);
}
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);
}
Aggregations