[FEAT] Enable init args for stateful UDFs

[kevinzwang]

No description provided.

[codspeed-hq]

CodSpeed Performance Report

Merging #2956 will not alter performance

_{Comparing kevin/stateful-actor-init-args (05cb70f) with main (c5a6d88)}

Summary

✅ 17 untouched benchmarks

[jaychia]

LGTM. We should also ensure the other configurations are passed along (I believe the only other one we need to think about is .with_batch_size())

[jaychia]

We should remember to reset the config too to the old one after the yield point.

We do this for the execution context (see: daft.context.py)

@contextlib.contextmanager
def execution_config_ctx(**kwargs):
    """Context manager that wraps set_execution_config to reset the config to its original setting afternwards"""
    original_config = get_context().daft_execution_config
    try:
        set_execution_config(**kwargs)
        yield
    finally:
        set_execution_config(config=original_config)

We can probably just make one for the planning config too!

[kevinzwang]

@jaychia the batch size is passed in already in the rust side that actually calls the stateful UDF: https://github.com/Eventual-Inc/Daft/blob/kevin/stateful-actor-init-args/src/daft-dsl/src/functions/python/udf.rs#L197

[codecov]

Codecov Report

Attention: Patch coverage is 76.31579% with 9 lines in your changes missing coverage. Please review.

Project coverage is 78.49%. Comparing base (46c5a7e) to head (05cb70f).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
daft/runners/pyrunner.py	0.00%	6 Missing ⚠️
daft/runners/ray_runner.py	66.66%	2 Missing ⚠️
src/daft-dsl/src/functions/python/mod.rs	91.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2956      +/-   ##
==========================================
+ Coverage   76.34%   78.49%   +2.14%     
==========================================
  Files         597      597              
  Lines       71388    70091    -1297     
==========================================
+ Hits        54504    55017     +513     
+ Misses      16884    15074    -1810     

Flag	Coverage Δ
	78.49% <76.31%> (+2.14%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
daft/udf.py	95.52% <100.00%> (+1.02%)	⬆️
src/common/daft-config/src/python.rs	77.34% <100.00%> (+5.50%)	⬆️
src/daft-dsl/src/python.rs	91.10% <100.00%> (+0.91%)	⬆️
src/daft-dsl/src/functions/python/mod.rs	79.19% <91.66%> (+24.19%)	⬆️
daft/runners/ray_runner.py	80.62% <66.66%> (+4.91%)	⬆️
daft/runners/pyrunner.py	85.25% <0.00%> (+3.82%)	⬆️

... and 38 files with indirect coverage changes