Java Code Examples for org.apache.thrift.TBase#write()

private void serializeArguments(RpcInvocation rpcInvocation, TBinaryProtocol protocol) throws TException {
    String argsClassName = ThriftUtil.generateArgsClassName(
            rpcInvocation.getServiceInterface().getName(), rpcInvocation.getMethod().getName());
    TBase argsClassObj = getClazzInstance(argsClassName);

    if (rpcInvocation.getArguments() == null) {
        argsClassObj.write(protocol);
        return;
    }

    String argsFieldsClassName = ThriftUtil.generateIDLFieldsClassName(argsClassName);
    Class<?> argsFieldsClazz = getClazz(argsFieldsClassName);

    Object[] fieldEnums = argsFieldsClazz.getEnumConstants();
    if (fieldEnums == null && fieldEnums.length != rpcInvocation.getArguments().length) {
        throw new ProtocolException("argument num is " + rpcInvocation.getArguments().length + " is not match with fields in " + argsFieldsClassName);
    }
    for (int i = 0; i < fieldEnums.length; i++) {
        TFieldIdEnum fieldIdEnum = (TFieldIdEnum) fieldEnums[i];
        if (fieldIdEnum == null) {
            continue;
        }
        argsClassObj.setFieldValue(fieldIdEnum, rpcInvocation.getArguments()[i]);
    }
    argsClassObj.write(protocol);
}
 

@SuppressWarnings("rawtypes")
private void writeResult(final TProtocol out, final TMessage msg, final WriterHandler onComplete, TBase args,
		final TBase result) {
	try {
		onComplete.beforeWrite(msg, args, result);
		// if (!isOneway()) {
		out.writeMessageBegin(new TMessage(msg.name, TMessageType.REPLY, msg.seqid));
		if (result != null) {
			result.write(out);
		} else {
			out.writeStructBegin(null);
			out.writeFieldStop();
			out.writeStructEnd();
		}
		out.writeMessageEnd();
		out.getTransport().flush();
		// }
		onComplete.afterWrite(msg, null, TMessageType.REPLY, args, result);
	} catch (Throwable e) {
		onComplete.afterWrite(msg, e, TMessageType.EXCEPTION, args, result);
	}
}
 

/**
 * Encodes a thrift object into a DEFLATE-compressed binary array.
 *
 * @param tBase Object to encode.
 * @return Deflated, encoded object.
 * @throws CodingException If the object could not be encoded.
 */
public static byte[] deflateNonNull(TBase<?, ?> tBase) throws CodingException {
  requireNonNull(tBase);

  // NOTE: Buffering is needed here for performance.
  // There are actually 2 buffers in play here - the BufferedOutputStream prevents thrift from
  // causing a call to deflate() on every encoded primitive. The DeflaterOutputStream buffer
  // allows the underlying Deflater to operate on a larger chunk at a time without stopping to
  // copy the intermediate compressed output to outBytes.
  // See http://bugs.java.com/bugdatabase/view_bug.do?bug_id=4986239
  ByteArrayOutputStream outBytes = new ByteArrayOutputStream();
  TTransport transport = new TIOStreamTransport(
      new BufferedOutputStream(
          new DeflaterOutputStream(outBytes, new Deflater(DEFLATE_LEVEL), DEFLATER_BUFFER_SIZE),
          DEFLATER_BUFFER_SIZE));
  try {
    TProtocol protocol = PROTOCOL_FACTORY.getProtocol(transport);
    tBase.write(protocol);
    transport.close(); // calls finish() on the underlying stream, completing the compression
    return outBytes.toByteArray();
  } catch (TException e) {
    throw new CodingException("Failed to serialize: " + tBase, e);
  } finally {
    transport.close();
  }
}
 

private void write(final TBase<?, ?> base) throws TException {
    final TProtocol protocol = protocolFactory.getProtocol(transport);

    // write chunk header
    writeChunkHeader(protocol);

    // write header
    final Header header = locator.headerLookup(base);
    if (header == null) {
        throw new TException("header must not be null base:" + base);
    }
    HeaderUtils.writeHeader(protocol, header);

    base.write(protocol);

    if (isNeedFlush()) {
        flush();
    }
}
 

