Apply review feedback

kou · Jun 29, 2023 · 3172ba8 · 3172ba8
1 parent c94dbed
commit 3172ba8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc
@@ -1018,8 +1018,9 @@ class RecordBatchStreamReaderImpl : public RecordBatchStreamReader,
   }
 
   Result<RecordBatchWithMetadata> ReadNext() override {
-    auto collect_listener = static_cast<CollectListener*>(raw_listener());
-    while (collect_listener->num_record_batches() == 0) {
+    auto collect_listener = checked_cast<CollectListener*>(raw_listener());
+    while (collect_listener->num_record_batches() == 0 &&
+           state() != StreamDecoderInternal::State::EOS) {
       ARROW_ASSIGN_OR_RAISE(auto message, message_reader_->ReadNextMessage());
       if (!message) {  // End of stream
         if (state() == StreamDecoderInternal::State::INITIAL_DICTIONARIES) {

diff --git a/cpp/src/arrow/ipc/reader.h b/cpp/src/arrow/ipc/reader.h
@@ -313,15 +313,17 @@ class ARROW_EXPORT CollectListener : public Listener {
   std::shared_ptr<Schema> schema() const { return schema_; }
 
   /// \return the all decoded record batches
-  std::vector<std::shared_ptr<RecordBatch>> record_batches() const {
+  const std::vector<std::shared_ptr<RecordBatch>>& record_batches() const {
     return record_batches_;
   }
 
   /// \return the all decoded metadatas
-  std::vector<std::shared_ptr<KeyValueMetadata>> metadatas() const { return metadatas_; }
+  const std::vector<std::shared_ptr<KeyValueMetadata>>& metadatas() const {
+    return metadatas_;
+  }
 
   /// \return the number of collected record batches
-  size_t num_record_batches() const { return record_batches_.size(); }
+  int64_t num_record_batches() const { return record_batches_.size(); }
 
   /// \return the last decoded record batch and remove it from
   /// record_batches