From e16c3db5d7ead3b9ce8af9307b798b93370984b7 Mon Sep 17 00:00:00 2001
From: Rui Mo <rui.mo@intel.com>
Date: Fri, 15 Apr 2022 17:13:51 +0800
Subject: [PATCH] create reader from iter (#94)

---
 cpp/src/arrow/record_batch.cc | 9 +++++++++
 cpp/src/arrow/record_batch.h  | 7 +++++++
 2 files changed, 16 insertions(+)

diff --git a/cpp/src/arrow/record_batch.cc b/cpp/src/arrow/record_batch.cc
index 045f5eb829b4d..be81c4646d053 100644
--- a/cpp/src/arrow/record_batch.cc
+++ b/cpp/src/arrow/record_batch.cc
@@ -382,4 +382,13 @@ Result<std::shared_ptr<RecordBatchReader>> RecordBatchReader::Make(
   return std::make_shared<SimpleRecordBatchReader>(std::move(batches), schema);
 }
 
+Result<std::shared_ptr<RecordBatchReader>> RecordBatchReader::Make(
+    Iterator<std::shared_ptr<RecordBatch>> it, std::shared_ptr<Schema> schema) {
+  if (schema == nullptr) {
+      return Status::Invalid("Cannot infer schema from nullptr");
+  }
+
+  return std::make_shared<SimpleRecordBatchReader>(std::move(it), schema);
+}
+
 }  // namespace arrow
diff --git a/cpp/src/arrow/record_batch.h b/cpp/src/arrow/record_batch.h
index 79510d60452e6..368aaf231e625 100644
--- a/cpp/src/arrow/record_batch.h
+++ b/cpp/src/arrow/record_batch.h
@@ -309,6 +309,13 @@ class ARROW_EXPORT RecordBatchReader {
   ///            element if not provided.
   static Result<std::shared_ptr<RecordBatchReader>> Make(
       RecordBatchVector batches, std::shared_ptr<Schema> schema = NULLPTR);
+
+  /// \brief Create a RecordBatchReader from an iterator of RecordBatch.
+  ///
+  /// \param[in] iterator of RecordBatch to read from
+  /// \param[in] schema schema to conform to. A valid schema should be provided.
+  static Result<std::shared_ptr<RecordBatchReader>> Make(
+      Iterator<std::shared_ptr<RecordBatch>> it, std::shared_ptr<Schema> schema);
 };
 
 }  // namespace arrow