Repository: apache/arrow
File: cpp/src/arrow/dataset/file_parquet_test.cc (C++)

Code before:
                                 int64_t expected_rows, int64_t expected_batches) {
    return CountRowsAndBatchesInScan(fragment.get(), expected_rows, expected_batches);
  }

  void CountRowGroupsInFragment(const std::shared_ptr<Fragment>& fragment,
                                std::vector<int> expected_row_groups,
                                compute::Expression filter) {
    SetFilter(filter);

    auto parquet_fragment = checked_pointer_cast<ParquetFileFragment>(fragment);
    ASSERT_OK_AND_ASSIGN(auto fragments, parquet_fragment->SplitByRowGroup(opts_->filter))

    EXPECT_EQ(fragments.size(), expected_row_groups.size());
    for (size_t i = 0; i < fragments.size(); i++) {
      auto expected = expected_row_groups[i];
      auto parquet_fragment = checked_pointer_cast<ParquetFileFragment>(fragments[i]);

      EXPECT_EQ(parquet_fragment->row_groups(), std::vector<int>{expected});
      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));

Review comment (suggestion):
These comments should probably remain in `TestParquetFileFormat.MultithreadedScanRegression` instead of here.

Code after:
    return CountRowsAndBatchesInScan(fragment.get(), expected_rows, expected_batches);
  }

  void CountRowGroupsInFragment(const std::shared_ptr<Fragment>& fragment,
                                std::vector<int> expected_row_groups,
                                compute::Expression filter) {
    SetFilter(filter);

    auto parquet_fragment = checked_pointer_cast<ParquetFileFragment>(fragment);
    ASSERT_OK_AND_ASSIGN(auto fragments, parquet_fragment->SplitByRowGroup(opts_->filter))

    EXPECT_EQ(fragments.size(), expected_row_groups.size());
    for (size_t i = 0; i < fragments.size(); i++) {
      auto expected = expected_row_groups[i];
      auto parquet_fragment = checked_pointer_cast<ParquetFileFragment>(fragments[i]);

      EXPECT_EQ(parquet_fragment->row_groups(), std::vector<int>{expected});
      EXPECT_EQ(SingleBatch(parquet_fragment.get())->num_rows(), expected + 1);
    }
  }

  void TestMultithreadedRegression(CustomizeScanOptionsWithThreadPool customizer) {
    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.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));