Parallelize Serialization of Columns within Parquet RowGroups

[devinjdangelo]

This PR is draft as it depends on apache/arrow-rs#4871

Which issue does this PR close?

Closes #7591
Closes #7589
Closes #7590
Closes #7632
Related to apache/arrow-rs#1718

Rationale for this change

The original parallel parquet writer independently serialized rowgroups. This has two major disadvantages:

1. High memory usage as each parallel RowGroup is entirely buffered in memory before being flushed to ObjectStore.
2. No easy way to respect the WriterProperties.max_row_group_size setting as the number of RowGroups is determined by the parallelism.

This PR implements a new approach which parallelizes parquet serialization primarily by serializing the different columns of a single row group in parallel. Different row groups can still be serialized in parallel if RecordBatches are being produced fast enough and the configured buffer size is large enough.

What changes are included in this PR?

1. New process to break apart an ArrowRowGroupWriter into ArrowColumnWriter components and send each to dedicated parallel task.
2. Respects the WriterProperties.max_row_group_size setting, closing a row group as soon as this threshold is reached.
3. If RecordBatches are being produced fast enough, the next row group can begin serializing before the prior has finished.
4. User Configurable limits on the amount of RowGroups and RecordBatches that are allowed to accumulate in memory before backpressure kicks in. These are made user configurable since optimal values will depend on the systems mix of cpu, memory, and I/O resources.
5. Refactor parallel parquet code to improve readabiltiy

Are these changes tested?

Yes, by existing tests.

Benchmarking Script

use datafusion::{
    dataframe::DataFrameWriteOptions,
    parquet::{basic::ZstdLevel, file::properties::WriterProperties},
    prelude::*,
};
use datafusion_common::DataFusionError;
use object_store::local::LocalFileSystem;
use peak_alloc::PeakAlloc;
use std::{sync::Arc, time::Instant};
use url::Url;

#[global_allocator]
static PEAK_ALLOC: PeakAlloc = PeakAlloc;

const FILENAME: &str =
    "/home/dev/arrow-datafusion/benchmarks/data/tpch_sf10/lineitem/part-0.parquet";

#[tokio::main]
async fn main() -> Result<(), DataFusionError> {
    let mut runtimes = Vec::new();
    let mut max_memory = Vec::new();
    for parallelism in ["on", "off"] {
        for single_file_output in [true, false] {
            PEAK_ALLOC.reset_peak_usage();
            let _ctx = if parallelism == "on" {
                let config = SessionConfig::new().set(
                    "datafusion.execution.parquet.allow_single_file_parallelism",
                    true.into(),
                );
                SessionContext::with_config(config)
            } else {
                let config = SessionConfig::new().set(
                    "datafusion.execution.parquet.allow_single_file_parallelism",
                    false.into(),
                );
                SessionContext::with_config(config)
            };
            let local = Arc::new(LocalFileSystem::new());
            let local_url = Url::parse("file://local").unwrap();
            _ctx.runtime_env().register_object_store(&local_url, local);

            let _read_options = ParquetReadOptions::default();

            let _df = _ctx
                .read_parquet(FILENAME, _read_options.clone())
                .await?;
            let _df = _df.clone()
                .union(_df.clone())?
                .union(_df.clone())?
                .union(_df.clone())?
                ;

            let out_path = if single_file_output{
                format!(
                "file://local/home/dev/arrow-datafusion/test_out/bench{parallelism}.parquet"
            )} else{
                format!(
                    "file://local/home/dev/arrow-datafusion/test_out/{parallelism}/"
                )
            };

            let start = Instant::now();
            _df.clone()
                .write_parquet(
                    out_path.as_str(),
                    DataFrameWriteOptions::new().with_single_file_output(single_file_output),
                    None,
                )
                .await?;
            let elapsed = Instant::now() - start;
            println!("write as parquet with parallelism {parallelism} to disk took -> {elapsed:?}");
            runtimes.push(elapsed);

            let peak_mem = PEAK_ALLOC.peak_usage_as_mb();
            println!(
                "Peak memory usage with parallelism {parallelism} is: {}MB",
                peak_mem
            );
            max_memory.push(peak_mem);

            let test_read = _ctx.read_parquet(out_path, _read_options).await.unwrap();

            // Ensure a query over all rows produces a correct result
            test_read
                .aggregate(vec![], vec![avg(col("l_partkey")), count(col("l_partkey"))])?
                .show()
                .await?;
        }
    }

    println!("Runtimes: {:?}", runtimes);
    println!("Peak memory: {:?}", max_memory);

    Ok(())
}

Benchmarking Results

Single File Output = True

	Single File Parallelism Off	Single File Parallelism On	Diff (%)
Execution Time (s)	154.98	22.65	-85.39%
Peak Memory Usage (MB)	449.7	684.4	52.19%

Single File Output = False

	Single File Parallelism Off	Single File Parallelism On
Execution Time (s)	37.76	25.22	-33.21%
Peak Memory Usage (MB)	1210.0	1017.6	-15.90%

Are there any user-facing changes?

Faster parquet serialization

[devinjdangelo]

@alamb @tustvold I updated this branch to build against arrow-rs master branch while we wait for the next release. Currently getting a compile error, which I opened a PR to fix here apache/arrow-rs#4893.

Once that PR is merged, the checks on this branch should pass and if you want to try it locally it should work (so long as there are no breaking changes pushed to arrow-rs master).

[alamb]

Thank you @devinjdangelo

[alamb]

