| 918 | } // namespace |
| 919 | |
| 920 | Status WriteTensor(const Tensor& tensor, io::OutputStream* dst, int32_t* metadata_length, |
| 921 | int64_t* body_length) { |
| 922 | const int elem_size = tensor.type()->byte_width(); |
| 923 | |
| 924 | *body_length = tensor.size() * elem_size; |
| 925 | |
| 926 | // Tensor metadata accounts for padding |
| 927 | if (tensor.is_contiguous()) { |
| 928 | RETURN_NOT_OK(WriteTensorHeader(tensor, dst, metadata_length)); |
| 929 | auto data = tensor.data(); |
| 930 | if (data && data->data()) { |
| 931 | RETURN_NOT_OK(dst->Write(data->data(), *body_length)); |
| 932 | } else { |
| 933 | *body_length = 0; |
| 934 | } |
| 935 | } else { |
| 936 | // The tensor written is made contiguous |
| 937 | Tensor dummy(tensor.type(), nullptr, tensor.shape()); |
| 938 | RETURN_NOT_OK(WriteTensorHeader(dummy, dst, metadata_length)); |
| 939 | |
| 940 | // TODO: Do we care enough about this temporary allocation to pass in a |
| 941 | // MemoryPool to this function? |
| 942 | ARROW_ASSIGN_OR_RAISE(auto scratch_space, |
| 943 | AllocateBuffer(tensor.shape()[tensor.ndim() - 1] * elem_size)); |
| 944 | |
| 945 | RETURN_NOT_OK(WriteStridedTensorData(0, 0, elem_size, tensor, |
| 946 | scratch_space->mutable_data(), dst)); |
| 947 | } |
| 948 | |
| 949 | return Status::OK(); |
| 950 | } |
| 951 | |
| 952 | Result<std::unique_ptr<Message>> GetTensorMessage(const Tensor& tensor, |
| 953 | MemoryPool* pool) { |