Introduce DistributedExt trait that extends the capabilities of DataFusion's session building tools

[gabotechs]

Introduces the DistributedExt trait, that contributes new methods to upstream DataFusion SessionConfig, SessionStateBuilder, SessionState and SessionContext for working with Distributed DataFusion.

The goal is that just by including the following import:

use datafusion_distributed::DistributedExt;

Some new methods will magically pop up in SessionConfig, SessionStateBuilder, SessionState and SessionContext:

let state = SessionStateBuilder::new()
    .with_default_features()
    .with_runtime_env(runtime_env)
    .with_distributed_option_extension(custom_option_extension)? // <- this magically appears
    .with_distributed_user_codec(custom_user_codec)              // <- this magically appears
    .with_distributed_channel_resolver(custom_channel_resolver)  // <- this magically appears
    .with_scalar_functions()
    .build();

This makes adding Distributed capabilities to an existing DataFusion program a seamless experience, and users will have the perception that they are working with normal DataFusion.

Previously, the experience for providing distributed capabilities was all over the place:

• adding users codecs was done with pure functions that mutated DataFusion objects
• providing a channel manager was done either in the ArrowFlightEndpoint::new() method or by setting it in a SessionState
• providing custom config extensions was done as extension traits of DataFusion session building structures.

With this PR, the experience of providing the Distributed extensions is done all as extension trait methods in DataFusion's SessionConfig, SessionStateBuilder, SessionState and SessionContext, and with all these methods well documented and centralized in one place (distributed_ext.rs)

[gabotechs]

This is pretty much just some trait trickery for allowing users to ergonomically build on top of existing DistributedSessionBuilder implementations adding further capabilities.

We use it in our localhost implementation for adding ourselves the LocalHostChannelResolver on top of whatever DistributedSessionBuilder each individual test provided.

[NGA-TRAN]

Nice refactor. I only have a few minor comments for documentation and a set of questions for future thinking

[NGA-TRAN]

I like this interface and I know you are refactoring it so it is not new but because I am here, I am trying to think out loud whether these are possible if we want to consider them in the future. I am talking about these in the context of providing API for scheduling, workload balancing and fairness for queries.

IIUC, these URLs are set per session, right? This means all queries of the same session will get the same set of workers. However, the number of workers each stage of a query uses depends on the needs of each stage itself and the stage will be able to control that which is very nice.

So a few thinkings:

1. If we have a large pool of workers, a large query may use all of them to paralyze the execution but a medium query may only need a few workers. Let us talk about a medium query that only needs a subset of workers. Should we provide an interface to return the same subset of workers different stages of a query requests? Or we prefer different stage of a query uses different subset of workers? Or let them use random subset of workers? These have pros and cons and also depends on the nature of the queries. Maybe all three options are good for different use cases and we want to provide APIs for all options? 🤔
2. How can we balance the workload for the workers if only a subset of the workers are used by many queries? If the queries keep using the first workers in the vector, will those workers be overused while the rest of the workers idle?

[gabotechs]

However, the number of workers each stage of a query uses depends on the needs of each stage itself and the stage will be able to control that which is very nice.

Not really, only the head stage in a plan has control over which URLs get assigned to all tasks in the plan, so the worker URLs are just decided once per query (at the beginning of it)

Should we provide an interface to return the same subset of workers different stages of a query requests? Or let them use random subset of workers?

User's right now need to implement this trait for providing the URLs:

pub trait ChannelResolver {
    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError>;
}

As they are free to return any Vec<Url> they want, they can choose to either always return the same, randomize it, etc..

Although at some point we probably want to extend the API to let people more control over the worker assignation itself.

How can we balance the workload for the workers if only a subset of the workers are used by many queries? If the queries keep using the first workers in the vector, will those workers be overused while the rest of the workers idle?

If the user implements something that does not shuffle the vector before returning it, then yes, probably only the first subset of workers will be used.

A first initial approach could be to just document this behavior so that users can inform their implementation, but extending the API somehow to allow finer grained control over worker assignation could also be a good approach.

[NGA-TRAN]

I had a hard time to see why we need this build_state function just to include just an Ok(...) but then I found you wrote build_state function and then provided it in start_localhost_context which makes sense. Maybe add that kind of instructions here or an appropriate place to show how to us it?

[gabotechs]

Added some comments about how to use them.

For showing how to use it with ArrowFlightEndpoint, I'm still not sure if it makes sense to do it either here, in the ArrowFlightEndpoint docs that I have not yet done, or as examples in a still not existent examples/ folder.

I'll make a follow up PR with this.

[NGA-TRAN]

All these instructions are similar to above. Since these are all functions of the same trait, do you think it is possible to make one set of instructions for all of them?

[gabotechs]

🤔 not sure what do you mean by "set of instructions". If you are talking about the amount of repetition happening in this documentation, it's a bit annoying, but I don't think Rust offers any way of reusing docs from one method to others, so doc generators and IDEs are expecting the full doc to be inlined in each function.

[robtandy]

Agree that the examples/ folder is now standing out as missing. It would be the entry point for interested users to see how to use the library.

This refactor looks ergonomically improved! PR approved. ty ty!