ARROW-439: [Python] Add option in "to_pandas" conversions to yield Categorical from...
[arrow.git] / cpp / src / arrow / util / parallel.h
1 // Licensed to the Apache Software Foundation (ASF) under one
2 // or more contributor license agreements. See the NOTICE file
3 // distributed with this work for additional information
4 // regarding copyright ownership. The ASF licenses this file
5 // to you under the Apache License, Version 2.0 (the
6 // "License"); you may not use this file except in compliance
7 // with the License. You may obtain a copy of the License at
8 //
9 // http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing,
12 // software distributed under the License is distributed on an
13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 // KIND, either express or implied. See the License for the
15 // specific language governing permissions and limitations
16 // under the License.
17
18 #ifndef ARROW_UTIL_PARALLEL_H
19 #define ARROW_UTIL_PARALLEL_H
20
21 #include <atomic>
22 #include <mutex>
23 #include <thread>
24 #include <vector>
25
26 #include "arrow/status.h"
27
28 namespace arrow {
29
30 template <class FUNCTION>
31 Status ParallelFor(int nthreads, int num_tasks, FUNCTION&& func) {
32 std::vector<std::thread> thread_pool;
33 thread_pool.reserve(nthreads);
34 std::atomic<int> task_counter(0);
35
36 std::mutex error_mtx;
37 bool error_occurred = false;
38 Status error;
39
40 for (int thread_id = 0; thread_id < nthreads; ++thread_id) {
41 thread_pool.emplace_back(
42 [&num_tasks, &task_counter, &error, &error_occurred, &error_mtx, &func]() {
43 int task_id;
44 while (!error_occurred) {
45 task_id = task_counter.fetch_add(1);
46 if (task_id >= num_tasks) {
47 break;
48 }
49 Status s = func(task_id);
50 if (!s.ok()) {
51 std::lock_guard<std::mutex> lock(error_mtx);
52 error_occurred = true;
53 error = s;
54 break;
55 }
56 }
57 });
58 }
59 for (auto&& thread : thread_pool) {
60 thread.join();
61 }
62 if (error_occurred) {
63 return error;
64 }
65 return Status::OK();
66 }
67
68 } // namespace arrow
69
70 #endif