private void write(final TBase<?, ?> base, final String fieldName, final List<ByteArrayOutput> list) throws TException {
    final TReplaceListProtocol protocol = new TReplaceListProtocol(protocolFactory.getProtocol(transport));

    // write chunk header
    writeChunkHeader(protocol);

    // write header
    final Header header = locator.headerLookup(base);
    if (header == null) {
        throw new TException("header must not be null base:" + base);
    }
    HeaderUtils.writeHeader(protocol, header);
    if (CollectionUtils.hasLength(list)) {
        protocol.addReplaceField(fieldName, list);
    }

    base.write(protocol);

    if (isNeedFlush()) {
        flush();
    }
}
 

public static void serialize(ThriftProtocol protocol, TBase<?, ?> source, OutputStream stream)
    throws ThriftException {
  try (TTransport transport = new TIOStreamTransport(stream)) {
    TProtocol thriftProtocol = getProtocolFactory(protocol).getProtocol(transport);
    try {
      source.write(thriftProtocol);
    } catch (TException e) {
      throw new ThriftException(e);
    }
  }
}
 

@Override
public byte[] serialize(TBase<?, ?> base, HeaderEntity headerEntity) throws TException {
    baos.reset();

    writeHeader(base, headerEntity);
    base.write(protocol);
    return baos.toByteArray();
}
 

public void serialize(TBase<?, ?> base, OutputStream outputStream) throws TException {
    tOutputStreamTransport.open(outputStream);
    try {
        final Header header = tBaseLocator.headerLookup(base);
        if (header == null) {
            throw new TException("header must not be null base:" + base);
        }
        HeaderUtils.writeHeader(protocol, header);
        base.write(protocol);
    } finally {
        tOutputStreamTransport.close();
    }
}
 

public void write(final TBase<?, ?> base) throws TException {
    final TBaseStreamNode node = new TBaseStreamNode(transport);
    node.setClassName(base.getClass().getName());
    node.setBeginPosition(transport.getBufferPosition());

    final TProtocol protocol = protocolFactory.getProtocol(transport);
    base.write(protocol);

    node.setEndPosition(transport.getBufferPosition());
    nodes.add(node);
}
 

private void writeReadCompare(TBase<?, ?> a)
        throws TException, InstantiationException, IllegalAccessException {
  ProtocolPipe[] pipes = {new ProtocolReadToWrite(), new BufferedProtocolReadToWrite(ThriftSchemaConverter.toStructType((Class<TBase<?, ?>>)a.getClass()))};
  for (ProtocolPipe p : pipes) {
    final ByteArrayOutputStream in = new ByteArrayOutputStream();
    final ByteArrayOutputStream out = new ByteArrayOutputStream();
    a.write(protocol(in));
    p.readOne(protocol(new ByteArrayInputStream(in.toByteArray())), protocol(out));
    TBase<?, ?> b = a.getClass().newInstance();
    b.read(protocol(new ByteArrayInputStream(out.toByteArray())));

    assertEquals(p.getClass().getSimpleName(), a, b);
  }
}
 

private void validateThrift(String[] expectations, TBase<?, ?> a)
      throws TException {
    final ThriftSchemaConverter thriftSchemaConverter = new ThriftSchemaConverter();
//      System.out.println(a);
    final Class<TBase<?,?>> class1 = (Class<TBase<?,?>>)a.getClass();
    final MessageType schema = thriftSchemaConverter.convert(class1);
    LOG.info("{}", schema);
    final StructType structType = thriftSchemaConverter.toStructType(class1);
    ExpectationValidatingRecordConsumer recordConsumer = new ExpectationValidatingRecordConsumer(new ArrayDeque<String>(Arrays.asList(expectations)));
    final MessageColumnIO columnIO = new ColumnIOFactory().getColumnIO(schema);
    ParquetWriteProtocol p = new ParquetWriteProtocol(new RecordConsumerLoggingWrapper(recordConsumer), columnIO, structType);
    a.write(p);
  }
 

