Use FlightStreamChunk in Flight/C++

cpp/src/arrow/flight/client.cc

@@ -147,9 +147,7 @@ class GrpcStreamReader : public FlightStreamReader {
                      std::unique_ptr<grpc::ClientReader<pb::FlightData>> stream,
                      std::unique_ptr<GrpcStreamReader>* out);
   std::shared_ptr<Schema> schema() const override;
-  Status ReadNext(std::shared_ptr<RecordBatch>* out) override;
-  Status ReadWithMetadata(std::shared_ptr<RecordBatch>* out,
-                          std::shared_ptr<Buffer>* app_metadata) override;
+  Status Next(FlightStreamChunk* out) override;
   void Cancel() override;
 
  private:
@@ -224,15 +222,10 @@ std::shared_ptr<Schema> GrpcStreamReader::schema() const {
   return batch_reader_->schema();
 }
 
-Status GrpcStreamReader::ReadNext(std::shared_ptr<RecordBatch>* out) {
-  std::shared_ptr<Buffer> app_metadata;
-  return ReadWithMetadata(out, &app_metadata);
-}
-
-Status GrpcStreamReader::ReadWithMetadata(std::shared_ptr<RecordBatch>* out,
-                                          std::shared_ptr<Buffer>* app_metadata) {
-  RETURN_NOT_OK(batch_reader_->ReadNext(out));
-  *app_metadata = std::move(last_app_metadata_);
+Status GrpcStreamReader::Next(FlightStreamChunk* out) {
+  out->app_metadata = nullptr;
+  RETURN_NOT_OK(batch_reader_->ReadNext(&out->data));
+  out->app_metadata = std::move(last_app_metadata_);
   return Status::OK();
 }
 

cpp/src/arrow/flight/flight-benchmark.cc

@@ -109,7 +109,7 @@ Status RunPerformanceTest(const std::string& hostname, const int port) {
     std::unique_ptr<FlightStreamReader> reader;
     RETURN_NOT_OK(client->DoGet(endpoint.ticket, &reader));
 
-    std::shared_ptr<RecordBatch> batch;
+    FlightStreamChunk batch;
 
     // This is hard-coded for right now, 4 columns each with int64
     const int bytes_per_record = 32;
@@ -120,26 +120,26 @@ Status RunPerformanceTest(const std::string& hostname, const int port) {
     int64_t num_bytes = 0;
     int64_t num_records = 0;
     while (true) {
-      RETURN_NOT_OK(reader->ReadNext(&batch));
-      if (!batch) {
+      RETURN_NOT_OK(reader->Next(&batch));
+      if (!batch.data) {
         break;
       }
 
       if (verify) {
-        auto values =
-            reinterpret_cast<const int64_t*>(batch->column_data(0)->buffers[1]->data());
+        auto values = reinterpret_cast<const int64_t*>(
+            batch.data->column_data(0)->buffers[1]->data());
         const int64_t start = token.start() + num_records;
-        for (int64_t i = 0; i < batch->num_rows(); ++i) {
+        for (int64_t i = 0; i < batch.data->num_rows(); ++i) {
           if (values[i] != start + i) {
             return Status::Invalid("verification failure");
           }
         }
       }
 
-      num_records += batch->num_rows();
+      num_records += batch.data->num_rows();
 
       // Hard-coded
-      num_bytes += batch->num_rows() * bytes_per_record;
+      num_bytes += batch.data->num_rows() * bytes_per_record;
     }
     stats.Update(num_records, num_bytes);
     return Status::OK();

cpp/src/arrow/flight/flight-test.cc

@@ -214,16 +214,16 @@ class TestFlightClient : public ::testing::Test {
     std::unique_ptr<FlightStreamReader> stream;
     ASSERT_OK(client_->DoGet(ticket, &stream));
 
-    std::shared_ptr<RecordBatch> chunk;
+    FlightStreamChunk chunk;
     for (int i = 0; i < num_batches; ++i) {
-      ASSERT_OK(stream->ReadNext(&chunk));
-      ASSERT_NE(nullptr, chunk);
-      ASSERT_BATCHES_EQUAL(*expected_batches[i], *chunk);
+      ASSERT_OK(stream->Next(&chunk));
+      ASSERT_NE(nullptr, chunk.data);
+      ASSERT_BATCHES_EQUAL(*expected_batches[i], *chunk.data);
     }
 
     // Stream exhausted
-    ASSERT_OK(stream->ReadNext(&chunk));
-    ASSERT_EQ(nullptr, chunk);
+    ASSERT_OK(stream->Next(&chunk));
+    ASSERT_EQ(nullptr, chunk.data);
   }
 
  protected:
@@ -284,18 +284,17 @@ class MetadataTestServer : public FlightServerBase {
   Status DoPut(const ServerCallContext& context,
                std::unique_ptr<FlightMessageReader> reader,
                std::unique_ptr<FlightMetadataWriter> writer) override {
-    std::shared_ptr<arrow::RecordBatch> chunk;
-    std::shared_ptr<Buffer> app_metadata;
+    FlightStreamChunk chunk;
     int counter = 0;
     while (true) {
-      RETURN_NOT_OK(reader->ReadWithMetadata(&chunk, &app_metadata));
-      if (chunk == nullptr) break;
-      if (app_metadata == nullptr) {
+      RETURN_NOT_OK(reader->Next(&chunk));
+      if (chunk.data == nullptr) break;
+      if (chunk.app_metadata == nullptr) {
         return Status::Invalid("Expected application metadata to be provided");
       }
-      if (std::to_string(counter) != app_metadata->ToString()) {
+      if (std::to_string(counter) != chunk.app_metadata->ToString()) {
         return Status::Invalid("Expected metadata value: " + std::to_string(counter) +
-                               " but got: " + app_metadata->ToString());
+                               " but got: " + chunk.app_metadata->ToString());
       }
       auto metadata = Buffer::FromString(std::to_string(counter));
       RETURN_NOT_OK(writer->WriteMetadata(*metadata));
@@ -770,18 +769,17 @@ TEST_F(TestMetadata, DoGet) {
   BatchVector expected_batches;
   ASSERT_OK(ExampleIntBatches(&expected_batches));
 
-  std::shared_ptr<RecordBatch> chunk;
-  std::shared_ptr<Buffer> metadata;
+  FlightStreamChunk chunk;
   auto num_batches = static_cast<int>(expected_batches.size());
   for (int i = 0; i < num_batches; ++i) {
-    ASSERT_OK(stream->ReadWithMetadata(&chunk, &metadata));
-    ASSERT_NE(nullptr, chunk);
-    ASSERT_NE(nullptr, metadata);
-    ASSERT_BATCHES_EQUAL(*expected_batches[i], *chunk);
-    ASSERT_EQ(std::to_string(i), metadata->ToString());
+    ASSERT_OK(stream->Next(&chunk));
+    ASSERT_NE(nullptr, chunk.data);
+    ASSERT_NE(nullptr, chunk.app_metadata);
+    ASSERT_BATCHES_EQUAL(*expected_batches[i], *chunk.data);
+    ASSERT_EQ(std::to_string(i), chunk.app_metadata->ToString());
   }
-  ASSERT_OK(stream->ReadNext(&chunk));
-  ASSERT_EQ(nullptr, chunk);
+  ASSERT_OK(stream->Next(&chunk));
+  ASSERT_EQ(nullptr, chunk.data);
 }
 
 TEST_F(TestMetadata, DoPut) {

cpp/src/arrow/flight/server.cc

@@ -137,19 +137,10 @@ class FlightMessageReaderImpl : public FlightMessageReader {
 
   std::shared_ptr<Schema> schema() const override { return batch_reader_->schema(); }
 
-  Status ReadNext(std::shared_ptr<RecordBatch>* out) override {
-    return ReadWithMetadata(out, nullptr);
-  }
-
-  Status ReadWithMetadata(std::shared_ptr<RecordBatch>* out,
-                          std::shared_ptr<Buffer>* app_metadata) override {
-    if (app_metadata) {
-      *app_metadata = nullptr;
-    }
-    RETURN_NOT_OK(batch_reader_->ReadNext(out));
-    if (app_metadata) {
-      *app_metadata = std::move(last_metadata_);
-    }
+  Status Next(FlightStreamChunk* out) override {
+    out->app_metadata = nullptr;
+    RETURN_NOT_OK(batch_reader_->ReadNext(&out->data));
+    out->app_metadata = std::move(last_metadata_);
     return Status::OK();
   }
 

cpp/src/arrow/flight/test-integration-client.cc

@@ -50,17 +50,16 @@ arrow::Status ReadToTable(arrow::flight::MetadataRecordBatchReader& reader,
   // For integration testing, we expect the server numbers the
   // batches, to test the application metadata part of the spec.
   std::vector<std::shared_ptr<arrow::RecordBatch>> retrieved_chunks;
-  std::shared_ptr<arrow::RecordBatch> chunk;
-  std::shared_ptr<arrow::Buffer> metadata_chunk;
+  arrow::flight::FlightStreamChunk chunk;
   int counter = 0;
   while (true) {
-    RETURN_NOT_OK(reader.ReadWithMetadata(&chunk, &metadata_chunk));
-    if (chunk == nullptr) break;
-    retrieved_chunks.push_back(chunk);
-    if (std::to_string(counter) != metadata_chunk->ToString()) {
+    RETURN_NOT_OK(reader.Next(&chunk));
+    if (!chunk.data) break;
+    retrieved_chunks.push_back(chunk.data);
+    if (std::to_string(counter) != chunk.app_metadata->ToString()) {
       return arrow::Status::Invalid(
           "Expected metadata value: " + std::to_string(counter) +
-          " but got: " + metadata_chunk->ToString());
+          " but got: " + chunk.app_metadata->ToString());
     }
     counter++;
   }

cpp/src/arrow/flight/test-integration-server.cc

@@ -100,14 +100,13 @@ class FlightIntegrationTestServer : public FlightServerBase {
     std::string key = descriptor.path[0];
 
     std::vector<std::shared_ptr<arrow::RecordBatch>> retrieved_chunks;
-    std::shared_ptr<arrow::RecordBatch> chunk;
-    std::shared_ptr<arrow::Buffer> metadata;
+    arrow::flight::FlightStreamChunk chunk;
     while (true) {
-      RETURN_NOT_OK(reader->ReadWithMetadata(&chunk, &metadata));
-      if (chunk == nullptr) break;
-      retrieved_chunks.push_back(chunk);
-      if (metadata) {
-        RETURN_NOT_OK(writer->WriteMetadata(*metadata));
+      RETURN_NOT_OK(reader->Next(&chunk));
+      if (chunk.data == nullptr) break;
+      retrieved_chunks.push_back(chunk.data);
+      if (chunk.app_metadata) {
+        RETURN_NOT_OK(writer->WriteMetadata(*chunk.app_metadata));
       }
     }
     std::shared_ptr<arrow::Table> retrieved_data;

cpp/src/arrow/flight/types.cc

@@ -25,6 +25,7 @@
 #include "arrow/ipc/dictionary.h"
 #include "arrow/ipc/reader.h"
 #include "arrow/status.h"
+#include "arrow/table.h"
 #include "arrow/util/uri.h"
 
 namespace arrow {
@@ -122,6 +123,24 @@ bool Location::Equals(const Location& other) const {
   return ToString() == other.ToString();
 }
 
+Status MetadataRecordBatchReader::ReadAll(
+    std::vector<std::shared_ptr<RecordBatch>>* batches) {
+  FlightStreamChunk chunk;
+
+  while (true) {
+    RETURN_NOT_OK(Next(&chunk));
+    if (!chunk.data) break;
+    batches->emplace_back(std::move(chunk.data));
+  }
+  return Status::OK();
+}
+
+Status MetadataRecordBatchReader::ReadAll(std::shared_ptr<Table>* table) {
+  std::vector<std::shared_ptr<RecordBatch>> batches;
+  RETURN_NOT_OK(ReadAll(&batches));
+  return Table::FromRecordBatches(schema(), batches, table);
+}
+
 SimpleFlightListing::SimpleFlightListing(const std::vector<FlightInfo>& flights)
     : position_(0), flights_(flights) {}
 

cpp/src/arrow/flight/types.h

@@ -27,14 +27,15 @@
 #include <vector>
 
 #include "arrow/flight/visibility.h"
-#include "arrow/ipc/reader.h"
 #include "arrow/ipc/writer.h"
 
 namespace arrow {
 
 class Buffer;
+class RecordBatch;
 class Schema;
 class Status;
+class Table;
 
 namespace ipc {
 
@@ -280,27 +281,30 @@ class ARROW_FLIGHT_EXPORT ResultStream {
   virtual Status Next(std::unique_ptr<Result>* info) = 0;
 };
 
-// Silence warning
-// "non dll-interface class RecordBatchReader used as base for dll-interface class"
-#ifdef _MSC_VER
-#pragma warning(push)
-#pragma warning(disable : 4275)
-#endif
+/// \brief A holder for a RecordBatch with associated Flight metadata.
+struct ARROW_FLIGHT_EXPORT FlightStreamChunk {
+ public:
+  std::shared_ptr<RecordBatch> data;
+  std::shared_ptr<Buffer> app_metadata;
+};
 
-/// \brief A RecordBatchReader that also exposes application-defined
-/// metadata from the Flight protocol.
-class ARROW_FLIGHT_EXPORT MetadataRecordBatchReader : public RecordBatchReader {
+/// \brief An interface to read Flight data with metadata.
+class ARROW_FLIGHT_EXPORT MetadataRecordBatchReader {
  public:
   virtual ~MetadataRecordBatchReader() = default;
 
-  virtual Status ReadWithMetadata(std::shared_ptr<RecordBatch>* out,
-                                  std::shared_ptr<Buffer>* app_metadata) = 0;
+  /// \brief Get the schema for this stream.
+  virtual std::shared_ptr<Schema> schema() const = 0;
+  /// \brief Get the next message from Flight. If the stream is
+  /// finished, then the members of \a FlightStreamChunk will be
+  /// nullptr.
+  virtual Status Next(FlightStreamChunk* next) = 0;
+  /// \brief Consume entire stream as a vector of record batches
+  virtual Status ReadAll(std::vector<std::shared_ptr<RecordBatch>>* batches);
+  /// \brief Consume entire stream as a Table
+  virtual Status ReadAll(std::shared_ptr<Table>* table);
 };
 
-#ifdef _MSC_VER
-#pragma warning(pop)
-#endif
-
 // \brief Create a FlightListing from a vector of FlightInfo objects. This can
 // be iterated once, then it is consumed
 class ARROW_FLIGHT_EXPORT SimpleFlightListing : public FlightListing {

java/flight/src/main/java/org/apache/arrow/flight/example/integration/IntegrationTestClient.java

@@ -130,8 +130,9 @@ class IntegrationTestClient {
           });
       int counter = 0;
       while (reader.read(root)) {
-        final ArrowBuf metadata = allocator.buffer(0);
-        metadata.writeBytes(Integer.toString(counter).getBytes(StandardCharsets.UTF_8));
+        final byte[] rawMetadata = Integer.toString(counter).getBytes(StandardCharsets.UTF_8);
+        final ArrowBuf metadata = allocator.buffer(rawMetadata.length);
+        metadata.writeBytes(rawMetadata);
         // Transfers ownership of the buffer, so do not release it ourselves
         stream.putNext(metadata);
         jsonLoader.load(unloader.getRecordBatch());

python/pyarrow/_flight.pyx

@@ -450,52 +450,82 @@ cdef class FlightInfo:
         return result
 
 
-_AppMetadataRecordBatch = collections.namedtuple(
-    '_BatchWithMetadata',
-    ['record_batch', 'app_metadata'],
-)
-
-
-class AppMetadataRecordBatch(_AppMetadataRecordBatch):
+cdef class FlightStreamChunk:
     """A RecordBatch with application metadata on the side."""
+    cdef:
+        CFlightStreamChunk chunk
+
+    @property
+    def data(self):
+        if self.chunk.data == NULL:
+            return None
+        return pyarrow_wrap_batch(self.chunk.data)
+
+    @property
+    def app_metadata(self):
+        if self.chunk.app_metadata == NULL:
+            return None
+        return pyarrow_wrap_buffer(self.chunk.app_metadata)
+
+    def __iter__(self):
+        return iter((self.data, self.app_metadata))
 
 
-cdef class MetadataRecordBatchReader(_CRecordBatchReader, _ReadPandasOption):
-    """A RecordBatchReader that also allows reading application metadata."""
+cdef class _MetadataRecordBatchReader:
+    """A reader for Flight streams."""
+
+    # Needs to be separate class so the "real" class can subclass the
+    # pure-Python mixin class
 
     cdef dict __dict__
+    cdef shared_ptr[CMetadataRecordBatchReader] reader
 
     cdef readonly:
         Schema schema
 
-    def read_with_metadata(self):
+
+cdef class MetadataRecordBatchReader(_MetadataRecordBatchReader,
+                                     _ReadPandasOption):
+    """A reader for Flight streams."""
+
+    def __iter__(self):
+        while True:
+            yield self.read_chunk()
+
+    def read_all(self):
+        """Read the entire contents of the stream as a Table."""
+        cdef:
+            shared_ptr[CTable] c_table
+        with nogil:
+            check_status(self.reader.get().ReadAll(&c_table))
+        return pyarrow_wrap_table(c_table)
+
+    def read_chunk(self):
         """Read the next RecordBatch along with any metadata.
 
         Returns
         -------
-        batch : RecordBatch
+        data : RecordBatch
             The next RecordBatch in the stream.
-        metadata : Buffer or None
+        app_metadata : Buffer or None
             Application-specific metadata for the batch as defined by
             Flight.
+
+        Raises
+        ------
+        StopIteration
+            when the stream is finished
         """
         cdef:
-            shared_ptr[CRecordBatch] batch
-            shared_ptr[CBuffer] app_metadata
+            FlightStreamChunk chunk = FlightStreamChunk()
 
         with nogil:
-            check_status((<CMetadataRecordBatchReader*> self.reader.get())
-                         .ReadWithMetadata(&batch, &app_metadata))
+            check_status(self.reader.get().Next(&chunk.chunk))
 
-        if batch.get() == NULL:
+        if chunk.chunk.data == NULL:
             raise StopIteration
 
-        if app_metadata != NULL:
-            metadata = pyarrow_wrap_buffer(app_metadata)
-        else:
-            metadata = None
-
-        return AppMetadataRecordBatch(pyarrow_wrap_batch(batch), metadata)
+        return chunk
 
 
 cdef class FlightStreamReader(MetadataRecordBatchReader):

python/pyarrow/includes/libarrow_flight.pxd

@@ -121,10 +121,16 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
         shared_ptr[CSchema] schema()
         CStatus Next(CFlightPayload*)
 
+    cdef cppclass CFlightStreamChunk" arrow::flight::FlightStreamChunk":
+        CFlightStreamChunk()
+        shared_ptr[CRecordBatch] data
+        shared_ptr[CBuffer] app_metadata
+
     cdef cppclass CMetadataRecordBatchReader \
-            " arrow::flight::MetadataRecordBatchReader"(CRecordBatchReader):
-        CStatus ReadWithMetadata(shared_ptr[CRecordBatch]* out,
-                                 shared_ptr[CBuffer]* app_metadata)
+            " arrow::flight::MetadataRecordBatchReader":
+        shared_ptr[CSchema] schema()
+        CStatus Next(CFlightStreamChunk* out)
+        CStatus ReadAll(shared_ptr[CTable]* table)
 
     cdef cppclass CFlightStreamReader \
             " arrow::flight::FlightStreamReader"(CMetadataRecordBatchReader):

python/pyarrow/tests/test_flight.py

@@ -133,7 +133,7 @@ class MetadataFlightServer(flight.FlightServerBase):
         expected_data = [-10, -5, 0, 5, 10]
         while True:
             try:
-                batch, buf = reader.read_with_metadata()
+                batch, buf = reader.read_chunk()
                 assert batch.equals(pa.RecordBatch.from_arrays(
                     [pa.array([expected_data[counter]])],
                     ['a']
@@ -424,6 +424,17 @@ def test_flight_do_get_ints():
         assert data.equals(table)
 
 
+@pytest.mark.pandas
+def test_do_get_ints_pandas():
+    """Try a simple do_get call."""
+    table = simple_ints_table()
+
+    with flight_server(ConstantFlightServer) as server_location:
+        client = flight.FlightClient.connect(server_location)
+        data = client.do_get(flight.Ticket(b'ints')).read_pandas()
+        assert list(data['some_ints']) == table.column(0).to_pylist()
+
+
 def test_flight_do_get_dicts():
     table = simple_dicts_table()
 
@@ -683,7 +694,7 @@ def test_flight_do_get_metadata():
         idx = 0
         while True:
             try:
-                batch, metadata = reader.read_with_metadata()
+                batch, metadata = reader.read_chunk()
                 batches.append(batch)
                 server_idx, = struct.unpack('<i', metadata.to_pybytes())
                 assert idx == server_idx
@@ -724,7 +735,7 @@ def test_cancel_do_get():
         reader = client.do_get(flight.Ticket(b'ints'))
         reader.cancel()
         with pytest.raises(pa.ArrowIOError, match=".*Cancel.*"):
-            reader.read_next_batch()
+            reader.read_chunk()
 
 
 @pytest.mark.slow
@@ -740,11 +751,11 @@ def test_cancel_do_get_threaded():
         raised_proper_exception = threading.Event()
 
         def block_read():
-            reader.read_next_batch()
+            reader.read_chunk()
             read_first_message.set()
             stream_canceled.wait(timeout=5)
             try:
-                reader.read_next_batch()
+                reader.read_chunk()
             except pa.ArrowIOError:
                 with result_lock:
                     raised_proper_exception.set()