| 1031 | } // namespace |
| 1032 | |
| 1033 | Result<std::unique_ptr<RecordBatchReader>> DeclarationToReader(Declaration declaration, |
| 1034 | QueryOptions options) { |
| 1035 | if (options.custom_cpu_executor != nullptr) { |
| 1036 | return Status::Invalid("Cannot use synchronous methods with a custom CPU executor"); |
| 1037 | } |
| 1038 | std::shared_ptr<Schema> schema; |
| 1039 | std::shared_ptr<ExecPlan> plan; |
| 1040 | auto batch_iterator = std::make_unique<Iterator<std::shared_ptr<RecordBatch>>>( |
| 1041 | ::arrow::internal::IterateSynchronously<std::shared_ptr<RecordBatch>>( |
| 1042 | [&](::arrow::internal::Executor* executor) |
| 1043 | -> Result<AsyncGenerator<std::shared_ptr<RecordBatch>>> { |
| 1044 | ExecContext exec_ctx(options.memory_pool, executor, |
| 1045 | options.function_registry); |
| 1046 | return DeclarationToRecordBatchGenerator(declaration, std::move(options), |
| 1047 | executor, &schema, &plan); |
| 1048 | }, |
| 1049 | options.use_threads)); |
| 1050 | |
| 1051 | struct PlanReader : RecordBatchReader { |
| 1052 | PlanReader(std::shared_ptr<ExecPlan> plan, std::shared_ptr<Schema> schema, |
| 1053 | std::unique_ptr<Iterator<std::shared_ptr<RecordBatch>>> iterator) |
| 1054 | : plan_(std::move(plan)), |
| 1055 | schema_(std::move(schema)), |
| 1056 | iterator_(std::move(iterator)) {} |
| 1057 | |
| 1058 | std::shared_ptr<Schema> schema() const override { return schema_; } |
| 1059 | |
| 1060 | Status ReadNext(std::shared_ptr<RecordBatch>* record_batch) override { |
| 1061 | if (!iterator_) { |
| 1062 | return Status::Invalid("call to ReadNext on already closed reader"); |
| 1063 | } |
| 1064 | return iterator_->Next().Value(record_batch); |
| 1065 | } |
| 1066 | |
| 1067 | Status Close() override { |
| 1068 | if (!iterator_) { |
| 1069 | // Already closed |
| 1070 | return Status::OK(); |
| 1071 | } |
| 1072 | // End plan and read from generator until finished |
| 1073 | plan_->StopProducing(); |
| 1074 | std::shared_ptr<RecordBatch> batch; |
| 1075 | do { |
| 1076 | Status st = ReadNext(&batch); |
| 1077 | if (!st.ok()) { |
| 1078 | if (st.IsCancelled()) break; // plan cancelled, so closing is done |
| 1079 | return st; |
| 1080 | } |
| 1081 | } while (batch != nullptr); |
| 1082 | iterator_.reset(); |
| 1083 | return Status::OK(); |
| 1084 | } |
| 1085 | |
| 1086 | std::shared_ptr<ExecPlan> plan_; |
| 1087 | std::shared_ptr<Schema> schema_; |
| 1088 | std::unique_ptr<Iterator<std::shared_ptr<RecordBatch>>> iterator_; |
| 1089 | }; |
| 1090 | |