[FEAT] Cap parallelism on local parquet reader

[colin-ho]

Implement a dynamically parallel local streaming parquet reader.

Background

The current streaming local parquet reader, while fast and streaming, has some problems:

• It reads and deserializes ALL row groups and ALL columns in parallel.
• It does not respect downstream back-pressure (the crossbeam channels are all bounded by max chunks, it's free to fill it up).

This leads to unnecessarily high memory usage, and it potentially starves downstream tasks.

Solution

Instead of launching all tasks at once, we can cap the number of parallel tasks based on certain factors:

• Number of CPUs
• Number of Columns.

Results

Most glaringly, the benefits of these are in memory usage of streaming queries, for example:

next(daft.read_parquet("data/tpch-dbgen/1_0/1/parquet/lineitem").iter_partitions()) # read lineitem tpch sf1

The new implementation hits a peak of 300mb, while the old goes over 1gb.

Another example, where we stream the entire file, but the consumption is slow:

for _ in daft.read_parquet("/Users/colinho/Desktop/Daft/z/daft_tpch_100g_32part_64RG.parquet").iter_partitions():
    time.sleep(0.1)

The new implementation hits a peak of 1.2gb, while the old goes over 3gb.

To maintain perfomance parity, I also wrote some benchmarks for parquet files with differing rows / cols / row groups, the results show that the new implementation is pretty much on par, with some slight differences.

On reading a tpch sf-1 lineitem table though: the results are pretty much the same: (~0.2s)

[codspeed-hq]

CodSpeed Performance Report

Merging #3310 will improve performances by ×2.2

_{Comparing colin/dynamic-parquet (6fed790) with main (6d30e30)}

Summary

⚡ 2 improvements
✅ 15 untouched benchmarks

Benchmarks breakdown

	Benchmark	main	colin/dynamic-parquet	Change
⚡	test_iter_rows_first_row[100 Small Files]	340.5 ms	212.2 ms	+60.44%
⚡	test_show[100 Small Files]	34.9 ms	15.6 ms	×2.2

[codecov]

Codecov Report

Attention: Patch coverage is 88.18182% with 26 lines in your changes missing coverage. Please review.

Project coverage is 77.34%. Comparing base (6d30e30) to head (6fed790).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/daft-parquet/src/stream_reader.rs	88.58%	25 Missing ⚠️
src/daft-parquet/src/read.rs	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3310      +/-   ##
==========================================
+ Coverage   77.00%   77.34%   +0.33%     
==========================================
  Files         696      696              
  Lines       86039    84849    -1190     
==========================================
- Hits        66256    65628     -628     
+ Misses      19783    19221     -562     

Files with missing lines	Coverage Δ
src/daft-parquet/src/read.rs	75.22% <0.00%> (-0.08%)	⬇️
src/daft-parquet/src/stream_reader.rs	89.85% <88.58%> (+1.57%)	⬆️

... and 10 files with indirect coverage changes

[colin-ho]

Some ideas to consider:

• Maybe we don't need to consider IO time and just consider the wait time.
• Can this semaphore become generic and used for other I/O code, as well as the local executor? Would be cool if we could dynamically adjust degree of operator parallelism as well.
• Can we decrease the permit count in addition to increase?
• Can we add memory pressure to the semaphore?

[samster25]

forgot submitting my review from yesterday lol

[colin-ho]

Tested this new implementation on TPCH SF1: