Add context to IR.do_evaluate

[TomAugspurger]

Description

This adds a keyword-only context argument to cudf_polars IR.do_evaluate method. The purpose to provide access to special pieces of data that might be necessary for controlling an IR nodes' execution, but doesn't belong on the IR node itself as a non-child argument. Specifically, we'd like to provide a CUDA stream argument as part of #20228, but we generalize that slightly and provide a system for providing arbitrary data.

A few notes on the implementation:

• For now, the context is just an empty dataclass. I suspect its design might change in the future.
• I've opted to push the creation of the context as high as possible. For now it's created in _callback and passed into ir.evaluate / evaluate_streaming and from there to all the methods that require it.
• There's some awkwardness between how our IR nodes and Dask's task graph treat arguments. I've opted to make context keyword only in IR.do_evaluate(..., context). However, Dask's task graph doesn't really deal with that. It wants a tuple of (function, arg1, arg2, ...). So that requires using functools.partial(function, context=context)(arg1, arg2, ...).
• After implementing this, I realized that Expr.evaluate also takes a context, and its a different type ExecutionContext :( I can rename the IR variant if we want.

Just a draft for now, and probably not worth reviewing until I have a branch somewhere that combines CUDA streams with this to verify it meets our needs.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[TomAugspurger]

5a77a80 has a POC for how this can be used. We add a new_stream: Callable[[], Stream] member to the context dataclass. Inside of do_evaluate we can call context.new_stream(). Once rapidsai/rapidsmpf#592 is done, we should be able to pass in a context that uses the stream pool from rapidsmpf.

Alternatively, rather than giving a Callable[[], Stream] we could attach a stream directly to the context and relying on dataclasses.replace() with new streams as needed. I'm not sure which is better at the moment.

Finally, we could drop the dataclass and just make it a dictionary. But I'd prefer to keep things structured where possible, so that both the functions and the callers of the function know what belongs in the context. We can attach an extra field to the dataclass that's just a dictionary if we need to pass arbitrary things in.

[TomAugspurger]

/merge