GH-48076: [C++][Flight] fix GeneratorStream for Tables (#48082)

### Rationale for this change

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.

### What changes are included in this PR?

Removal of the problematic assertion and initialization of the writer object when it is needed first.
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 or DICTIONARY_BATCH one.

### Are these changes tested?

Yes, via CI. Tests for the GeneratorStreams were extended so that they test GeneratorStreams backed by Tables and RecordBatchReaders, not just RecordBatches. 

### Are there any user-facing changes?

No, just a fix for a regression restoring the functionality from version 21.0.0 and earlier.

* GitHub Issue: #48076

Authored-by: Dan Homola <[email protected]>
Signed-off-by: David Li <[email protected]>

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,101 @@ def test_flight_generator_stream():
         assert result.equals(data)
 
 
+def test_flight_generator_stream_of_batches_with_dict():
+    """
+    Try downloading a flight of RecordBatches with dictionaries
+    in a GeneratorStream.
+    """
+    data = pa.Table.from_arrays([
+        pa.array(["foo", "bar", "baz", "foo", "foo"],
+                 pa.dictionary(pa.int64(), pa.utf8())),
+        pa.array([123, 234, 345, 456, 567])
+    ], names=['a', 'b'])
+
+    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_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_table_with_dict():
+    """
+    Try downloading a flight of Table with dictionaries
+    in a GeneratorStream.
+    """
+    data = pa.Table.from_arrays([
+        pa.array(["foo", "bar", "baz", "foo", "foo"],
+                 pa.dictionary(pa.int64(), pa.utf8())),
+        pa.array([123, 234, 345, 456, 567])
+    ], names=['a', 'b'])
+
+    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_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_generator_stream_of_record_batch_reader_with_dict():
+    """
+    Try downloading a flight of RecordBatchReader with dictionaries
+    in a GeneratorStream.
+    """
+    data = pa.Table.from_arrays([
+        pa.array(["foo", "bar", "baz", "foo", "foo"],
+                 pa.dictionary(pa.int64(), pa.utf8())),
+        pa.array([123, 234, 345, 456, 567])
+    ], names=['a', 'b'])
+
+    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, \