Feedback on high memory usage when merging N parquet files

[ndchandar]

Hello,
I am writing a program that takes N parquet files (where N = 40). Each source parquet file is about ~6 to ~8 MB in size and are Zstd compressed. They are compacted/combined to produce a bigger sized parquet file (~220 to ~250 MB). It appears that we need as much as ~24 GB of memory to have a successful compaction. This gist https://gist.github.com/ndchandar/3900558ff719cefeb8b058e36a18f8be#file-parquet_rewriter-rs-L32-L138 (export_with_datafusion) is the interesting bit. It basically lists all files in a directory, takes N files and compacts them. I tried giving hints to the optimizer that the sources are already sorted but it doesn't seem to help. The row group size is set to 1M.

Giving less memory (E.g 12 or 16 gb), I am running into the below issue

Caused by:
    Resources exhausted: Failed to allocate additional 2.0 MB for ExternalSorterMerge[4] with 49.8 MB already allocated for this reservation - 1826.2 KB remain available for the total pool

I am trying to understand why the spill is not happening efficiently (I am relatively new to DataFusion). Looking for any help/hints to reduce the memory utilization

[ndchandar]

I was able to reduce parallelism by tuning datafusion.execution.target_partitions for our workloads. This resulted in lesser memory and cpu usage. I also bumpled datafusion.execution.parquet.write_batch_size to a much higher number (from the default 8192 to 65536. Are there other parameters that I could tune? I am trying to find the balance between optimal memory/cpu usage versus being reasonably quick with regards to compaction/merging

[ryux1]

The error points at the explicit output sort, not the Parquet writer itself:

.with_sort_by(vec![signal_name_sort_expr, timestamp_sort_expr])

That introduces the ExternalSorterMerge named in the error. Also, file_sort_order is a promise that each input file has that ordering; it does not make the concatenation of 40 files globally ordered when their key ranges overlap. So DataFusion may still need a merge/sort for the requested output ordering.

The first change I would try is replacing GreedyMemoryPool with FairSpillPool:

use datafusion::execution::memory_pool::FairSpillPool;

let memory_pool = Arc::new(FairSpillPool::new(limit_bytes));

In 50.3, the greedy pool is documented as appropriate for no spill or a single spillable operator, while the fair pool is intended when multiple spillable operators compete for the limit: 50.3 source.

A few other changes should reduce the peak:

• Keep target_partitions low, as you found.
• Move parquet.write_batch_size back down. Raising it from 8192 to 65536 increases writer buffering; it is a throughput tradeoff in the wrong direction for a memory-constrained run.
• Reduce max_row_group_size from 1M while diagnosing. DataFusion's own 50.3 config docs note that larger row groups require more memory to write: source.
• Because you partition output by year/month/day, bound writer concurrency too. In 50.3 the relevant settings are datafusion.execution.minimum_parallel_output_files (default 4) and datafusion.execution.max_buffered_batches_per_output_file (default 2): source.

Finally, 6-8 MB Zstd-compressed input size is not a useful estimate of Arrow working memory: sorting operates on decoded arrays, plus sort buffers and active Parquet writers.

I would test one change at a time in this order: fair pool, smaller write batch, smaller row groups, then lower output-writer concurrency. If removing with_sort_by makes memory collapse, that confirms the global output sort is the dominant consumer; only remove it if downstream readers do not require that physical ordering.