problem parallelizing reading of multiple parquet files using S3FileSystem

[aavbsouza]

Hello. I am encountering some difficulties to make the S3 filesystem to work in parallel. The problem that I have is that I need to read several hundreds of Parquet files, from each file I need a segment given by a index column on the first stage I read some metadata columns and on a second stage I read a payload from a data column (a LargeList float array).

For the reading of the metadata I am using a dataset that scans only one file. For the the payload I am using a parquet/arrow reader with readRowGroup API, finding the necessary row groups using the metadata/stats of the given parquet file

I have tried some approaches:

• Use an outer parallelization loop over files using oneAPI TBB, for each file I create all the arrow machinery (filesystems, readers, datasets...). With this approach I got the expected parallelism/performance when using the local filesystem, but when using the S3Filesystem I got very low throughput and cpu utilization even tough some of post processing of the data is cpu bound.

• Sharing the S3Filesystem instance. I built outside the parallelization loop a S3FileSystem instance with its associated io_context. That I shared among the tasks of the ONEAPI tbb parallel for loop. This approach is exemplified here:

std::shared_ptr<parquet::ReaderProperties> parquetProperties()
{
    // Configure general Parquet reader settings
    auto reader_properties = std::make_shared<parquet::ReaderProperties>(arrow::default_memory_pool());
    reader_properties->set_buffer_size(64 * 1024 * 1024); // 64 MB buffer size
    reader_properties->enable_buffered_stream();
    return reader_properties;
}

std::shared_ptr<parquet::ArrowReaderProperties> arrowReaderProperties()
{
    auto arrow_reader_props = std::make_shared<parquet::ArrowReaderProperties>();
    arrow_reader_props->set_batch_size(10000); // default 64 * 1024// Configure general Parquet reader settings
    arrow_reader_props->set_use_threads(true); // Enable multithreading
    arrow_reader_props->set_pre_buffer(true);  // Enable pre-buffering
    arrow_reader_props->set_cache_options(arrow::io::CacheOptions::LazyDefaults());
    return arrow_reader_props;
}

arrow::fs::S3Options CreateS3Options()
{
    arrow::fs::S3Options options = arrow::fs::S3Options::Defaults();
    // Set the region from environment variable or default to us-east-1
    options.region = std::getenv("AWS_DEFAULT_REGION") ? std::getenv("AWS_DEFAULT_REGION") : "us-east-1";
    // Increase timeouts for dev container environment
    options.connect_timeout = 15000; // 15 seconds
    options.request_timeout = 60000; // 60 seconds
    // Add retry configuration
    options.retry_strategy = arrow::fs::S3RetryStrategy::GetAwsStandardRetryStrategy(1000);
    return options;
}
arrow::Result<std::shared_ptr<arrow::fs::FileSystem>> createFilesystem(const bool s3_enabled,
                                                                       arrow::io::IOContext io_context) {

    // Create a filesystem based on whether S3 is enabled
    std::shared_ptr<arrow::fs::FileSystem> filesystem;

    if (s3_enabled) {
        ARROW_ASSIGN_OR_RAISE(filesystem, arrow::fs::S3FileSystem::Make(CreateS3Options(), io_context));
    } else {
        filesystem = std::make_shared<arrow::fs::LocalFileSystem>();
    }

    return filesystem;
}

