[C++][Flight] DoExchange doesn't support dictionary replacement

[asfimport]

Looking at the server DoExchangeMessageWriter implementation, it seems to assume that the dictionary values for a given dictionary-encoded field will never change accross record batches.

Reporter: Antoine Pitrou / @pitrou

Related issues:

• [C++] Support sending delta DictionaryBatch or replacement DictionaryBatch in IPC stream writer class (is related to)
• [C++][FlightRPC] Implement a UCX transport (is related to)

_{Note: This issue was originally created as ARROW-10787. Please see the migration documentation for further details.}

[asfimport]

Antoine Pitrou / @pitrou:
cc @lidavidm

[raulcd]

The DoExchange code that generates the messages and the payloads does not reuse the code changed on server.cc here:

• GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload #47115

It uses its own implementation here:

arrow/cpp/src/arrow/flight/transport_server.cc

Lines 163 to 256 in 066b216

	// TODO(ARROW-10787): this should use the same writer/ipc trick as client
	class TransportMessageWriter final : public FlightMessageWriter {
	public:
	explicit TransportMessageWriter(ServerDataStream* stream)
	: stream_(stream), ipc_options_(::arrow::ipc::IpcWriteOptions::Defaults()) {}
	Status Begin(const std::shared_ptr<Schema>& schema,
	const ipc::IpcWriteOptions& options) override {
	if (started_) {
	return Status::Invalid("This writer has already been started.");
	}
	started_ = true;
	ipc_options_ = options;
	RETURN_NOT_OK(mapper_.AddSchemaFields(*schema));
	FlightPayload schema_payload;
	RETURN_NOT_OK(ipc::GetSchemaPayload(*schema, ipc_options_, mapper_,
	&schema_payload.ipc_message));
	return WritePayload(schema_payload);
	}
	Status WriteRecordBatch(const RecordBatch& batch) override {
	return WriteWithMetadata(batch, nullptr);
	}
	Status WriteMetadata(std::shared_ptr<Buffer> app_metadata) override {
	FlightPayload payload{};
	payload.app_metadata = app_metadata;
	return WritePayload(payload);
	}
	Status WriteWithMetadata(const RecordBatch& batch,
	std::shared_ptr<Buffer> app_metadata) override {
	RETURN_NOT_OK(CheckStarted());
	RETURN_NOT_OK(EnsureDictionariesWritten(batch));
	FlightPayload payload{};
	if (app_metadata) {
	payload.app_metadata = app_metadata;
	}
	RETURN_NOT_OK(ipc::GetRecordBatchPayload(batch, ipc_options_, &payload.ipc_message));
	RETURN_NOT_OK(WritePayload(payload));
	++stats_.num_record_batches;
	return Status::OK();
	}
	Status Close() override {
	// It's fine to Close() without writing data
	return Status::OK();
	}
	ipc::WriteStats stats() const override { return stats_; }
	private:
	Status WritePayload(const FlightPayload& payload) {
	ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload));
	if (!success) {
	return MakeFlightError(FlightStatusCode::Internal,
	"Could not write metadata to stream (client disconnect?)");
	}
	++stats_.num_messages;
	return Status::OK();
	}
	Status CheckStarted() {
	if (!started_) {
	return Status::Invalid("This writer is not started. Call Begin() with a schema");
	}
	return Status::OK();
	}
	Status EnsureDictionariesWritten(const RecordBatch& batch) {
	if (dictionaries_written_) {
	return Status::OK();
	}
	dictionaries_written_ = true;
	ARROW_ASSIGN_OR_RAISE(const auto dictionaries,
	ipc::CollectDictionaries(batch, mapper_));
	for (const auto& pair : dictionaries) {
	FlightPayload payload{};
	RETURN_NOT_OK(ipc::GetDictionaryPayload(pair.first, pair.second, ipc_options_,
	&payload.ipc_message));
	RETURN_NOT_OK(WritePayload(payload));
	++stats_.num_dictionary_batches;
	}
	return Status::OK();
	}
	ServerDataStream* stream_;
	::arrow::ipc::IpcWriteOptions ipc_options_;
	ipc::DictionaryFieldMapper mapper_;
	ipc::WriteStats stats_;
	bool started_ = false;
	bool dictionaries_written_ = false;
	};

I think we can do the same trick we did for DoGet were we use an internal ipc::RecordBatchWriter and our own IpcPayloadWriter to generate the FlightPayload so we can get rid of the manual dictionary collection and generation.
@lidavidm what do you think?

[lidavidm]

Yes, well, that's why there's a comment there saying we should do that :)

[raulcd]

Issue resolved by pull request 47410
#47410