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