Allow to get number of rows in record batch

ClickHouse · Sep 19, 2023 · 9d9c464 · 9d9c464
1 parent 36195a8
commit 9d9c464
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git 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
@@ -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.