refactor tests

Author: scott-routledge2

cpp/src/arrow/dataset/file_parquet_test.cc

@@ -17,6 +17,7 @@
 
 #include "arrow/dataset/file_parquet.h"
 
+#include <functional>
 #include <memory>
 #include <thread>
 #include <utility>
@@ -25,6 +26,7 @@
 #include "arrow/compute/api_scalar.h"
 #include "arrow/dataset/dataset_internal.h"
 #include "arrow/dataset/parquet_encryption_config.h"
+#include "arrow/dataset/scanner.h"
 #include "arrow/dataset/test_util_internal.h"
 #include "arrow/io/interfaces.h"
 #include "arrow/io/memory.h"
@@ -133,6 +135,30 @@ class ParquetFormatHelper {
   }
 };
 
+class DelayedBufferReader : public ::arrow::io::BufferReader {
+ public:
+  explicit DelayedBufferReader(const std::shared_ptr<::arrow::Buffer>& buffer)
+      : ::arrow::io::BufferReader(buffer) {}
+
+  ::arrow::Future<std::shared_ptr<Buffer>> ReadAsync(
+      const ::arrow::io::IOContext& io_context, int64_t position,
+      int64_t nbytes) override {
+    read_async_count.fetch_add(1);
+    auto self = std::dynamic_pointer_cast<DelayedBufferReader>(shared_from_this());
+    return DeferNotOk(::arrow::io::internal::SubmitIO(
+        io_context, [self, position, nbytes]() -> Result<std::shared_ptr<Buffer>> {
+          std::this_thread::sleep_for(std::chrono::seconds(1));
+          return self->DoReadAt(position, nbytes);
+        }));
+  }
+
+  std::atomic<int> read_async_count{0};
+};
+
+using OptionsCustomizer =
+    std::function<void(std::shared_ptr<ScanOptions>&,
+                       std::vector<std::shared_ptr<arrow::internal::ThreadPool>>&)>;
+
 class TestParquetFileFormat : public FileFormatFixtureMixin<ParquetFormatHelper> {
  public:
   RecordBatchIterator Batches(Fragment* fragment) {
@@ -183,6 +209,55 @@ class TestParquetFileFormat : public FileFormatFixtureMixin<ParquetFormatHelper>
       EXPECT_EQ(SingleBatch(parquet_fragment.get())->num_rows(), expected + 1);
     }
   }
+
+  void TestMultithreadedRegression(OptionsCustomizer customizer) {
+    // GH-38438: This test is similar to MultithreadedScan, but it try to use self
+    // designed Executor and DelayedBufferReader to mock async execution to make
+    // the state machine more complex.
+
+    auto reader = MakeGeneratedRecordBatch(schema({field("utf8", utf8())}), 10000, 100);
+    ASSERT_OK_AND_ASSIGN(auto buffer, ParquetFormatHelper::Write(reader.get()));
+
+    std::vector<Future<>> completes;
+    std::vector<std::shared_ptr<arrow::internal::ThreadPool>> pools;
+
+    for (int idx = 0; idx < 2; ++idx) {
+      auto buffer_reader = std::make_shared<DelayedBufferReader>(buffer);
+      auto source = std::make_shared<FileSource>(buffer_reader, buffer->size());
+      auto fragment = MakeFragment(*source);
+      std::shared_ptr<Scanner> scanner;
+
+      {
+        auto options = std::make_shared<ScanOptions>();
+        ASSERT_OK_AND_ASSIGN(auto thread_pool, arrow::internal::ThreadPool::Make(1));
+        pools.emplace_back(thread_pool);
+        customizer(options, pools);
+        auto fragment_scan_options = std::make_shared<ParquetFragmentScanOptions>();
+        fragment_scan_options->arrow_reader_properties->set_pre_buffer(true);
+
+        options->fragment_scan_options = fragment_scan_options;
+        ScannerBuilder builder(ArithmeticDatasetFixture::schema(), fragment, options);
+
+        ASSERT_OK(builder.UseThreads(true));
+        ASSERT_OK(builder.BatchSize(10000));
+        ASSERT_OK_AND_ASSIGN(scanner, builder.Finish());
+      }
+
+      ASSERT_OK_AND_ASSIGN(auto batch, scanner->Head(10000));
+      [[maybe_unused]] auto fut = scanner->ScanBatchesUnorderedAsync();
+      // Random ReadAsync calls, generate some futures to make the state machine
+      // more complex.
+      for (int yy = 0; yy < 16; yy++) {
+        completes.emplace_back(
+            buffer_reader->ReadAsync(::arrow::io::IOContext(), 0, 1001));
+      }
+      scanner = nullptr;
+    }
+
+    for (auto& f : completes) {
+      f.Wait();
+    }
+  }
 };
 
 TEST_F(TestParquetFileFormat, InspectFailureWithRelevantError) {
@@ -904,120 +979,25 @@ TEST(TestParquetStatistics, NoNullCount) {
   }
 }
 
