pitrou commented on code in PR #45268:
URL: https://github.com/apache/arrow/pull/45268#discussion_r1942803312
##########
cpp/src/arrow/acero/task_util.cc:
##########
@@ -369,17 +369,25 @@ Status TaskSchedulerImpl::ScheduleMore(size_t thread_id,
int num_tasks_finished)
int group_id = tasks[i].first;
int64_t task_id = tasks[i].second;
RETURN_NOT_OK(schedule_impl_([this, group_id, task_id](size_t thread_id)
-> Status {
- RETURN_NOT_OK(ScheduleMore(thread_id, 1));
-
bool task_group_finished = false;
- RETURN_NOT_OK(ExecuteTask(thread_id, group_id, task_id,
&task_group_finished));
+ // PostExecuteTask must be called later if any error ocurres during task
execution
+ // (including ScheduleMore), so we preserve the status.
+ auto status = [&]() {
Review Comment:
Why the lambda? This does not seem to be needed.
##########
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));
Review Comment:
/Can we annotate non-obvious parameters with their names?
```suggestion
ASSERT_OK(scheduler->StartScheduling(
/*xxx=*/0, [](TaskScheduler::TaskGroupContinuationImpl) { return
Status::OK(); },
/*xxx=*/1, /*xxx=*/true));
```
##########
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:
Hmm, so the `Status::Invalid` that's returned above is never reported?
##########
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);
Review Comment:
```suggestion
int num_abort_calls = 0;
auto abort_cont = [&]() {
++num_abort_calls;
};
scheduler->Abort(abort_cont);
ASSERT_EQ(num_abort_calls, 1);
```
##########
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;
Review Comment:
Perhaps use a counter as suggested above as well?
##########
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));
Review Comment:
```suggestion
ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Task failed",
scheduler->StartTaskGroup(/*xxx=*/0,
task_group, kNumTasks));
```
--
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]
