Java Code Examples for com.google.flatbuffers.FlatBufferBuilder#finish()

public FragmentWritableBatch(
    final QueryId queryId,
    final int sendMajorFragmentId,
    final int sendMinorFragmentId,
    final int receiveMajorFragmentId,
    ArrowRecordBatch recordBatch,
    final int... receiveMinorFragmentId){
  this.buffers = recordBatch.getBuffers().stream().map(buf -> buf.asNettyBuffer()).collect
    (Collectors.toList()).toArray(new ByteBuf[0]);
  this.recordCount = recordBatch.getLength();
  FlatBufferBuilder fbbuilder = new FlatBufferBuilder();
  fbbuilder.finish(recordBatch.writeTo(fbbuilder));
  ByteBuffer arrowRecordBatch = fbbuilder.dataBuffer();
  final FragmentRecordBatch.Builder builder = FragmentRecordBatch.newBuilder()
      .setArrowRecordBatch(ByteString.copyFrom(arrowRecordBatch))
      .setQueryId(queryId)
      .setReceivingMajorFragmentId(receiveMajorFragmentId)
      .setSendingMajorFragmentId(sendMajorFragmentId)
      .setSendingMinorFragmentId(sendMinorFragmentId);

  for(final int i : receiveMinorFragmentId){
    builder.addReceivingMinorFragmentId(i);
  }

  this.header = builder.build();
}
 

/**
 * Register the event name - "accuracy", "loss", etc for later use in recording events.
 * @param name Name to register
 * @return Number of bytes written
 */
public long registerEventName(String name) throws IOException {
    Preconditions.checkState(endStaticInfoOffset >= 0, "Cannot write name - have not written end of static info marker");

    FlatBufferBuilder fbb = new FlatBufferBuilder(0);
    long time = System.currentTimeMillis();
    int offset = UIEvent.createUIEvent(fbb, UIEventType.ADD_NAME, UIEventSubtype.NONE, -1, time, 0, 0, (short)-1, 0, 0);
    fbb.finish(offset);

    FlatBufferBuilder fbb2 = new FlatBufferBuilder(0);
    int idx = nameIndexCounter.getAndIncrement();
    nameIndexMap.put(idx, name);
    indexNameMap.put(name, idx);
    int strOffset = fbb2.createString(name);
    int offset2 = UIAddName.createUIAddName(fbb2, idx, strOffset);
    fbb2.finish(offset2);

    long l = append(fbb, fbb2);
    return l;
}
 

public long writeHistogramEventCustomBins(String name, EventSubtype subtype, long time, int iteration, int epoch, INDArray bins, INDArray y) throws IOException {
    Preconditions.checkState(y.rank() == 1, "Y array must be rank 1, got Y array with shape %ndShape", y);
    Preconditions.checkState(bins.rank() == 2, "Bins array must have shape [2,numBins], got bins array with shape %ndShape", bins);
    Preconditions.checkState(y.length() == bins.size(1), "Bins array must have shape [2,numBins], where numBins must match y.length()=%s, got bins array with shape %ndShape", y.length(), bins);

    //TODO add support for plugin, variable and frame/iter
    //TODO: Code duplication for histogram methods...
    Preconditions.checkState(indexNameMap.containsKey(name), "Name \"%s\" not yet registered", name);
    int idx = indexNameMap.get(name);

    FlatBufferBuilder fbb = new FlatBufferBuilder(0);
    int offset = UIEvent.createUIEvent(fbb, UIEventType.HISTOGRAM, subtype.asUIEventSubtype(), idx, time, iteration, epoch, (short)-1, 0, 0);
    fbb.finish(offset);

    FlatBufferBuilder fbb2 = new FlatBufferBuilder(0);
    int yOffset = y.toFlatArray(fbb2);

    int binRangesOffset = bins.toFlatArray(fbb2);

    int offset2 = UIHistogram.createUIHistogram(fbb2, UIHistogramType.CUSTOM, y.length(), binRangesOffset, yOffset, 0);
    fbb2.finish(offset2);

    return append(fbb, fbb2);
}
 

/**
 * Converting old Arrow Schema to new one based on Arrow version used
 * in Dremio as of 2.1.0
 * @param oldSchema
 * @return
 */
