From 9d9c464ce6883f52aaca9f913eec4cd50006c767 Mon Sep 17 00:00:00 2001
From: avogar <avogar@clickhouse.com>
Date: Wed, 9 Aug 2023 07:05:31 +0000
Subject: [PATCH] Allow to get number of rows in record batch

---
 cpp/src/arrow/ipc/reader.cc | 17 +++++++++++++++++
 cpp/src/arrow/ipc/reader.h  |  2 ++
 2 files changed, 19 insertions(+)

diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc
index eadba69f05612..241dbd1523d2e 100644
--- a/cpp/src/arrow/ipc/reader.cc
+++ b/cpp/src/arrow/ipc/reader.cc
@@ -1305,6 +1305,23 @@ class RecordBatchFileReaderImpl : public RecordBatchFileReader {
     return total;
   }
 
+  Result<int64_t> RecordBatchCountRows(int i) override {
+    DCHECK_GE(i, 0);
+    DCHECK_LT(i, num_record_batches());
+    ARROW_ASSIGN_OR_RAISE(auto outer_message,
+                          ReadMessageFromBlock(GetRecordBatchBlock(i)));
+    auto metadata = outer_message->metadata();
+    const flatbuf::Message* message = nullptr;
+    RETURN_NOT_OK(
+        internal::VerifyMessage(metadata->data(), metadata->size(), &message));
+    auto batch = message->header_as_RecordBatch();
+    if (batch == nullptr) {
+      return Status::IOError(
+          "Header-type of flatbuffer-encoded Message is not RecordBatch.");
+    }
+    return batch->length();
+  }
+
   Status Open(const std::shared_ptr<io::RandomAccessFile>& file, int64_t footer_offset,
               const IpcReadOptions& options) {
     owned_file_ = file;
diff --git a/cpp/src/arrow/ipc/reader.h b/cpp/src/arrow/ipc/reader.h
index 1698abd14b48f..80c660f0d5cf7 100644
--- a/cpp/src/arrow/ipc/reader.h
+++ b/cpp/src/arrow/ipc/reader.h
@@ -203,6 +203,8 @@ class ARROW_EXPORT RecordBatchFileReader
   /// \brief Computes the total number of rows in the file.
   virtual Result<int64_t> CountRows() = 0;
 
+  virtual Result<int64_t> RecordBatchCountRows(int i) = 0;
+
   /// \brief Begin loading metadata for the desired batches into memory.
   ///
   /// This method will also begin loading all dictionaries messages into memory.