| 674 | } |
| 675 | |
| 676 | Future<BatchesWithCommonSchema> DeclarationToExecBatchesImpl( |
| 677 | Declaration declaration, QueryOptions options, |
| 678 | ::arrow::internal::Executor* cpu_executor) { |
| 679 | std::shared_ptr<Schema> out_schema; |
| 680 | AsyncGenerator<std::optional<ExecBatch>> sink_gen; |
| 681 | ExecContext exec_ctx(options.memory_pool, cpu_executor, options.function_registry); |
| 682 | ARROW_ASSIGN_OR_RAISE(std::shared_ptr<ExecPlan> exec_plan, |
| 683 | ExecPlan::Make(options, exec_ctx)); |
| 684 | SinkNodeOptions sink_options(&sink_gen, &out_schema); |
| 685 | sink_options.sequence_output = options.sequence_output; |
| 686 | Declaration with_sink = Declaration::Sequence({declaration, {"sink", sink_options}}); |
| 687 | ARROW_RETURN_NOT_OK(with_sink.AddToPlan(exec_plan.get())); |
| 688 | if (!options.field_names.empty()) { |
| 689 | ARROW_ASSIGN_OR_RAISE(out_schema, |
| 690 | out_schema->WithNames(std::move(options.field_names))); |
| 691 | } |
| 692 | ARROW_RETURN_NOT_OK(exec_plan->Validate()); |
| 693 | exec_plan->StartProducing(); |
| 694 | auto collected_fut = CollectAsyncGenerator(sink_gen); |
| 695 | return exec_plan->finished().Then( |
| 696 | [collected_fut, exec_plan, |
| 697 | schema = std::move(out_schema)]() -> Result<BatchesWithCommonSchema> { |
| 698 | if (!collected_fut.is_finished()) { |
| 699 | return Status::Invalid( |
| 700 | "Plan finished but it did not emit the expected number of batches."); |
| 701 | } |
| 702 | ARROW_ASSIGN_OR_RAISE(auto collected, collected_fut.result()); |
| 703 | std::vector<ExecBatch> exec_batches = ::arrow::internal::MapVector( |
| 704 | [](std::optional<ExecBatch> batch) { return batch.value_or(ExecBatch()); }, |
| 705 | std::move(collected)); |
| 706 | return BatchesWithCommonSchema{std::move(exec_batches), schema}; |
| 707 | }); |
| 708 | } |
| 709 | |
| 710 | Future<> DeclarationToStatusImpl(Declaration declaration, QueryOptions options, |
| 711 | ::arrow::internal::Executor* cpu_executor) { |
no test coverage detected