@VisibleForTesting
byte[] convertFromOldSchema(OldSchema oldSchema) {
  FlatBufferBuilder builder = new FlatBufferBuilder();
  int[] fieldOffsets = new int[oldSchema.fieldsLength()];
  for (int i = 0; i < oldSchema.fieldsLength(); i++) {
    fieldOffsets[i] = convertFromOldField(oldSchema.fields(i), builder);
  }
  int fieldsOffset = org.apache.arrow.flatbuf.Schema.createFieldsVector(builder, fieldOffsets);
  int[] metadataOffsets = new int[oldSchema.customMetadataLength()];
  for (int i = 0; i < metadataOffsets.length; i++) {
    int keyOffset = builder.createString(oldSchema.customMetadata(i).key());
    int valueOffset = builder.createString(oldSchema.customMetadata(i).value());
    KeyValue.startKeyValue(builder);
    KeyValue.addKey(builder, keyOffset);
    KeyValue.addValue(builder, valueOffset);
    metadataOffsets[i] = KeyValue.endKeyValue(builder);
  }
  int metadataOffset = org.apache.arrow.flatbuf.Field.createCustomMetadataVector(builder, metadataOffsets);
  org.apache.arrow.flatbuf.Schema.startSchema(builder);
  org.apache.arrow.flatbuf.Schema.addFields(builder, fieldsOffset);
  org.apache.arrow.flatbuf.Schema.addCustomMetadata(builder, metadataOffset);
  builder.finish(org.apache.arrow.flatbuf.Schema.endSchema(builder));
  return builder.sizedByteArray();
}
 

private boolean QueryRealms(Socket socket) throws IOException {
  FlatBufferBuilder builder = new FlatBufferBuilder();

  int[] realms = new int[REALMS.length];
  for (int i = 0; i < REALMS.length; i++) {
    realms[i] = Realm.createRealm(builder, builder.createString(REALMS[i][0]), builder.createString(REALMS[i][1]));
  }
  int realmsVec = QueryRealms.createRealmsVector(builder, realms);

  QueryRealms.startQueryRealms(builder);
  QueryRealms.addRealms(builder, realmsVec);
  int realmId = QueryRealms.endQueryRealms(builder);

  int id = com.riiablo.net.packet.bnls.BNLS.createBNLS(builder, BNLSData.QueryRealms, realmId);

  builder.finish(id);

  ByteBuffer data = builder.dataBuffer();

  OutputStream out = socket.getOutputStream();
  WritableByteChannel channel = Channels.newChannel(out);
  channel.write(data);
  Gdx.app.log(TAG, "returning realms list...");
  return false;
}
 

private boolean StartInstance(Socket socket, com.riiablo.net.packet.msi.MSI packet) throws IOException {
  Gdx.app.debug(TAG, "Starting instance...");

  try {
    File outFile = new File("D2GS.tmp");
    ProcessBuilder processBuilder = new ProcessBuilder("java", "-jar", "server/d2gs/build/libs/d2gs-1.0.jar");
    processBuilder.redirectOutput(ProcessBuilder.Redirect.to(outFile));
    processBuilder.redirectError(ProcessBuilder.Redirect.to(outFile));
    Process process = processBuilder.start();
    instances.add(process);
  } catch (Throwable t) {
    Gdx.app.error(TAG, t.getMessage(), t);
  }

  int ip     = 2130706433; // 127.0.0.1
  short port = 6114;

  FlatBufferBuilder builder = new FlatBufferBuilder();
  StartInstance.startStartInstance(builder);
  StartInstance.addResult(builder, Result.SUCCESS);
  StartInstance.addIp(builder, ip);
  StartInstance.addPort(builder, port);
  int startInstanceOffset = StartInstance.endStartInstance(builder);
  int id = com.riiablo.net.packet.msi.MSI.createMSI(builder, MSIData.StartInstance, startInstanceOffset);
  builder.finish(id);

  ByteBuffer data = builder.dataBuffer();
  OutputStream out = socket.getOutputStream();
  WritableByteChannel channel = Channels.newChannel(out);
  channel.write(data);
  Gdx.app.debug(TAG, "Returning instance at " + InetAddress.getByAddress(ByteBuffer.allocate(4).putInt(ip).array()));
  return true;
}
 

/**
 * Encode the header as a UIStaticInfoRecord instance for the specific {@link UIEventType}
 * @param type UIEventType
 */
