GH-48076: [C++][Flight] fix GeneratorStream for Tables

After the changes in #47115, GeneratorStreams backed by anything else
than RecordBatches failed. This includes Tables and RecordBatchReaders.

This was caused by a too strict assumption that the
RecordBatchStream#GetSchemaPayload would always get called, which
is not the case when the GeneratorStream is backed by a Table
or a RecordBatchReader.

So to fix this, remove the assertion and instead initialize the writer
on first access.

Author: no23reason

cpp/src/arrow/flight/server.cc

@@ -285,12 +285,7 @@ class RecordBatchStream::RecordBatchStreamImpl {
 
   Status GetSchemaPayload(FlightPayload* payload) {
     if (!writer_) {
-      // Create the IPC writer on first call
-      auto payload_writer =
-          std::make_unique<ServerRecordBatchPayloadWriter>(&payload_deque_);
-      ARROW_ASSIGN_OR_RAISE(
-          writer_, ipc::internal::OpenRecordBatchWriter(std::move(payload_writer),
-                                                        reader_->schema(), options_));
+      RETURN_NOT_OK(InitializeWriter());
     }
 
     // Return the expected schema payload.
@@ -317,8 +312,7 @@ class RecordBatchStream::RecordBatchStreamImpl {
         return Status::OK();
       }
       if (!writer_) {
-        return Status::UnknownError(
-            "Writer should be initialized before reading Next batches");
+        RETURN_NOT_OK(InitializeWriter());
       }
       // One WriteRecordBatch call might generate multiple payloads, so we
       // need to collect them in a deque.
@@ -370,6 +364,15 @@ class RecordBatchStream::RecordBatchStreamImpl {
   ipc::IpcWriteOptions options_;
   std::unique_ptr<ipc::RecordBatchWriter> writer_;
   std::deque<FlightPayload> payload_deque_;
+
+  Status InitializeWriter() {
+    auto payload_writer =
+        std::make_unique<ServerRecordBatchPayloadWriter>(&payload_deque_);
+    ARROW_ASSIGN_OR_RAISE(
+        writer_, ipc::internal::OpenRecordBatchWriter(std::move(payload_writer),
+                                                      reader_->schema(), options_));
+    return Status::OK();
+  }
 };
 
 FlightMetadataWriter::~FlightMetadataWriter() = default;

python/pyarrow/tests/test_flight.py

@@ -246,6 +246,40 @@ def do_action(self, context, action):
         raise NotImplementedError
 
 
+class EchoTableStreamFlightServer(EchoFlightServer):
+    """An echo server that streams the whole table."""
+
+    def do_get(self, context, ticket):
+        return flight.GeneratorStream(
+            self.last_message.schema,
+            [self.last_message])
+
+    def list_actions(self, context):
+        return []
+
+    def do_action(self, context, action):
+        if action.type == "who-am-i":
+            return [context.peer_identity(), context.peer().encode("utf-8")]
+        raise NotImplementedError
+
+
+class EchoRecordBatchReaderStreamFlightServer(EchoFlightServer):
+    """An echo server that streams the whole table as a RecordBatchReader."""
+
+    def do_get(self, context, ticket):
+        return flight.GeneratorStream(
+            self.last_message.schema,
+            [self.last_message.to_reader()])
+
+    def list_actions(self, context):
+        return []
+
+    def do_action(self, context, action):
+        if action.type == "who-am-i":
+            return [context.peer_identity(), context.peer().encode("utf-8")]
+        raise NotImplementedError
+
+
 class GetInfoFlightServer(FlightServerBase):
     """A Flight server that tests GetFlightInfo."""
 
@@ -1362,7 +1396,7 @@ def test_flight_large_message():
         assert result.equals(data)
 
 
-def test_flight_generator_stream():
+def test_flight_generator_stream_of_batches():
     """Try downloading a flight of RecordBatches in a GeneratorStream."""
     data = pa.Table.from_arrays([
         pa.array(range(0, 10 * 1024))
@@ -1378,6 +1412,38 @@ def test_flight_generator_stream():
         assert result.equals(data)
 
 
+def test_flight_generator_stream_of_table():
+    """Try downloading a flight of Table in a GeneratorStream."""
+    data = pa.Table.from_arrays([
+        pa.array(range(0, 10 * 1024))
+    ], names=['a'])
+
+    with EchoTableStreamFlightServer() as server, \
+        FlightClient(('localhost', server.port)) as client:
+        writer, _ = client.do_put(flight.FlightDescriptor.for_path('test'),
+                                  data.schema)
+        writer.write_table(data)
+        writer.close()
+        result = client.do_get(flight.Ticket(b'')).read_all()
+        assert result.equals(data)
+
+
+def test_flight_generator_stream_of_record_batch_reader():
+    """Try downloading a flight of RecordBatchReader in a GeneratorStream."""
+    data = pa.Table.from_arrays([
+        pa.array(range(0, 10 * 1024))
+    ], names=['a'])
+
+    with EchoRecordBatchReaderStreamFlightServer() as server, \
+        FlightClient(('localhost', server.port)) as client:
+        writer, _ = client.do_put(flight.FlightDescriptor.for_path('test'),
+                                  data.schema)
+        writer.write_table(data)
+        writer.close()
+        result = client.do_get(flight.Ticket(b'')).read_all()
+        assert result.equals(data)
+
+
 def test_flight_invalid_generator_stream():
     """Try streaming data with mismatched schemas."""
     with InvalidStreamFlightServer() as server, \