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

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, \