protected Pair<Integer, FlatBufferBuilder> encodeStaticHeader(byte type) {
    FlatBufferBuilder fbb = new FlatBufferBuilder(12);

    int staticInfoOffset = UIStaticInfoRecord.createUIStaticInfoRecord(fbb, type);
    fbb.finish(staticInfoOffset);
    int lengthHeader = fbb.offset();       //MUST be called post finish to get real length
    return new Pair<>(lengthHeader, fbb);
}
 

private void CreateGame(GameSession session, ResponseListener listener) {
  Gdx.app.debug(TAG, "Creating game " + session);
  FlatBufferBuilder builder = new FlatBufferBuilder();
  int gameNameOffset = builder.createString(session.name);
  int descriptionOffset = builder.createString(session.desc);
  CreateGame.startCreateGame(builder);
  CreateGame.addGameName(builder, gameNameOffset);
  CreateGame.addDescription(builder, descriptionOffset);
  int createGameOffset = CreateGame.endCreateGame(builder);
  int id = MCP.createMCP(builder, MCPData.CreateGame, createGameOffset);
  builder.finish(id);
  ByteBuffer data = builder.dataBuffer();
  connection.sendRequest(data, listener);
}
 

private void ListGames(ResponseListener listener) {
  Gdx.app.debug(TAG, "Requesting games list");
  FlatBufferBuilder builder = new FlatBufferBuilder();
  ListGames.startListGames(builder);
  int listGamesOffset = ListGames.endListGames(builder);
  int id = MCP.createMCP(builder, MCPData.ListGames, listGamesOffset);
  builder.finish(id);
  ByteBuffer data = builder.dataBuffer();
  connection.sendRequest(data, listener);
}
 

private boolean ConnectionAccepted(Socket socket) throws IOException {
  Gdx.app.debug(TAG, "Connection accepted!");
  FlatBufferBuilder builder = new FlatBufferBuilder();
  ConnectionAccepted.startConnectionAccepted(builder);
  int connectionAcceptedId = ConnectionAccepted.endConnectionAccepted(builder);
  int id = com.riiablo.net.packet.bnls.BNLS.createBNLS(builder, BNLSData.ConnectionAccepted, connectionAcceptedId);
  builder.finish(id);

  ByteBuffer data = builder.dataBuffer();
  OutputStream out = socket.getOutputStream();
  WritableByteChannel channel = Channels.newChannel(out);
  channel.write(data);
  return false;
}
 

public long writeHistogramEventDiscrete(@NonNull String name, EventSubtype subtype, long time, int iteration, int epoch, List<String> binLabels, @NonNull INDArray y) throws IOException {
    Preconditions.checkState(binLabels == null || binLabels.size() == y.length(), "Number of bin labels (if present) must " +
            "be same as Y array length - got %s bins, array shape %ndShape", (binLabels == null ? 0L : binLabels.size()), y.length());
    Preconditions.checkState(y.rank() == 1, "Y array must be rank 1, got Y array with shape %ndShape", y);

    //TODO add support for plugin, variable and frame/iter
    Preconditions.checkState(indexNameMap.containsKey(name), "Name \"%s\" not yet registered", name);
    int idx = indexNameMap.get(name);

    FlatBufferBuilder fbb = new FlatBufferBuilder(0);
    int offset = UIEvent.createUIEvent(fbb, UIEventType.HISTOGRAM, subtype.asUIEventSubtype(), idx, time, iteration, epoch, (short)-1, 0, 0);
    fbb.finish(offset);

    FlatBufferBuilder fbb2 = new FlatBufferBuilder(0);
    int yOffset = y.toFlatArray(fbb2);
    int binLabelsOffset = 0;
    if(binLabels != null){
        int[] str = new int[binLabels.size()];
        for( int i=0; i<binLabels.size(); i++ ){
            String s = binLabels.get(i);
            if(s == null)
                s = "";
            str[i] = fbb2.createString(s);
        }
        binLabelsOffset = UIHistogram.createBinlabelsVector(fbb2, str);
    }
    int offset2 = UIHistogram.createUIHistogram(fbb2, UIHistogramType.DISCRETE, y.length(), 0, yOffset, binLabelsOffset);
    fbb2.finish(offset2);

    return append(fbb, fbb2);
}
 

