ARROW-8299: [C++] Reusable "optional ParallelFor" function for option…

…al use of multithreading We often see code like ``` if (use_threads) { return ::arrow::internal::ParallelFor(n, Func); } else { for (size_t i = 0; i < n; ++i) { RETURN_NOT_OK(Func(i)); } return Status::OK(); ``` This patch adds a helper function to do this. Closes #6870 from liyafan82/fly_0408_pf Authored-by: liyafan82 <fan_li_ya@foxmail.com> Signed-off-by: Antoine Pitrou <antoine@python.org>
apache · Apr 8, 2020 · e39326b · e39326b
1 parent c49c47d
commit e39326b
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 35 deletions.
diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc
@@ -384,15 +384,8 @@ Status DecompressBuffers(Compression::type compression, const IpcReadOptions& op
     return Status::OK();
   };
 
-  if (options.use_threads) {
-    return ::arrow::internal::ParallelFor(static_cast<int>(fields->size()),
-                                          DecompressOne);
-  } else {
-    for (int i = 0; i < static_cast<int>(fields->size()); ++i) {
-      RETURN_NOT_OK(DecompressOne(i));
-    }
-    return Status::OK();
-  }
+  return ::arrow::internal::OptionalParallelFor(
+      options.use_threads, static_cast<int>(fields->size()), DecompressOne);
 }
 
 Result<std::shared_ptr<RecordBatch>> LoadRecordBatchSubset(

diff --git a/cpp/src/arrow/ipc/writer.cc b/cpp/src/arrow/ipc/writer.cc
@@ -199,15 +199,8 @@ class RecordBatchSerializer {
       return Status::OK();
     };
 
-    if (options_.use_threads) {
-      return ::arrow::internal::ParallelFor(static_cast<int>(out_->body_buffers.size()),
-                                            CompressOne);
-    } else {
-      for (size_t i = 0; i < out_->body_buffers.size(); ++i) {
-        RETURN_NOT_OK(CompressOne(i));
-      }
-      return Status::OK();
-    }
+    return ::arrow::internal::OptionalParallelFor(
+        options_.use_threads, static_cast<int>(out_->body_buffers.size()), CompressOne);
   }
 
   Status Assemble(const RecordBatch& batch) {

diff --git a/cpp/src/arrow/python/arrow_to_pandas.cc b/cpp/src/arrow/python/arrow_to_pandas.cc
@@ -63,7 +63,7 @@ namespace arrow {
 class MemoryPool;
 
 using internal::checked_cast;
-using internal::ParallelFor;
+using internal::OptionalParallelFor;
 
 // ----------------------------------------------------------------------
 // PyCapsule code for setting ndarray base to reference C++ object
@@ -1925,14 +1925,7 @@ class ConsolidatedBlockCreator : public PandasBlockCreator {
       return block->Write(std::move(arrays_[i]), i, this->column_block_placement_[i]);
     };
 
-    if (options_.use_threads) {
-      return ParallelFor(num_columns_, WriteColumn);
-    } else {
-      for (int i = 0; i < num_columns_; ++i) {
-        RETURN_NOT_OK(WriteColumn(i));
-      }
-      return Status::OK();
-    }
+    return OptionalParallelFor(options_.use_threads, num_columns_, WriteColumn);
   }
 
  private:
@@ -2032,14 +2025,8 @@ Status ConvertCategoricals(const PandasOptions& options,
       }
     }
   }
-  if (options.use_threads) {
-    return ParallelFor(static_cast<int>(columns_to_encode.size()), EncodeColumn);
-  } else {
-    for (auto i : columns_to_encode) {
-      RETURN_NOT_OK(EncodeColumn(i));
-    }
-    return Status::OK();
-  }
+  return OptionalParallelFor(options.use_threads,
+                             static_cast<int>(columns_to_encode.size()), EncodeColumn);
 }
 
 Status ConvertArrayToPandas(const PandasOptions& options, std::shared_ptr<Array> arr,

diff --git a/cpp/src/arrow/util/parallel.h b/cpp/src/arrow/util/parallel.h
@@ -44,5 +44,21 @@ Status ParallelFor(int num_tasks, FUNCTION&& func) {
   return st;
 }
 
+// A parallelizer that takes a `Status(int)` function and calls it with
+// arguments between 0 and `num_tasks - 1`, in sequence or in parallel,
+// depending on the input boolean.
+
+template <class FUNCTION>
+Status OptionalParallelFor(bool use_threads, int num_tasks, FUNCTION&& func) {
+  if (use_threads) {
+    return ParallelFor(num_tasks, std::forward<FUNCTION>(func));
+  } else {
+    for (int i = 0; i < num_tasks; ++i) {
+      RETURN_NOT_OK(func(i));
+    }
+    return Status::OK();
+  }
+}
+
 }  // namespace internal
 }  // namespace arrow