arrow::Result<std::shared_ptr<arrow::RecordBatchReader>> ReadFilteredParquetRecordReaderS3(
                                                            std::shared_ptr<arrow::fs::FileSystem> filesystem,
                                                            const std::string &parquet_path,
                                                            std::vector<std::string> &columns,
                                                            const std::int64_t start, const std::int64_t end)
{

  // Create a Parquet file format
  auto format = std::make_shared<ds::ParquetFileFormat>();
  auto parquet_scan_options = std::make_shared<arrow::dataset::ParquetFragmentScanOptions>();

  parquet_scan_options->reader_properties = parquetProperties();
  parquet_scan_options->arrow_reader_properties = arrowReaderProperties();

  auto scan_options = std::make_shared<arrow::dataset::ScanOptions>();

  arrow::dataset::FileSystemFactoryOptions options;
  ARROW_ASSIGN_OR_RAISE(auto factory, arrow::dataset::FileSystemDatasetFactory::Make(
                                          filesystem, {parquet_path}, format, options));

  ARROW_ASSIGN_OR_RAISE(auto dataset, factory->Finish());

  // Create a scanner builder
  ARROW_ASSIGN_OR_RAISE(auto scan_builder, dataset->NewScan());
  ARROW_RETURN_NOT_OK(scan_builder->FragmentScanOptions(parquet_scan_options));

  // Set a filter: e.g., column start <= "INDEX" <= end
  auto filter = cp::and_(cp::greater_equal(cp::field_ref("INDEX"), cp::literal(start)),
                         cp::less(cp::field_ref("INDEX"), cp::literal(end)));

  ARROW_RETURN_NOT_OK(scan_builder->Filter(filter));

  // Specify the columns to read, e.g., "TRACE_DATA"
  ARROW_RETURN_NOT_OK(scan_builder->Project(columns));

  // Finish the scanner
  ARROW_ASSIGN_OR_RAISE(auto scanner, scan_builder->Finish());

  scanner->options()->use_threads = true;     // Enable multithreading for better performance
  scanner->options()->cache_metadata = true; // Enable async reading
  scanner->options()->batch_readahead = 16;   // Set batch readahead for performance
  return scanner->ToRecordBatchReader();
}
void test(std::shared_ptr<arrow::fs::FileSystem> filesystem, const std::vector<arrow::fs::FileInfo> &file_infos, const std::int64_t start, const std::int64_t end) {

    std::vector<std::string> columns = {"INDEX", "METADATA1", "METADATA2", "METADATA3", "METADATA4"};
    std::vector<std::array<double, 4>> data(file_infos.size()*(end-start));
    // Process files in parallel using TBB
    tbb::parallel_for(tbb::blocked_range<size_t>(0, file_infos.size()),
        [&](const tbb::blocked_range<size_t>& range) {
            for (size_t i = range.begin(); i != range.end(); ++i) {
                const auto& file_info = file_infos[i];
                if (file_info.type() == arrow::fs::FileType::File) {
                    // Process each file here
                    std::cout << "Processing file: " << file_info.path() << std::endl;

                    auto reader_result = ReadFilteredParquetRecordReaderS3(filesystem, file_info.path(), columns, start, end);
                    if (!reader_result.ok()) {
                        std::cerr << "Error reading Parquet files: " << reader_result.status().ToString() << std::endl;
                        return;
                    }
                    auto reader = reader_result.ValueOrDie();

                    while (true) {
                            std::shared_ptr<arrow::RecordBatch> batch;
                            auto rb_read_status = reader->ReadNext(&batch);
                            if (!rb_read_status.ok()) {
                                std::cerr<<"rb_read_status="<<rb_read_status.ToString()<<std::endl;
                                return;
                            }
                            if (!batch) break; // End of stream
                            std::cout << "Processed batch with " << batch->num_rows() << " rows from file: " << file_info.path() << std::endl;
                    }
                }
            }
        });
}
int main(int argc, char* argv[]) {
    absl::SetProgramUsageMessage("Usage: arrow_tbb_dataset_s3 [options]");
    absl::ParseCommandLine(argc, argv);
    auto statusS3 = arrow::fs::InitializeS3(arrow::fs::S3GlobalOptions{});
    if(!statusS3.ok()) {
        std::cerr << "Error initializing S3: " << statusS3.ToString() << std::endl;
        return 1;
    }
    auto s3_path = absl::GetFlag(FLAGS_s3_path);
    auto num_threads = absl::GetFlag(FLAGS_num_threads);

    auto start = absl::GetFlag(FLAGS_start);
    auto end = absl::GetFlag(FLAGS_end);
    if (start < 0 || end < 0 || start >= end) {
        std::cerr << "Invalid start or end trace values." << std::endl;
        return 1;
    }

    // Create filesystem
    auto io_pool_result = arrow::internal::ThreadPool::Make(num_threads);
    if (!io_pool_result.ok()) {
        std::cerr << "Error creating IO thread pool: " << io_pool_result.status().ToString() << std::endl;
        return 1;
    }
    auto io_pool = io_pool_result.ValueOrDie();
    arrow::io::IOContext io_context(io_pool.get(), arrow::io::default_io_context().stop_token());
    auto filesystem = createFilesystem(true, io_context);
    if (!filesystem.ok()) {
        std::cerr << "Error creating filesystem: " << filesystem.status().ToString() << std::endl;
        return 1;
    }

    auto fs = filesystem.ValueOrDie();
    // List directory contents
    arrow::fs::FileSelector selector;
    selector.base_dir = s3_path;
    selector.recursive = true;
    selector.max_recursion = 1;

    arrow::Result<std::vector<arrow::fs::FileInfo>> result = fs->GetFileInfo(selector);
    if (!result.ok()) {
        std::cerr << "Error listing directory: " << result.status().ToString() << std::endl;
        return 1;
    }
    auto file_infos = result.ValueOrDie();
    auto ret = test(fs, file_infos, start, end );
    if (!ret.ok()) {
        std::cerr << "Error in test3: " << ret.status().ToString() << std::endl;
        return 1;
    }
    auto statusS3Finish = arrow::fs::FinalizeS3();
    if (!statusS3Finish.ok()) {
        std::cerr << "Error finalizing S3: " << statusS3Finish.ToString() << std::endl;
        return 1;
    }
    return 0;
}

With this approach I still got very low cpu utilization/throughput

Any pointers or caveats woudl be useful

thanks =)

[felipecrv]

How far are you from saturating your available network bandwidth?

You should try disabling "use threads" in the scanner and parquet reader since you're managing the scheduling now. These options means more thread creation and scheduler overhead when you already have tbb::parallel_for trying to schedule the tasks at a more coarse-grained level. It might not improve the performance, but it's worth trying.

Other things to experiment tweaking is increasing the block size (in number of files) of the parallel-for. Spawning too many tasks might create too much competition for I/O and CPU which might turn out to make things worse.

A profiler can show where the bottlenecks are.