private void StartInstance(CreateGame createGame, ResponseListener listener) {
  Gdx.app.debug(TAG, "Requesting game instance for " + createGame);
  FlatBufferBuilder builder = new FlatBufferBuilder();
  StartInstance.startStartInstance(builder);
  int startInstanceOffset = StartInstance.endStartInstance(builder);
  int id = MSI.createMSI(builder, MSIData.StartInstance, startInstanceOffset);
  builder.finish(id);
  ByteBuffer data = builder.dataBuffer();

  Socket socket = null;
  try {
    socket = Gdx.net.newClientSocket(Net.Protocol.TCP, "localhost", com.riiablo.server.mcp.MSI.PORT, null);

    OutputStream out = socket.getOutputStream();
    WritableByteChannel channelOut = Channels.newChannel(out);
    channelOut.write(data);

    buffer.clear();
    buffer.mark();
    InputStream in = socket.getInputStream();
    ReadableByteChannel channelIn = Channels.newChannel(in);
    channelIn.read(buffer);
    buffer.limit(buffer.position());
    buffer.reset();

    MSI packet = MSI.getRootAsMSI(buffer);
    Gdx.app.log(TAG, "packet type " + MCPData.name(packet.dataType()));
    listener.handleResponse(packet);
  } catch (Throwable t) {
    listener.failed(t);
  } finally {
    if (socket != null) socket.dispose();
  }
}
 

private boolean ConnectionAccepted(Socket socket) throws IOException {
  Gdx.app.debug(TAG, "Connection accepted!");
  FlatBufferBuilder builder = new FlatBufferBuilder();
  ConnectionAccepted.startConnectionAccepted(builder);
  int connectionAcceptedId = ConnectionAccepted.endConnectionAccepted(builder);
  int id = com.riiablo.net.packet.mcp.MCP.createMCP(builder, MCPData.ConnectionAccepted, connectionAcceptedId);
  builder.finish(id);

  ByteBuffer data = builder.dataBuffer();
  OutputStream out = socket.getOutputStream();
  WritableByteChannel channel = Channels.newChannel(out);
  channel.write(data);
  return false;
}
 

/**
 * Write a single scalar event to the file
 * @param name      Name of the event. Must be registered by {@link #registerEventName(String)} first!
 * @param time      Timestamp
 * @param iteration Iteration for the event
 * @param epoch     Epoch for the event
 * @param scalar    Scalar value to write
 * @return          Number of bytes written
 */
public long writeScalarEvent(String name, EventSubtype subtype, long time, int iteration, int epoch, Number scalar) throws IOException {
    //TODO add support for plugin, variable and frame/iter
    Preconditions.checkState(indexNameMap.containsKey(name), "Name \"%s\" not yet registered", name);
    int idx = indexNameMap.get(name);
    FlatBufferBuilder fbb = new FlatBufferBuilder(0);
    int offset = UIEvent.createUIEvent(fbb, UIEventType.SCALAR, subtype.asUIEventSubtype(), idx, time, iteration, epoch, (short)-1, 0, 0);
    fbb.finish(offset);

    FlatBufferBuilder fbb2 = new FlatBufferBuilder(0);
    int offset2 = Nd4j.scalar(scalar).toFlatArray(fbb2);
    fbb2.finish(offset2);

    return append(fbb, fbb2);
}
 

@Test
public void testDataSetSchema() throws Exception {
  Field field1 = new Field("a", true, new Int(32, true), null);
  Field child1 = new Field("c", true, Utf8.INSTANCE, null);
  Field field2 = new Field("b", true, Struct.INSTANCE, ImmutableList.of(child1));
  Schema schema = new Schema(ImmutableList.of(field1, field2));
  FlatBufferBuilder builder = new FlatBufferBuilder();
  schema.getSchema(builder);
  builder.finish(schema.getSchema(builder));
  NamespaceTestUtils.addSource(namespaceService, "s");
  NamespaceTestUtils.addPhysicalDS(namespaceService, "s.foo", builder.sizedByteArray());
  ByteBuffer bb = ByteBuffer.wrap(DatasetHelper.getSchemaBytes(namespaceService.getDataset(new NamespaceKey(PathUtils.parseFullPath("s.foo")))).toByteArray());
  Schema returnedSchema = Schema.convertSchema(org.apache.arrow.flatbuf.Schema.getRootAsSchema(bb));
  assertEquals(schema, returnedSchema);
}
 

