| 79 | } |
| 80 | |
| 81 | void CheckRunOutput(const std::vector<BatchesWithSchema>& batches, |
| 82 | const BatchesWithSchema& exp_batches, bool parallel = false) { |
| 83 | SCOPED_TRACE(parallel ? "parallel" : "single threaded"); |
| 84 | |
| 85 | ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make()); |
| 86 | |
| 87 | Declaration union_decl{"union", ExecNodeOptions{}}; |
| 88 | |
| 89 | for (const auto& batch : batches) { |
| 90 | union_decl.inputs.emplace_back(Declaration{ |
| 91 | "source", SourceNodeOptions{batch.schema, batch.gen(parallel, |
| 92 | /*slow=*/false)}}); |
| 93 | } |
| 94 | AsyncGenerator<std::optional<ExecBatch>> sink_gen; |
| 95 | |
| 96 | // Test UnionNode::Make with zero inputs |
| 97 | if (batches.size() == 0) { |
| 98 | ASSERT_RAISES(Invalid, Declaration::Sequence( |
| 99 | {union_decl, {"sink", SinkNodeOptions{&sink_gen}}}) |
| 100 | .AddToPlan(plan.get())); |
| 101 | return; |
| 102 | } else { |
| 103 | ASSERT_OK(Declaration::Sequence({union_decl, {"sink", SinkNodeOptions{&sink_gen}}}) |
| 104 | .AddToPlan(plan.get())); |
| 105 | } |
| 106 | |
| 107 | Future<std::vector<ExecBatch>> actual = StartAndCollect(plan.get(), sink_gen); |
| 108 | |
| 109 | auto expected_matcher = |
| 110 | Finishes(ResultWith(UnorderedElementsAreArray(exp_batches.batches))); |
| 111 | ASSERT_THAT(actual, expected_matcher); |
| 112 | |
| 113 | // union node with multiple inputs should produce unordered batches |
| 114 | if (batches.size() > 1) { |
| 115 | for (const auto& batch : *actual.result()) { |
| 116 | ASSERT_EQ(batch.index, compute::kUnsequencedIndex); |
| 117 | } |
| 118 | } |
| 119 | } |
| 120 | |
| 121 | void CheckUnionExecNode(size_t num_input_nodes, size_t num_batches, bool parallel) { |
| 122 | auto random_schema = GenerateRandomSchema(num_input_nodes); |
nothing calls this directly
no test coverage detected