GH-26727: [C++][Flight] Use ipc::RecordBatchWriter with custom IpcPayloadWriter for TransportMessageWriter (DoExchange) (#47410)

### Rationale for this change

The Flight Server DoExchange method currently does not support Dictionary replacement or Dictionary Deltas, similar to how the client currently behaves or how we do for DoGet we should use an `ipc::RecordBatchWriter` with a custom `IpcPayloadWriter` instead of reimplementing Dictionary Replacement / Deltas logic.

### What changes are included in this PR?

Removes manually generation of individual ipc Payloads and uses an `ipc::RecordBatchWriter` and a custom  `TransportMessagePayloadWriter` to modify the `IpcPayloads` into `FlightPayloads`.

### Are these changes tested?

Yes, existing tests cover the DoExchange functionality and new test for Python has been added where Dictionary deltas are being send via DoExchange. The test was failing before this change because the dictionary wasn't updated:
```python
            received_table = reader.read_all()
            expected_table = simple_dicts_table()
>           assert received_table.equals(expected_table)
E           assert False
E            +  where False = equals(pyarrow.Table\nsome_dicts: dictionary<values=string, indices=int64, ordered=0>\n----\nsome_dicts: [  -- dictionary:\n["foo... -- dictionary:\n["foo","baz","quux"]  -- indices:\n[2,1],  -- dictionary:\n["foo","baz","quux","new"]  -- indices:\n[0,3]])
E            +    where equals = pyarrow.Table\nsome_dicts: dictionary<values=string, indices=int64, ordered=0>\n----\nsome_dicts: [  -- dictionary:\n["foo...ull],  -- dictionary:\n["foo","baz","quux"]  -- indices:\n[2,1],  -- dictionary:\n["foo","baz","quux"]  -- indices:\n[0,3]].equals

pyarrow/tests/test_flight.py:2596: AssertionError
================================================================================== short test summary info ===================================================================================
FAILED pyarrow/tests/test_flight.py::test_flight_dictionary_deltas_do_exchange - assert False

```

### Are there any user-facing changes?

No, only that the expected dictionary replacement/deltas will work for DoExchange.

* GitHub Issue: #26727

Authored-by: Raúl Cumplido <[email protected]>
Signed-off-by: Raúl Cumplido <[email protected]>

Author: raulcd

cpp/src/arrow/flight/transport_server.cc

@@ -26,6 +26,7 @@
 #include "arrow/ipc/reader.h"
 #include "arrow/result.h"
 #include "arrow/status.h"
+#include "arrow/util/logging_internal.h"
 
 namespace arrow {
 namespace flight {
@@ -160,25 +161,63 @@ class TransportMessageReader final : public FlightMessageReader {
   std::shared_ptr<Buffer> app_metadata_;
 };
 
-// TODO(ARROW-10787): this should use the same writer/ipc trick as client
+/// \brief An IpcPayloadWriter for ServerDataStream.
+///
+/// To support app_metadata and reuse the existing IPC infrastructure,
+/// this takes a pointer to a buffer to be combined with the IPC
+/// payload when writing a Flight payload.
+class TransportMessagePayloadWriter : public ipc::internal::IpcPayloadWriter {
+ public:
+  TransportMessagePayloadWriter(ServerDataStream* stream,
+                                std::shared_ptr<Buffer>* app_metadata)
+      : stream_(stream), app_metadata_(app_metadata) {}
+
+  Status Start() override { return Status::OK(); }
+  Status WritePayload(const ipc::IpcPayload& ipc_payload) override {
+    FlightPayload payload;
+    payload.ipc_message = ipc_payload;
+
+    if (ipc_payload.type == ipc::MessageType::RECORD_BATCH && *app_metadata_) {
+      payload.app_metadata = std::move(*app_metadata_);
+    }
+    ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload));
+    if (!success) {
+      return MakeFlightError(
+          FlightStatusCode::Internal,
+          "Could not write record batch to stream (client disconnect?)");
+    }
+    return arrow::Status::OK();
+  }
+  Status Close() override {
+    // Closing is handled one layer up in TransportMessageWriter::Close
+    return Status::OK();
+  }
+
+ private:
+  ServerDataStream* stream_;
+  std::shared_ptr<Buffer>* app_metadata_;
+};
+
 class TransportMessageWriter final : public FlightMessageWriter {
  public:
   explicit TransportMessageWriter(ServerDataStream* stream)
-      : stream_(stream), ipc_options_(::arrow::ipc::IpcWriteOptions::Defaults()) {}
+      : stream_(stream),
+        app_metadata_(nullptr),
+        ipc_options_(::arrow::ipc::IpcWriteOptions::Defaults()) {}
 
   Status Begin(const std::shared_ptr<Schema>& schema,
                const ipc::IpcWriteOptions& options) override {
-    if (started_) {
+    if (batch_writer_) {
       return Status::Invalid("This writer has already been started.");
     }
-    started_ = true;
     ipc_options_ = options;
+    std::unique_ptr<ipc::internal::IpcPayloadWriter> payload_writer(
+        new TransportMessagePayloadWriter(stream_, &app_metadata_));
 
-    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);
+    ARROW_ASSIGN_OR_RAISE(batch_writer_,
+                          ipc::internal::OpenRecordBatchWriter(std::move(payload_writer),
+                                                               schema, ipc_options_));
+    return Status::OK();
   }
 
   Status WriteRecordBatch(const RecordBatch& batch) override {
@@ -188,71 +227,56 @@ class TransportMessageWriter final : public FlightMessageWriter {
   Status WriteMetadata(std::shared_ptr<Buffer> app_metadata) override {
     FlightPayload payload{};
     payload.app_metadata = app_metadata;
-    return WritePayload(payload);
+    ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload));
+    if (!success) {
+      ARROW_RETURN_NOT_OK(Close());
+      return MakeFlightError(FlightStatusCode::Internal,
+                             "Could not write metadata to stream (client disconnect?)");
+    }
+    // Those messages are not written through the batch writer,
+    // count them separately to include them in the stats.
+    extra_messages_++;
+    return Status::OK();
   }
 
   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;
+    app_metadata_ = app_metadata;
+    auto status = batch_writer_->WriteRecordBatch(batch);
+    if (!status.ok()) {
+      ARROW_RETURN_NOT_OK(Close());
     }
-    RETURN_NOT_OK(ipc::GetRecordBatchPayload(batch, ipc_options_, &payload.ipc_message));
-    RETURN_NOT_OK(WritePayload(payload));
-    ++stats_.num_record_batches;
-    return Status::OK();
+    return status;
   }
 
   Status Close() override {
-    // It's fine to Close() without writing data
+    if (batch_writer_) {
+      RETURN_NOT_OK(batch_writer_->Close());
+    }
     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();
+  ipc::WriteStats stats() const override {
+    ARROW_CHECK_NE(batch_writer_, nullptr);
+    auto write_stats = batch_writer_->stats();
+    write_stats.num_messages += extra_messages_;
+    return write_stats;
   }
 
+ private:
   Status CheckStarted() {
-    if (!started_) {
+    if (!batch_writer_) {
       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_;
+  std::unique_ptr<ipc::RecordBatchWriter> batch_writer_;
+  std::shared_ptr<Buffer> app_metadata_;
   ::arrow::ipc::IpcWriteOptions ipc_options_;
-  ipc::DictionaryFieldMapper mapper_;
-  ipc::WriteStats stats_;
-  bool started_ = false;
-  bool dictionaries_written_ = false;
+  int64_t extra_messages_ = 0;
 };
 
 /// \brief Adapt TransportDataStream to the FlightMetadataWriter

python/pyarrow/tests/test_flight.py

@@ -114,12 +114,12 @@ def simple_ints_table():
 
 def simple_dicts_table():
     dict_values = pa.array(["foo", "baz", "quux"], type=pa.utf8())
-    new_dict_values = pa.array(["bar", "qux"], type=pa.utf8())
+    new_dict_values = pa.array(["foo", "baz", "quux", "new"], type=pa.utf8())
     data = [
         pa.chunked_array([
             pa.DictionaryArray.from_arrays([1, 0, None], dict_values),
             pa.DictionaryArray.from_arrays([2, 1], dict_values),
-            pa.DictionaryArray.from_arrays([0, 1], new_dict_values)
+            pa.DictionaryArray.from_arrays([0, 3], new_dict_values)
         ])
     ]
     return pa.Table.from_arrays(data, names=['some_dicts'])
@@ -2536,3 +2536,39 @@ def received_headers(self, headers):
         assert ("x-header-bin", b"header\x01value") in factory.headers
         assert ("x-trailer", "trailer-value") in factory.headers
         assert ("x-trailer-bin", b"trailer\x01value") in factory.headers
+
+
+def test_flight_dictionary_deltas_do_exchange():
+    class DeltaFlightServer(ConstantFlightServer):
+        def do_exchange(self, context, descriptor, reader, writer):
+            if descriptor.command == b'dict_deltas':
+                expected_table = simple_dicts_table()
+                received_table = reader.read_all()
+                assert received_table.equals(expected_table)
+
+                options = pa.ipc.IpcWriteOptions(emit_dictionary_deltas=True)
+                writer.begin(expected_table.schema, options=options)
+                # TODO: GH-47422: Inspect ReaderStats once exposed and validate deltas
+                writer.write_table(expected_table)
+
+    with DeltaFlightServer() as server, \
+            FlightClient(('localhost', server.port)) as client:
+        expected_table = simple_dicts_table()
+
+        descriptor = flight.FlightDescriptor.for_command(b"dict_deltas")
+        writer, reader = client.do_exchange(descriptor,
+                                            options=flight.FlightCallOptions(
+                                                write_options=pa.ipc.IpcWriteOptions(
+                                                    emit_dictionary_deltas=True)
+                                            )
+                                            )
+        # Send client table with dictionary updates (deltas should be sent)
+        with writer:
+            writer.begin(expected_table.schema, options=pa.ipc.IpcWriteOptions(
+                emit_dictionary_deltas=True))
+            writer.write_table(expected_table)
+            writer.done_writing()
+            received_table = reader.read_all()
+
+        # TODO: GH-47422: Inspect ReaderStats once exposed and validate deltas
+        assert received_table.equals(expected_table)