zanmato1984 commented on code in PR #45268:
URL: https://github.com/apache/arrow/pull/45268#discussion_r1948436541
##########
cpp/src/arrow/acero/task_util_test.cc:
##########
@@ -231,5 +231,96 @@ TEST(TaskScheduler, StressTwo) {
}
}
+TEST(TaskScheduler, AbortContOnTaskErrorSerial) {
+ constexpr int kNumTasks = 16;
+
+ auto scheduler = TaskScheduler::Make();
+ auto task = [&](std::size_t, int64_t task_id) {
+ if (task_id == kNumTasks / 2) {
+ return Status::Invalid("Task failed");
+ }
+ return Status::OK();
+ };
+
+ int task_group =
+ scheduler->RegisterTaskGroup(task, [](std::size_t) { return
Status::OK(); });
+ scheduler->RegisterEnd();
+
+ ASSERT_OK(scheduler->StartScheduling(
+ 0, [](TaskScheduler::TaskGroupContinuationImpl) { return Status::OK();
}, 1, true));
+ ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Task failed",
+ scheduler->StartTaskGroup(0, task_group,
kNumTasks));
+
+ bool abort_cont_called = false;
+ auto abort_cont = [&]() {
+ ASSERT_FALSE(abort_cont_called);
+ abort_cont_called = true;
+ };
+
+ scheduler->Abort(abort_cont);
+
+ ASSERT_TRUE(abort_cont_called);
+}
+
+TEST(TaskScheduler, AbortContOnTaskErrorParallel) {
+#ifndef ARROW_ENABLE_THREADING
+ GTEST_SKIP() << "Test requires threading support";
+#endif
+ constexpr int kNumThreads = 16;
+
+ ThreadIndexer thread_indexer;
+ int num_threads = std::min(static_cast<int>(thread_indexer.Capacity()),
kNumThreads);
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<ThreadPool> thread_pool,
+ MakePrimedThreadPool(num_threads));
+ TaskScheduler::ScheduleImpl schedule =
+ [&](TaskScheduler::TaskGroupContinuationImpl task) {
+ return thread_pool->Spawn([&, task] {
+ std::size_t thread_id = thread_indexer();
+ auto status = task(thread_id);
+ ASSERT_TRUE(status.ok() || status.IsInvalid() ||
status.IsCancelled());
+ });
+ };
+
+ for (int num_tasks :
+ {2, num_threads - 1, num_threads, num_threads + 1, 2 * num_threads}) {
+ ARROW_SCOPED_TRACE("num_tasks = ", num_tasks);
+ for (int num_concurrent_tasks :
+ {1, num_tasks - 1, num_tasks, num_tasks + 1, 2 * num_tasks}) {
+ ARROW_SCOPED_TRACE("num_concurrent_tasks = ", num_concurrent_tasks);
+ for (int aborting_task_id = 0; aborting_task_id < num_tasks;
++aborting_task_id) {
+ ARROW_SCOPED_TRACE("aborting_task_id = ", aborting_task_id);
+ auto scheduler = TaskScheduler::Make();
+
+ bool abort_cont_called = false;
+ auto abort_cont = [&]() {
+ ASSERT_FALSE(abort_cont_called);
+ abort_cont_called = true;
+ };
+
+ auto task = [&](std::size_t, int64_t task_id) {
+ if (task_id == aborting_task_id) {
+ scheduler->Abort(abort_cont);
+ }
+ if (task_id % 2 == 0) {
+ return Status::Invalid("Task failed");
+ }
+ return Status::OK();
+ };
+
+ int task_group =
+ scheduler->RegisterTaskGroup(task, [](std::size_t) { return
Status::OK(); });
+ scheduler->RegisterEnd();
+
+ ASSERT_OK(scheduler->StartScheduling(0, schedule,
num_concurrent_tasks, false));
+ ASSERT_OK(scheduler->StartTaskGroup(0, task_group, num_tasks));
Review Comment:
At least not here. This `StartTaskGroup` merely schedules the parallel tasks
into the underlying executor (the raw thread pool in this specific case, and
the more sophisticated `AsyncScheduler` in official acero code path). The
status returned here is merely if the tasks are successfully scheduled. The
error of a specific task will be propagated thru the abort continuation of this
`TaskScheduler`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
