Add SessionStateBuilder and extract out the registration of defaults #11403

Omega359 · 2024-07-10T23:24:59Z

SessionStateBuilder has been added and the existing SessionState::new* functions updated to use the builder to construct a SessionState. Defaults previously set in the new* functions have been extracted out into SessionStateDefaults

Which issue does this PR close?

Closes #11320

Rationale for this change

Adds the ability to fine grain control how a SessionState is constructed.

What changes are included in this PR?

Code.

Are these changes tested?

All current tests pass

Are there any user-facing changes?

All the builder style functions on SessionState have been deprecated.

api change label required.

…asic SessionState.

…State and clone it.

datafusion/core/src/execution/session_state.rs

jayzhan211 · 2024-07-12T00:17:26Z

datafusion/core/src/execution/session_state.rs

-
-        new_self
+        SessionStateBuilder::new_with_config_rt(config, runtime)
+            .with_defaults(true)


we can make new_with_config_rt with default: false. And call with_defaults() to set defaults

I think it would be good to mention / point at SessionStateBuilder in the docs here

As a follow on PR it probably makes sense to deprecate SessionState::new_with_config and new_with_config... and redirect people to use SessionStateBuilder

I waffled back and forth as whether to have default set to false or true by, um, default and decided that it would likely be best to mirror the current system behavior. That being said, it's simple to change.

I think we need either session state with default or not. If we have builder that has default: true , we need another function disable_default() or something similar to get the plain session state. If we have builder that has default: false, we need function like with_default to get the session state with builtin. Both are fine to me, as long as builtin is optional

I agree 💯

alamb

Thank you so much @Omega359 -- I think this PR is looking really nice 😍

I had some API improvement suggestions, but I actually think they could be done as follow on PRs as this PR maintains backwards compatibility with the existing APIs.

Specifically I am thinking we could iterate on the internal implementation of SessionStateBuilder easily.

It might make sense to take a quick look at the suggestions for public API changes (e.g. new() vs new_with_config_and_rt() and how use_defaults works) but otherwise 👍

Let me know what you think

alamb · 2024-07-12T13:07:42Z

datafusion/core/src/execution/session_state.rs

-
-        new_self
+        SessionStateBuilder::new_with_config_rt(config, runtime)
+            .with_defaults(true)


I think it would be good to mention / point at SessionStateBuilder in the docs here

As a follow on PR it probably makes sense to deprecate SessionState::new_with_config and new_with_config... and redirect people to use SessionStateBuilder

alamb · 2024-07-12T13:07:49Z

datafusion/core/src/execution/session_state.rs

+        SessionStateBuilder::new_with_config_rt(config, runtime)
+            .with_defaults(true)
+            .with_catalog_list(catalog_list)
+            .build()


this is so nice ❤️

alamb · 2024-07-12T13:13:32Z

datafusion/core/src/execution/session_state.rs

