| 660 | } // namespace |
| 661 | |
| 662 | arrow::Result<arrow::BufferVector> SerializePayloadToBuffers(const FlightPayload& msg) { |
| 663 | // Size of the IPC body (protobuf: data_body) |
| 664 | size_t body_size = 0; |
| 665 | // Size of the Protobuf "header" (everything except for the body) |
| 666 | size_t header_size = 0; |
| 667 | // Size of IPC header metadata (protobuf: data_header) |
| 668 | int32_t metadata_size = 0; |
| 669 | |
| 670 | // Write the descriptor if present |
| 671 | int32_t descriptor_size = 0; |
| 672 | if (msg.descriptor != nullptr) { |
| 673 | DCHECK_LE(msg.descriptor->size(), kInt32Max); |
| 674 | descriptor_size = static_cast<int32_t>(msg.descriptor->size()); |
| 675 | header_size += 1 + WireFormatLite::LengthDelimitedSize(descriptor_size); |
| 676 | } |
| 677 | |
| 678 | // App metadata tag if appropriate |
| 679 | int32_t app_metadata_size = 0; |
| 680 | if (msg.app_metadata && msg.app_metadata->size() > 0) { |
| 681 | DCHECK_LE(msg.app_metadata->size(), kInt32Max); |
| 682 | app_metadata_size = static_cast<int32_t>(msg.app_metadata->size()); |
| 683 | header_size += 1 + WireFormatLite::LengthDelimitedSize(app_metadata_size); |
| 684 | } |
| 685 | |
| 686 | const arrow::ipc::IpcPayload& ipc_msg = msg.ipc_message; |
| 687 | // No data in this payload (metadata-only). |
| 688 | bool has_ipc = ipc_msg.type != ipc::MessageType::NONE; |
| 689 | bool has_body = has_ipc ? ipc::Message::HasBody(ipc_msg.type) : false; |
| 690 | |
| 691 | if (has_ipc) { |
| 692 | DCHECK(has_body || ipc_msg.body_length == 0); |
| 693 | ARROW_RETURN_NOT_OK( |
| 694 | IpcMessageHeaderSize(ipc_msg, has_body, &header_size, &metadata_size)); |
| 695 | body_size = static_cast<size_t>(ipc_msg.body_length); |
| 696 | } |
| 697 | |
| 698 | // TODO(wesm): messages over 2GB unlikely to be yet supported |
| 699 | // Validated in WritePayload since returning error here causes gRPC to fail an assertion |
| 700 | DCHECK_LE(body_size, kInt32Max); |
| 701 | |
| 702 | // Allocate and initialize buffers |
| 703 | arrow::BufferVector buffers; |
| 704 | ARROW_ASSIGN_OR_RAISE(auto header_buf, arrow::AllocateBuffer(header_size)); |
| 705 | |
| 706 | // Force the header_stream to be destructed, which actually flushes |
| 707 | // the data into the slice. |
| 708 | { |
| 709 | ArrayOutputStream header_writer(const_cast<uint8_t*>(header_buf->mutable_data()), |
| 710 | static_cast<int>(header_size)); |
| 711 | CodedOutputStream header_stream(&header_writer); |
| 712 | |
| 713 | // Write descriptor |
| 714 | if (msg.descriptor != nullptr) { |
| 715 | WireFormatLite::WriteTag(pb::FlightData::kFlightDescriptorFieldNumber, |
| 716 | WireFormatLite::WIRETYPE_LENGTH_DELIMITED, &header_stream); |
| 717 | header_stream.WriteVarint32(descriptor_size); |
| 718 | header_stream.WriteRawMaybeAliased(msg.descriptor->data(), |
| 719 | static_cast<int>(msg.descriptor->size())); |