I plan to review this tomorrow

[devinjdangelo]

With the latest release of arrow-rs, this PR is now ready for review. I have merged in all of the conflicting write related changes that have happened since this has been waiting, and updated the benchmarking/results as well (see PR description).

cc @alamb

[tustvold]

Makes sense to me, just some minor comments

[tustvold]

This seems very high

[devinjdangelo]

I lowered this one back to 2 with no noticeable impact on performance.

[tustvold]

This seems extraordinarily high, in order places we buffer up to 1

[devinjdangelo]

This one does have a significant impact on performance if lowered significantly. I spent some time testing and tuning the exact values. Setting max parallel row groups to 2 and maximum_buffered_record_batches_per_stream to 128 allows two row groups to run in parallel. If this is set too low, backpressure will kick in too long before a second row group can be spawned and everything will wait on just 1 rowgroup to write.

                 ┌─────────────────┐        ┌────────────────────────┐
 ExecPlan───────►│ RowGroup1 Queue ├───────►│Parallel Col Serializers│
   │             └─────────────────┘        └────────────────────────┘
   │
   │
   │
   │             ┌─────────────────┐        ┌────────────────────────┐
   └────────────►│RowGroup2 Queue  ├───────►│Parallel Col Serializers│
                 └─────────────────┘        └────────────────────────┘
Once max_rowgroup_rows
Sent to RowGroup1 Queue
Spawn a new Queue with
its own parallel writers

RowGroup2 Queue won't be created until RowGroup1 Queue has received the desired number of rows. The goal is to have two row groups serializing in parallel if RecordBatches are being produced fast enough. For a streaming plan reading from disk, we probably never need more than 2 in parallel. If we are writing already in-memory data on a system with many cores, it is highly beneficial to boost these queue sizes even more so we could have an arbitrarily large number of row groups serializing in parallel.

[tustvold]

Could we instead rely on input partitioning? Relying on being able to buffer an entire row group in memory as arrow data has caused major issues in the past. I don't feel comfortable merging this as is, as it will lead to a major regression in memory usage.

Edit: do we even really need row group parallelism, we already have parallelism at the file and column level, it seems a tad unnecessary tbh, and comes with major drawbacks

[devinjdangelo]

Lower the buffer size down to 2 would save a few 100MB of memory, but at the cost of more than double the serialization time. The decision comes down to which is a better default behavior. Imo, most users can spare the memory for the performance benefit, and for those that can't they can always lower the buffer size. We could instead default to a low buffer size (favoring minimum memory usage over execution time) and I could update the docs to suggest increasing the buffer for signficant performance gains on systems with many cores. Here are the numbers I gathered using the script in the description:

Parallel Parquet Writer, Varying Row Group Buffer Sizes

	Buffer Size=2	Buffer Size=64	Buffer Size=128
Execution Time (s)	62.6	35.6	25.02
Peak Memory Usage (MB)	495.0	606.8	712.05

For comparison, the non parallelized writer takes 155s and peak memory usage is 449.7MB for the same task.

[tustvold]

Is the memory usage of these buffered RecordBatch tracked with a MemoryReservation at all?

For context apache/arrow-rs#3871 was the issue from back when ArrowWriter did something similar, and this caused operational problems for us as writers could easily consume GB of memory when writing large files.

[devinjdangelo]

I haven't added anything special to track with a MemoryReservation. It is just a bounded channel size. Memory usage can't grow without bound, but it can grow up to however large 128 RecordBatches is in memory. So with a particularly wide table or an extra large batch_size setting, I could see it climbing into Gb of memory. If we are concerned about that, we could use buffer_size=2 as the default and leave it up to user's if it is worth the memory/performance tradeoff to increase the buffer size.

It is also true that the numbers above are gathered from a fairly extreme case of writing ~250 million rows to a single parquet file, and you could instead just write 2 or more files in parallel to close the performance gap. For a more reasonably sized ~50million rows and 1.5Gb the gap is smaller, but it is still there:

Parallel Parquet Writer, Varying Row Group Buffer Sizes

	Buffer Size=2	Buffer Size=64	Buffer Size=128
Execution Time (s)	19.76	13.41	13.68
Peak Memory Usage (MB)	272.2	277.2	281.8

[devinjdangelo]

Ok, I pushed a commit to change the buffer size back down to 2. I think this is a good default in most cases. I called out in the docs that if you have memory to spare and in particular if you are writing out cached in memory data (like a cached data frame) then you will likely benefit significantly from boosting the buffer sizes.

This optimization could perhaps be automatic at some point if we could automatically set the buffer size based on a known memory budget.

[tustvold]

Thanks, I think once you start buffering non-trivial numbers of batches, and therefore non-trivial amounts of data in memory, it is important that this is accounted for in the MemoryPool. Ideally we would account everywhere, but for practical reasons we just really need to get the places where non-trivial amounts can build up, so as to prevent OOMs.

[tustvold]

Am I correct in thinking this is still bounded by the input parallelism? Is it worth noting this?

[devinjdangelo]

It can be limited based on the queue sizes as described in my other comment.

[tustvold]

Should we be using a JoinSet here so that if a task fails it aborts them all?

[devinjdangelo]

I think that the order they are joined/appended matters for the SerializedRowGroupWriter, but I would have to double check.

[tustvold]

Should this check that current_rg_rows != 0

[devinjdangelo]

Yes good call, I updated this.

buffering

[tustvold]

Thank you

[alamb]

Thank you @devinjdangelo and @tustvold