+impl SessionStateBuilder {
+    /// Returns new [`SessionStateBuilder`] using the provided
+    /// [`SessionConfig`] and [`RuntimeEnv`].
+    pub fn new_with_config_rt(


Another way we could do this would be simply

let builder = SessionStateBuilder::new() .with_config(config) .with_runtime_env(runtime_env);

I think the original concern with an API like this is that this would likely create a new SessionConfig::default() just to replace it with config

However, I think we could optimize this eventually (e.g. by deferring the creation of SessionState until calling build() for example) and if we are going to change the API anyways I think we should consider doing the "better" thing

I had that originally however I decided to switch just to mirror the existing behavior. It's a simple change, just would need to document that .build() will fail if config and runtime_env are not provided (assuming that we don't just use defaults there)

I was thinking that if the config and runtime weren't provided then it would make the defaults

Something like this perhaps

fn build(self) -> SessionConfig { let Self { config, ... } = self; // config is None if not provided explicitly let config = config.unwrap_or_else(|| SessionConfig::default); ... }

alamb · 2024-07-12T13:15:14Z

datafusion/core/src/execution/session_state.rs

+        Self {
+            state: SessionState {
+                session_id,
+                analyzer: Analyzer::new(),
+                expr_planners: vec![],
+                optimizer: Optimizer::new(),
+                physical_optimizers: PhysicalOptimizer::new(),
+                query_planner: Arc::new(DefaultQueryPlanner {}),
+                catalog_list,
+                table_functions: HashMap::new(),
+                scalar_functions: HashMap::new(),
+                aggregate_functions: HashMap::new(),
+                window_functions: HashMap::new(),
+                serializer_registry: Arc::new(EmptySerializerRegistry),
+                file_formats: HashMap::new(),
+                table_options: TableOptions::default_from_session_config(
+                    config.options(),
+                ),
+                config,
+                execution_props: ExecutionProps::new(),
+                table_factories: HashMap::new(),
+                runtime_env,
+                function_factory: None,


Rather than using an inner SessionState, it might be possible to simply replicate the fields on the builder (possibly in Option) -- while this would involve replicated code, it would make doing things like deferred creation of session config and runtime env easier

However, since this is an implementation detail we could also change it as a follow on PR without much issue

I thought about that however then the builder would need accessors functions so that the end user could actually query into any defaults set and update as required. Not horrendous but not typical of a normal Builder, at least in languages I've used this pattern with in the past.

Sounds reasonable -- we could try it as a follow on PR too

alamb · 2024-07-12T13:16:20Z

datafusion/core/src/execution/session_state.rs

+    /// Returns a new [SessionStateBuilder] based on an existing [SessionState]
+    /// The session id for the new builder will be set to a unique value; all
+    /// other fields will be cloned from what is set in the provided session state
+    pub fn new_from_existing(existing: &SessionState) -> Self {


Since this clones SessionState anyways, I think it would make more sense for this to take the argument by value

Suggested change

pub fn new_from_existing(existing: &SessionState) -> Self {

pub fn new_from_existing(existing: SessionState) -> Self {

The rationale being there if the caller already has an owned SessionState they can use that and avoid a clone. This API basically requires them to do a clone

In addition to this function, another nice to have UX feature would be a From impl, so users could do

let state: SessionState = get_existing_state(); let state = SessionStateBuilder::from(state) .with_optimizer_pass(Arc::new(MyNewOptimizer{})) .build()

alamb · 2024-07-12T13:18:04Z

datafusion/core/src/execution/session_state.rs

+        }
+    }
+
+    /// Set to true (default = true) if defaults for table_factories, expr_planners, file formats


😍 this is perfect

I feel like defaults is a somewhat overloaded term. Perhaps we could improve the naming here. Some thoughts 🤔

with_built_ins()

with_install_defaults()

with_default_features()

I like .with_default_features() myself.

alamb · 2024-07-12T13:20:37Z

datafusion/core/src/execution/session_state.rs

+                self.state.table_factories =
+                    SessionStateDefaults::default_table_factories();
+            }
+            if self.state.expr_planners.is_empty() {


the semantics of "use_defaults" is effectively ignored if the user has provided their own pieces might be confusing

I think it would be easier to reason about if the defaults got installed as part of SessionBuilder::with_use_defaults() (rather than a flag that deferred the installation to the end) -- that way the user could decide if they wanted to modify the default lists or start from an empty list

What I am hoping long term is that we could pass in the defaults as an argument (and implement it as a trait) thus breaking the dependency between SessionState and core

let state = SessionStateBuilder::new() .with_defaults(SessionStateDefaults::new())? .build()

datafusion/core/src/execution/session_state.rs

More changes planned, will review then

Omega359 · 2024-07-12T23:15:57Z

I've updated the PR with the latest changes that I hope reflect all the feedback received thus far.

jayzhan211

👍

alamb

I think this is awesome -- thank you so much @Omega359 and @jayzhan211

I had a few comments, but I also think we could address them as follow on PRs as well

alamb · 2024-07-14T16:10:14Z

datafusion-cli/src/catalog.rs

@@ -178,13 +180,18 @@ impl SchemaProvider for DynamicFileSchemaProvider {
                // to any command options so the only choice is to use an empty collection
                match scheme {
                    "s3" | "oss" | "cos" => {
-                        state = state.add_table_options_extension(AwsOptions::default());
+                        if let Some(table_options) = builder.table_options() {


It seems like this code previously always registered the table options and now it only registers them if another table option was registered. Is that intentional?

Maybe it doesn't matter

Well .... I updated all of the existing code in the project to reflect the change but I agree that it may be a bit of a surprise to anyone expecting the old behavior when upgrading. 6 of one, 1/2 dozen of another ...

Wait, sorry, misread the comment. I think it works only because the builder was created from an existing session state.

Maybe we can clean it up as a follow on -- it is not critical in my opinion

alamb · 2024-07-14T16:14:00Z

datafusion/core/src/execution/context/mod.rs

@@ -315,7 +320,7 @@ impl SessionContext {
    }

    /// Creates a new `SessionContext` using the provided [`SessionState`]
-    #[deprecated(since = "32.0.0", note = "Use SessionState::new_with_state")]
+    #[deprecated(since = "32.0.0", note = "Use SessionContext::new_with_state")]


As a follow on PR perhaps we can remove this API entirely

We could probably change the note to say say "SessionStateBuilder"

alamb · 2024-07-14T16:14:28Z

datafusion/core/src/execution/context/mod.rs

-        let session_state =
-            SessionState::new_with_config_rt(SessionConfig::new(), runtime)
-                .with_query_planner(Arc::new(MyQueryPlanner {}));
+        let session_state = SessionStateBuilder::new()


this just looks so much nicer

alamb · 2024-07-14T16:15:36Z

datafusion/core/tests/user_defined/user_defined_plan.rs

        .with_query_planner(Arc::new(TopKQueryPlanner {}))
-        .add_optimizer_rule(Arc::new(TopKOptimizerRule {}));
-    state.add_analyzer_rule(Arc::new(MyAnalyzerRule {}));
+        .add_optimizer_rule(Arc::new(TopKOptimizerRule {}))


maybe we could make these names consistent (rather than add_optimizer_rule it could be with_optimizer_rule -- a follow on PR would be totally fine too

will do, easy change.

alamb · 2024-07-14T16:16:03Z

datafusion/core/src/execution/session_state.rs

@@ -89,9 +91,29 @@ use uuid::Uuid;
 /// Execution context for registering data sources and executing queries.
 /// See [`SessionContext`] for a higher level API.
 ///
+/// Use the [`SessionStateBuilder`] to build a SessionState object.


alamb · 2024-07-14T16:20:52Z

datafusion/core/src/execution/session_state.rs

+    /// The session id for the new builder will be unset; all other fields will
+    /// be cloned from what is set in the provided session state
+    pub fn new_from_existing(existing: SessionState) -> Self {
+        let cloned = existing.clone();


Since existing is owned, why do we need cloned here?

I tried removing this clone locally and it seemed to work

leftover from a previous iteration, I'll remove. thx for the catch

alamb · 2024-07-14T16:22:48Z

datafusion/core/src/execution/session_state.rs

+        }
+    }
+
+    /// Set defaults for table_factories, file formats, expr_planners and builtin


I think this is great for this PR

As part of the larger catalog refactor, I think we'll ned to figure out how to get SessioNStateDefaults` out of this file -- but I have some ideas (specifically we can make the defaults a trait that could be passed to the builder).

Not for this PR

alamb · 2024-07-14T16:23:33Z

datafusion/core/src/execution/session_state.rs

+
+    /// Add `physical_optimizer_rule` to the end of the list of
+    /// [`PhysicalOptimizerRule`]s used to rewrite queries.
+    pub fn add_physical_optimizer_rule(


I know this is just following the existing (inconsistent) API naming, but since we are making a new set of APIs, perhaps we could make these builder style APIs consistently start with with_

alamb · 2024-07-14T16:23:47Z

datafusion/core/src/execution/session_state.rs

+
+    /// Add `optimizer_rule` to the end of the list of
+    /// [`OptimizerRule`]s used to rewrite queries.
+    pub fn add_optimizer_rule(


likewise it would be great to call this with_optimizer_rule

alamb · 2024-07-14T16:25:59Z

datafusion/core/src/execution/session_state.rs

+        let config = self.config.unwrap_or_default();
+        let runtime_env = self.runtime_env.unwrap_or(Arc::new(RuntimeEnv::default()));
+
+        let mut state = SessionState {


Stylistically, another way to write this code that might be more boiler plate but would let the compiler ensure all fields are used would be:

let Self { session_id, analyzer, ... // by explicitly listing all these fields, the compiler will complain if any fields are missed } = self; let session_id = session_id.unwrap_or_else(|| Uuid::new_v4().to_string()); let analyzer = analyzer.unwrap_or_default(),

Interesting. That's not an approach I've seen used but I can see how it could be useful. Updating.

I switched a few places in the sql planner to use it after we got a few bugs due to by newly fields being (silently) ignored.

…e/11320_session_state_builder

Omega359 · 2024-07-14T17:58:41Z

Updated build() method to use suggested let Self{..} = self; approach. Add_xxx functions renamed to with_xxx.

alamb

love it 🚀

alamb · 2024-07-14T19:00:28Z

I think technically this is not an API change (though it does deprecate some existing apis) so not marking it

alamb · 2024-07-14T19:00:37Z

Thanks again everyone! this is great

alamb · 2024-07-16T11:20:14Z

Thanks @Omega359 and @jayzhan211

…pache#11403) * Create a SessionStateBuilder and use it for creating anything but a basic SessionState. * Updated new_from_existing to take a reference to the existing SessionState and clone it. * Minor documentation update. * SessionStateDefaults improvements. * Reworked how SessionStateBuilder works from PR feedback. * Bug fix for missing array_expressions cfg feature. * Review feedback updates + doc fixes for SessionStateDefaults * Cargo fmt update.

alamb · 2024-07-31T20:25:17Z

BTW @wiedld just moved InfluxDB 3.0 over to use this new API and it is 👌 -- thanks again @Omega359

SessionState::new_with_config_rt was deprecated in favor of SessionStateBuilder in apache#11403 which is not shipped in 40.x.

…#11839) SessionState::new_with_config_rt was deprecated in favor of SessionStateBuilder in #11403 which is not shipped in 40.x.

The old constructor is deprecated. Ref: apache/datafusion#11403

* update datafusion deps to point to githuc.com/apache/datafusion Datafusion 41 is not yet released on crates.io. * update TableProvider::scan Ref: apache/datafusion#11516 * use SessionStateBuilder The old constructor is deprecated. Ref: apache/datafusion#11403 * update AggregateFunction Upstream Changes: - The field name was switched from `func_name` to func. - AggregateFunctionDefinition was removed Ref: apache/datafusion#11803 * update imports in catalog Catlog API was extracted to a separate crate. Ref: apache/datafusion#11516 * use appropriate path for approx_distinct Ref: apache/datafusion#11644 * migrate AggregateExt to ExprFunctionExt Also removed `sqlparser` dependency since it's re-exported upstream. Ref: apache/datafusion#11550 * update regr_count tests for new return type Ref: apache/datafusion#11731 * migrate from function-array to functions-nested The package was renamed upstream. Ref: apache/datafusion#11602 * cargo fmt * lock datafusion deps to 41 * remove todo from cargo.toml All the datafusion dependencies are re-exported, but I still need to figure out *why*.

Omega359 added 5 commits July 10, 2024 16:59

Create a SessionStateBuilder and use it for creating anything but a b…

fab06da

…asic SessionState.

Updated new_from_existing to take a reference to the existing Session…

de39082

…State and clone it.

Minor documentation update.

7d1bc09

SessionStateDefaults improvements.

490214d

Merge branch 'apache:main' into feature/11320_session_state_builder

93a3a7d

github-actions bot added core Core DataFusion crate substrait labels Jul 10, 2024

Omega359 marked this pull request as ready for review July 11, 2024 13:11

jayzhan211 reviewed Jul 12, 2024

View reviewed changes

datafusion/core/src/execution/session_state.rs Outdated Show resolved Hide resolved

jayzhan211 reviewed Jul 12, 2024

View reviewed changes

alamb mentioned this pull request Jul 12, 2024

DataFusion weekly project plan (Andrew Lamb) - July 8, 2024 #11334

Closed

9 tasks

alamb previously approved these changes Jul 12, 2024

View reviewed changes

Omega359 marked this pull request as draft July 12, 2024 16:51

alamb mentioned this pull request Jul 12, 2024

[Epic] Extract catalog functionality from the core to make it more modular #10782

Open

Omega359 added 2 commits July 12, 2024 18:30

Reworked how SessionStateBuilder works from PR feedback.

e3fad6d

Bug fix for missing array_expressions cfg feature.

8a579cf

Omega359 marked this pull request as ready for review July 12, 2024 23:15

jayzhan211 approved these changes Jul 13, 2024

View reviewed changes

alamb approved these changes Jul 14, 2024

View reviewed changes

Omega359 added 2 commits July 14, 2024 13:19

Merge remote-tracking branch 'refs/remotes/upstream/main' into featur…

eb64e2f

…e/11320_session_state_builder

Review feedback updates + doc fixes for SessionStateDefaults

71decf6

Cargo fmt update.

380ce98

alamb approved these changes Jul 14, 2024

View reviewed changes

alamb merged commit 8475806 into apache:main Jul 14, 2024
24 checks passed

Omega359 deleted the feature/11320_session_state_builder branch July 14, 2024 23:33

jayzhan211 mentioned this pull request Jul 15, 2024

ExprPlanner not propagated to SqlToRel #11477

Closed

alamb mentioned this pull request Jul 20, 2024

Minor: move SessionStateDefaults into its own module #11566

Merged

kezhuw added a commit to kezhuw/datafusion that referenced this pull request Aug 6, 2024

Bump deprecated version of SessionState::new_with_config_rt to 41.0.0

6721e01

SessionState::new_with_config_rt was deprecated in favor of SessionStateBuilder in apache#11403 which is not shipped in 40.x.

kezhuw mentioned this pull request Aug 6, 2024

Bump deprecated version of SessionState::new_with_config_rt to 41.0.0 #11839

Merged

lewiszlw pushed a commit that referenced this pull request Aug 6, 2024

Bump deprecated version of SessionState::new_with_config_rt to 41.0.0 (…

d8bc7e2

…#11839) SessionState::new_with_config_rt was deprecated in favor of SessionStateBuilder in #11403 which is not shipped in 40.x.

Michael-J-Ward added a commit to Michael-J-Ward/datafusion-python that referenced this pull request Aug 10, 2024

use SessionStateBuilder

47239c7

The old constructor is deprecated. Ref: apache/datafusion#11403

Michael-J-Ward added a commit to Michael-J-Ward/datafusion-python that referenced this pull request Aug 20, 2024

use SessionStateBuilder

7207433

The old constructor is deprecated. Ref: apache/datafusion#11403

alamb mentioned this pull request Aug 23, 2024

Minor: Add example for configuring SessionContext #12139

Merged

alamb mentioned this pull request Sep 20, 2024

[EPIC] Easier extension configuration SessionState / SessionConfig #12550

Open

6 tasks

	pub fn new_from_existing(existing: &SessionState) -> Self {
	pub fn new_from_existing(existing: SessionState) -> Self {

Add SessionStateBuilder and extract out the registration of defaults #11403

Add SessionStateBuilder and extract out the registration of defaults #11403

Conversation

Omega359 commented Jul 10, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

jayzhan211 Jul 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jayzhan211 Jul 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Omega359 Jul 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Omega359 commented Jul 12, 2024

jayzhan211 left a comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Omega359 commented Jul 14, 2024

alamb left a comment

Choose a reason for hiding this comment

alamb commented Jul 14, 2024

alamb commented Jul 14, 2024

alamb commented Jul 16, 2024

alamb commented Jul 31, 2024

Omega359 commented Jul 10, 2024 •

edited

Loading

jayzhan211 Jul 12, 2024 •

edited

Loading

jayzhan211 Jul 12, 2024 •

edited

Loading

Omega359 Jul 12, 2024 •

edited

Loading