private void writeParquetFile(List<TBase> recordsToWrite, Configuration conf, Path parquetFile) throws IOException, InterruptedException, org.apache.thrift.TException {
  //create a test file
  final TProtocolFactory protocolFactory = new TCompactProtocol.Factory();
  final TaskAttemptID taskId = new TaskAttemptID("local", 0, true, 0, 0);
  Class writeClass = recordsToWrite.get(0).getClass();
  final ThriftToParquetFileWriter w = new ThriftToParquetFileWriter(parquetFile, ContextUtil.newTaskAttemptContext(conf, taskId), protocolFactory, writeClass);
  final ByteArrayOutputStream baos = new ByteArrayOutputStream();
  final TProtocol protocol = protocolFactory.getProtocol(new TIOStreamTransport(baos));
  for (TBase recordToWrite : recordsToWrite) {
    recordToWrite.write(protocol);
  }
  w.write(new BytesWritable(baos.toByteArray()));
  w.close();
}
 

private static void write(TBase<?, ?> tbase, OutputStream to) throws IOException {
  try {
    tbase.write(protocol(to));
  } catch (TException e) {
    throw new IOException("can not write " + tbase, e);
  }
}
 

private static void write(TBase<?, ?> tbase, OutputStream to) throws IOException {
  try {
    tbase.write(protocol(to));
  } catch (TException e) {
    throw new IOException("can not write " + tbase, e);
  }
}
 

private static HttpData encodeSuccess(ServiceRequestContext ctx,
                                      RpcResponse reply,
                                      SerializationFormat serializationFormat,
                                      String methodName, int seqId,
                                      TBase<?, ?> result) {

    final ByteBuf buf = ctx.alloc().buffer(128);
    boolean success = false;
    try {
        final TTransport transport = new TByteBufTransport(buf);
        final TProtocol outProto = ThriftProtocolFactories.get(serializationFormat).getProtocol(transport);
        final TMessage header = new TMessage(methodName, TMessageType.REPLY, seqId);
        outProto.writeMessageBegin(header);
        result.write(outProto);
        outProto.writeMessageEnd();

        ctx.logBuilder().responseContent(reply, new ThriftReply(header, result));

        final HttpData encoded = PooledHttpData.wrap(buf);
        success = true;
        return encoded;
    } catch (TException e) {
        throw new Error(e); // Should never reach here.
    } finally {
        if (!success) {
            buf.release();
        }
    }
}
 

/**
 * Encodes the given value of type {@code T} onto the given output stream using provided {@link
 * ThriftCoder#protocolFactory}.
 *
 * @param value {@link org.apache.thrift.TBase} to encode.
 * @param outStream stream to output encoded value to.
 * @throws IOException if writing to the {@code OutputStream} fails for some reason
 * @throws CoderException if the value could not be encoded for some reason
 */
@Override
public void encode(T value, OutputStream outStream) throws CoderException, IOException {
  ByteArrayOutputStream baos = new ByteArrayOutputStream();
  TProtocol protocol = protocolFactory.getProtocol(new TIOStreamTransport(baos));
  try {
    TBase<?, ?> tBase = (TBase<?, ?>) value;
    tBase.write(protocol);
  } catch (Exception te) {
    throw new CoderException("Could not write value. Error: " + te.getMessage());
  }
  outStream.write(baos.toByteArray());
}
 

public byte[] continueSerialize(TBase<?, ?> base) throws TException {
    writeHeader(base);
    base.write(protocol);
    return baos.toByteArray();
}
 

/**
 * Serialize the Thrift object into a byte array. The process is simple,
 * just clear the byte array output, write the object into it, and grab the
 * raw bytes.
 *
 * @param base The object to serialize
 * @return Serialized object in byte[] format
 */
public byte[] serialize(TBase base) throws TException {
  outputStream.reset();
  outputStream.write(SECRET_BYTE);
  outputStream.write(COMPACT_PROTOCOL_BYTE);
  base.write(protocol);
  return outputStream.toByteArray();
}