create reader from iter (apache#94)

zhouyuan · Apr 15, 2022 · e16c3db · e16c3db
1 parent a4b67ef
commit e16c3db
Show file tree

Hide file tree

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