Function WriteTensorMessage

cpp/src/arrow/ipc/metadata_internal.cc:1330–1362

}

Result<std::shared_ptr<Buffer>> WriteTensorMessage(const Tensor& tensor,
                                                   int64_t buffer_start_offset,
                                                   const IpcWriteOptions& options) {
  using TensorDimOffset = flatbuffers::Offset<flatbuf::TensorDim>;
  using TensorOffset = flatbuffers::Offset<flatbuf::Tensor>;

  FBB fbb;
  const int elem_size = tensor.type()->byte_width();

  flatbuf::Type fb_type_type;
  Offset fb_type;
  RETURN_NOT_OK(TensorTypeToFlatbuffer(fbb, *tensor.type(), &fb_type_type, &fb_type));

  std::vector<TensorDimOffset> dims;
  for (int i = 0; i < tensor.ndim(); ++i) {
    FBString name = fbb.CreateString(tensor.dim_name(i));
    dims.push_back(flatbuf::CreateTensorDim(fbb, tensor.shape()[i], name));
  }

  auto fb_shape = fbb.CreateVector(dims.data(), dims.size());

  flatbuffers::Offset<flatbuffers::Vector<int64_t>> fb_strides;
  fb_strides = fbb.CreateVector(tensor.strides().data(), tensor.strides().size());
  int64_t body_length = tensor.size() * elem_size;
  flatbuf::Buffer buffer(buffer_start_offset, body_length);

  TensorOffset fb_tensor =
      flatbuf::CreateTensor(fbb, fb_type_type, fb_type, fb_shape, fb_strides, &buffer);

  return WriteFBMessage(fbb, flatbuf::MessageHeader::MessageHeader_Tensor,
                        fb_tensor.Union(), body_length, options.metadata_version,
                        /*custom_metadata=*/nullptr, options.memory_pool);
}

Result<std::shared_ptr<Buffer>> WriteSparseTensorMessage(
    const SparseTensor& sparse_tensor, int64_t body_length,