-class DelayedBufferReader : public ::arrow::io::BufferReader {
- public:
-  explicit DelayedBufferReader(const std::shared_ptr<::arrow::Buffer>& buffer)
-      : ::arrow::io::BufferReader(buffer) {}
-
-  ::arrow::Future<std::shared_ptr<Buffer>> ReadAsync(
-      const ::arrow::io::IOContext& io_context, int64_t position,
-      int64_t nbytes) override {
-    read_async_count.fetch_add(1);
-    auto self = std::dynamic_pointer_cast<DelayedBufferReader>(shared_from_this());
-    return DeferNotOk(::arrow::io::internal::SubmitIO(
-        io_context, [self, position, nbytes]() -> Result<std::shared_ptr<Buffer>> {
-          std::this_thread::sleep_for(std::chrono::seconds(1));
-          return self->DoReadAt(position, nbytes);
-        }));
-  }
-
-  std::atomic<int> read_async_count{0};
-};
-
 TEST_F(TestParquetFileFormat, MultithreadedScanRegression) {
-  // GH-38438: This test is similar to MultithreadedScan, but it try to use self
-  // designed Executor and DelayedBufferReader to mock async execution to make
-  // the state machine more complex.
-  auto reader = MakeGeneratedRecordBatch(schema({field("utf8", utf8())}), 10000, 100);
-
-  ASSERT_OK_AND_ASSIGN(auto buffer, ParquetFormatHelper::Write(reader.get()));
-
-  std::vector<Future<>> completes;
-  std::vector<std::shared_ptr<arrow::internal::ThreadPool>> pools;
-
-  for (int idx = 0; idx < 2; ++idx) {
-    auto buffer_reader = std::make_shared<DelayedBufferReader>(buffer);
-    auto source = std::make_shared<FileSource>(buffer_reader, buffer->size());
-    auto fragment = MakeFragment(*source);
-    std::shared_ptr<Scanner> scanner;
-
-    {
-      auto options = std::make_shared<ScanOptions>();
-      ASSERT_OK_AND_ASSIGN(auto thread_pool, arrow::internal::ThreadPool::Make(1));
-      pools.emplace_back(thread_pool);
-      options->io_context =
-          ::arrow::io::IOContext(::arrow::default_memory_pool(), pools.back().get());
-      auto fragment_scan_options = std::make_shared<ParquetFragmentScanOptions>();
-      fragment_scan_options->arrow_reader_properties->set_pre_buffer(true);
-
-      options->fragment_scan_options = fragment_scan_options;
-      ScannerBuilder builder(ArithmeticDatasetFixture::schema(), fragment, options);
-
-      ASSERT_OK(builder.UseThreads(true));
-      ASSERT_OK(builder.BatchSize(10000));
-      ASSERT_OK_AND_ASSIGN(scanner, builder.Finish());
-    }
-
-    ASSERT_OK_AND_ASSIGN(auto batch, scanner->Head(10000));
-    [[maybe_unused]] auto fut = scanner->ScanBatchesUnorderedAsync();
-    // Random ReadAsync calls, generate some futures to make the state machine
-    // more complex.
-    for (int yy = 0; yy < 16; yy++) {
-      completes.emplace_back(buffer_reader->ReadAsync(::arrow::io::IOContext(), 0, 1001));
-    }
-    scanner = nullptr;
-  }
-
-  for (auto& f : completes) {
-    f.Wait();
-  }
+  OptionsCustomizer customize_io_context =
+      [](std::shared_ptr<ScanOptions>& options,
+         std::vector<std::shared_ptr<arrow::internal::ThreadPool>>& pools) {
+        options->io_context =
+            ::arrow::io::IOContext(::arrow::default_memory_pool(), pools.back().get());
+      };
+  TestMultithreadedRegression(customize_io_context);
 }
 
 TEST_F(TestParquetFileFormat, MultithreadedComputeRegression) {
   // GH-43694: Test similar situation as MultithreadedScanRegression but with
-  // the exec context instead
-
-  auto reader = MakeGeneratedRecordBatch(schema({field("utf8", utf8())}), 10000, 100);
-  ASSERT_OK_AND_ASSIGN(auto buffer, ParquetFormatHelper::Write(reader.get()));
-
-  std::vector<Future<>> completes;
-  std::vector<std::shared_ptr<arrow::internal::ThreadPool>> pools;
-
-  for (int idx = 0; idx < 2; ++idx) {
-    auto buffer_reader = std::make_shared<DelayedBufferReader>(buffer);
-    auto source = std::make_shared<FileSource>(buffer_reader, buffer->size());
-    auto fragment = MakeFragment(*source);
-    std::shared_ptr<Scanner> scanner;
-
-    {
-      auto options = std::make_shared<ScanOptions>();
-      ASSERT_OK_AND_ASSIGN(auto thread_pool, arrow::internal::ThreadPool::Make(1));
-      pools.emplace_back(thread_pool);
-      options->cpu_executor = pools.back().get();
-      auto fragment_scan_options = std::make_shared<ParquetFragmentScanOptions>();
-      fragment_scan_options->arrow_reader_properties->set_pre_buffer(true);
-
-      options->fragment_scan_options = fragment_scan_options;
-      ScannerBuilder builder(ArithmeticDatasetFixture::schema(), fragment, options);
-
-      ASSERT_OK(builder.UseThreads(true));
-      ASSERT_OK(builder.BatchSize(10000));
-      ASSERT_OK_AND_ASSIGN(scanner, builder.Finish());
-    }
-
-    ASSERT_OK_AND_ASSIGN(auto batch, scanner->Head(10000));
-    [[maybe_unused]] auto fut = scanner->ScanBatchesUnorderedAsync();
-    // Random ReadAsync calls, generate some futures to make the state machine
-    // more complex.
-    for (int yy = 0; yy < 16; yy++) {
-      completes.emplace_back(buffer_reader->ReadAsync(::arrow::io::IOContext(), 0, 1001));
-    }
-    scanner = nullptr;
-  }
-
-  for (auto& f : completes) {
-    f.Wait();
-  }
+  // the customized CPU executor instead
+  OptionsCustomizer customize_cpu_executor =
+      [](std::shared_ptr<ScanOptions>& options,
+         std::vector<std::shared_ptr<arrow::internal::ThreadPool>>& pools) {
+        options->cpu_executor = pools.back().get();
+      };
+  TestMultithreadedRegression(customize_cpu_executor);
 }
 
 }  // namespace dataset

cpp/src/arrow/dataset/scanner.h

@@ -107,8 +107,10 @@ struct ARROW_DS_EXPORT ScanOptions {
 
   /// Executor for any CPU tasks
   ///
+  /// If null, the global CPU executor will be used
+  ///
   /// Note: The Executor will be ignored if use_threads is set to false
-  arrow::internal::Executor* cpu_executor = nullptr;
+  arrow::internal::Executor* cpu_executor = NULLPTR;
 
   /// If true the scanner will scan in parallel
   ///

cpp/src/arrow/util/thread_pool.h

@@ -19,7 +19,6 @@
 
 #include <cstdint>
 #include <memory>
-#include <optional>
 #include <queue>
 #include <type_traits>
 #include <unordered_set>