diff --git a/cpp/src/arrow/dataset/dataset.cc b/cpp/src/arrow/dataset/dataset.cc index 597c0a256c098..a24b283de4c4f 100644 --- a/cpp/src/arrow/dataset/dataset.cc +++ b/cpp/src/arrow/dataset/dataset.cc @@ -118,11 +118,22 @@ InMemorySource::InMemorySource(std::shared_ptr schema, : Source(std::move(schema)), get_batches_(VectorRecordBatchGenerator(std::move(batches))) {} +struct TableRecordBatchGenerator { + explicit TableRecordBatchGenerator(std::shared_ptr table) + : table_(std::move(table)) {} + + RecordBatchIterator operator()() const { + auto reader = std::make_shared(*table_); + auto table = table_; + return MakeFunctionIterator([reader, table] { return reader->Next(); }); + } + + std::shared_ptr
table_; +}; + InMemorySource::InMemorySource(std::shared_ptr
table) - : Source(table->schema()), get_batches_([table] { - auto reader = std::make_shared(*table); - return MakeFunctionIterator([reader, table] { return reader->Next(); }); - }) {} + : Source(table->schema()), + get_batches_(TableRecordBatchGenerator(std::move(table))) {} FragmentIterator InMemorySource::GetFragmentsImpl( std::shared_ptr scan_options) {