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. Also, to accommodate for this case, drop the incoming
message when initializing the writer in Next, as the message there
is of the SCHEMA type and we want RECORD_BATCH one.

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,15 @@ class RecordBatchStream::RecordBatchStreamImpl {
         return Status::OK();
       }
       if (!writer_) {
-        return Status::UnknownError(
-            "Writer should be initialized before reading Next batches");
+        RETURN_NOT_OK(InitializeWriter());
+        // If the writer has not been initialized yet, the first batch in the payload
+        // queue is going to be a SCHEMA one. In this context, that is
+        // unexpected, so drop it from the queue so that there is a RECORD_BATCH
+        // message on the top (same as would be if the writer had been initialized
+        // in GetSchemaPayload).
+        if (payload_deque_.front().ipc_message.type == ipc::MessageType::SCHEMA) {
+          payload_deque_.pop_front();
+        }
       }
       // One WriteRecordBatch call might generate multiple payloads, so we
       // need to collect them in a deque.
@@ -370,6 +372,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, \