@Before
public void setUp() throws Exception {
  kvStoreProvider = new LocalKVStoreProvider(DremioTest.CLASSPATH_SCAN_RESULT, null, true, false).asLegacy();
  kvStoreProvider.start();

  namespaceService = new NamespaceServiceImpl(kvStoreProvider);

  NamespaceTestUtils.addSource(namespaceService, SOURCE);
  NamespaceTestUtils.addSpace(namespaceService, EMPTYSPACE);
  NamespaceTestUtils.addSpace(namespaceService, SPACE);
  NamespaceTestUtils.addFolder(namespaceService, FOLDER);
  NamespaceTestUtils.addFolder(namespaceService, SUBFOLDER);

  final Field field1 = CompleteType.INT.toField("a");
  final Field child1 = CompleteType.VARCHAR.toField("c");
  final Field field2 = CompleteType.struct(ImmutableList.of(child1)).toField("b");
  final org.apache.arrow.vector.types.pojo.Schema schema =
    new org.apache.arrow.vector.types.pojo.Schema(ImmutableList.of(field1, field2));
  final FlatBufferBuilder builder = new FlatBufferBuilder();
  schema.getSchema(builder);
  builder.finish(schema.getSchema(builder));
  NamespaceTestUtils.addPhysicalDS(namespaceService, PDS, builder.sizedByteArray());

  final String vds1 = "space.folder.vds1";
  NamespaceTestUtils.addDS(namespaceService, vds1, "select * from space.folder.pds");

  final String vds2 = "space.folder.subfolder.vds2";
  NamespaceTestUtils.addDS(namespaceService, vds2, "select * from space.folder.vds1");

  this.catalog = new InformationSchemaCatalogImpl(namespaceService);
}
 

public byte[] serialize() {
  FlatBufferBuilder builder = new FlatBufferBuilder();
  builder.finish(serialize(builder));
  return builder.sizedByteArray();
}
 

public static ByteBuffer serializeRecordToFlatBuffer(FlatBufferBuilder builder, Record<byte[]> record) {
    checkNotNull(record, "record-context can't be null");
    Optional<EncryptionContext> encryptionCtx = (record instanceof RecordWithEncryptionContext)
            ? ((RecordWithEncryptionContext<byte[]>) record).getEncryptionCtx()
            : Optional.empty();
    Map<String, String> properties = record.getProperties();

    int encryptionCtxOffset = -1;
    int propertiesOffset = -1;

    if (properties != null && !properties.isEmpty()) {
        int[] propertiesOffsetArray = new int[properties.size()];
        int i = 0;
        for (Entry<String, String> property : properties.entrySet()) {
            propertiesOffsetArray[i++] = KeyValue.createKeyValue(builder, builder.createString(property.getKey()),
                    builder.createString(property.getValue()));
        }
        propertiesOffset = Message.createPropertiesVector(builder, propertiesOffsetArray);
    }

    if (encryptionCtx.isPresent()) {
        encryptionCtxOffset = createEncryptionCtxOffset(builder, encryptionCtx);
    }

    int payloadOffset = Message.createPayloadVector(builder, record.getValue());
    Message.startMessage(builder);
    Message.addPayload(builder, payloadOffset);
    if (encryptionCtxOffset != -1) {
        Message.addEncryptionCtx(builder, encryptionCtxOffset);
    }
    if (propertiesOffset != -1) {
        Message.addProperties(builder, propertiesOffset);
    }
    int endMessage = Message.endMessage(builder);
    builder.finish(endMessage);
    ByteBuffer bb = builder.dataBuffer();

    // to avoid copying of data, use same byte[] wrapped by ByteBuffer. But, ByteBuffer.array() returns entire array
    // so, it requires to read from offset:
    // builder.sizedByteArray()=>copies buffer: sizedByteArray(space, bb.capacity() - space)
    int space = bb.capacity() - builder.offset();
    return ByteBuffer.wrap(bb.array(), space, bb.capacity() - space);
}
 

@Test
public void testVectorEncoding() {
    val scalar = Nd4j.trueVector(new float[]{1, 2, 3, 4, 5});

    FlatBufferBuilder bufferBuilder = new FlatBufferBuilder(0);
    val fb = scalar.toFlatArray(bufferBuilder);
    bufferBuilder.finish(fb);
    val db = bufferBuilder.dataBuffer();

    val flat = FlatArray.getRootAsFlatArray(db);


    val restored = Nd4j.createFromFlatArray(flat);

    assertEquals(scalar, restored);
}
 

public byte[] serialize() {
  FlatBufferBuilder builder = new FlatBufferBuilder();
  builder.finish(serialize(builder));
  return builder.sizedByteArray();
}