[FEAT] Compute pool for native executor

[colin-ho]

Create a multithreaded compute runtime for swordfish compute tasks. Switch query runtime to be single threaded, and use IO pool for scan task streams.

Additionally, adds in a tokio_select together with the tokio::signal::ctrlc and main async execution loop so that queries can be cancelled.

import os
import daft
import numpy
import time
import psutil

current_process = psutil.Process(os.getpid())

daft.set_execution_config(enable_native_executor=True, default_morsel_size=1)
dfs = [
    iter(
        daft.from_pydict({"a": numpy.random.rand(10)}).with_column(
            "plus_one", daft.col("a") + 1
        )
    )
    for _ in range(10)
]
while True:
    for i, df in enumerate(dfs):
        time.sleep(0.1)
        try:
            print("threads: ", current_process.num_threads())
            print(next(df))
        except StopIteration:
            dfs.pop(i)
    if not dfs:
        break

If you run this script you can see that the number of threads increases by only 1 per dataframe.

TODO:

• replace rayon with this -> [FEAT] Compute Pool 2 #3076

[codspeed-hq]

CodSpeed Performance Report

Merging #2986 will not alter performance

_{Comparing colin/compute-pool (7aa7b61) with main (4ec76ce)}

Summary

✅ 17 untouched benchmarks

[codecov]

Codecov Report

Attention: Patch coverage is 81.17647% with 64 lines in your changes missing coverage. Please review.

Project coverage is 78.70%. Comparing base (4ec76ce) to head (7aa7b61).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/daft-local-execution/src/sources/scan_task.rs	0.00%	39 Missing ⚠️
src/common/runtime/src/lib.rs	88.59%	13 Missing ⚠️
...-execution/src/intermediate_ops/intermediate_op.rs	87.50%	3 Missing ⚠️
.../src/intermediate_ops/anti_semi_hash_join_probe.rs	86.66%	2 Missing ⚠️
src/daft-local-execution/src/run.rs	80.00%	2 Missing ⚠️
src/daft-csv/src/metadata.rs	50.00%	1 Missing ⚠️
src/daft-io/src/python.rs	50.00%	1 Missing ⚠️
src/daft-json/src/schema.rs	50.00%	1 Missing ⚠️
src/daft-local-execution/src/sinks/concat.rs	90.90%	1 Missing ⚠️
src/daft-parquet/src/read.rs	90.90%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2986      +/-   ##
==========================================
+ Coverage   78.66%   78.70%   +0.03%     
==========================================
  Files         618      619       +1     
  Lines       73219    73412     +193     
==========================================
+ Hits        57601    57776     +175     
- Misses      15618    15636      +18     

Files with missing lines	Coverage Δ
src/daft-csv/src/read.rs	99.33% <100.00%> (ø)
src/daft-functions/src/tokenize/bpe.rs	96.35% <100.00%> (ø)
src/daft-functions/src/uri/download.rs	83.46% <100.00%> (ø)
src/daft-functions/src/uri/upload.rs	64.63% <100.00%> (ø)
src/daft-io/src/lib.rs	71.21% <ø> (-3.33%)	⬇️
src/daft-io/src/s3_like.rs	66.41% <ø> (ø)
src/daft-json/src/read.rs	95.42% <100.00%> (ø)
...-local-execution/src/intermediate_ops/aggregate.rs	100.00% <ø> (ø)
...ft-local-execution/src/intermediate_ops/explode.rs	100.00% <ø> (ø)
...aft-local-execution/src/intermediate_ops/filter.rs	100.00% <ø> (ø)
... and 23 more

... and 2 files with indirect coverage changes

[colin-ho]

Todo: look into thread priority: https://docs.rs/thread-priority/latest/thread_priority/, we could potentially set priorities for compute and io pool.

Would do it on https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.on_thread_start

[samster25]

we can prob put this as a method on StreamingSinkState that does the following:

• as_any
• downcasts
• error message
  like fn get_state_mut<T>(&self) -> &mut T

[colin-ho]

Opted for

pub(crate) fn with_state_mut<T: DynStreamingSinkState + 'static, F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut T) -> R,

since returning &mut T will outlive the lifetime of the mutex guard.