Merge pull request #39 from ClickHouse/count-from-record-batch

Allow to get number of rows in record batch
ClickHouse · Aug 9, 2023 · 1d93838 · 1d93838
2 parents 1f1b3d3 + 802ff81
commit 1d93838
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
@@ -1240,6 +1240,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.