Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add SessionStateBuilder and extract out the registration of defaults #11403

Merged
merged 10 commits into from
Jul 14, 2024
Merged

Add SessionStateBuilder and extract out the registration of defaults #11403

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
fab06da
Create a SessionStateBuilder and use it for creating anything but a b…
Omega359 Jul 10, 2024
de39082
Updated new_from_existing to take a reference to the existing Session…
Omega359 Jul 10, 2024
7d1bc09
Minor documentation update.
Omega359 Jul 10, 2024
490214d
SessionStateDefaults improvements.
Omega359 Jul 10, 2024
93a3a7d
Merge branch 'apache:main' into feature/11320_session_state_builder
Omega359 Jul 10, 2024
e3fad6d
Reworked how SessionStateBuilder works from PR feedback.
Omega359 Jul 12, 2024
8a579cf
Bug fix for missing array_expressions cfg feature.
Omega359 Jul 12, 2024
eb64e2f
Merge remote-tracking branch 'refs/remotes/upstream/main' into featur…
Omega359 Jul 14, 2024
71decf6
Review feedback updates + doc fixes for SessionStateDefaults
Omega359 Jul 14, 2024
380ce98
Cargo fmt update.
Omega359 Jul 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions datafusion-cli/src/catalog.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ use datafusion::datasource::listing::{
use datafusion::datasource::TableProvider;
use datafusion::error::Result;
use datafusion::execution::context::SessionState;
use datafusion::execution::session_state::SessionStateBuilder;

use async_trait::async_trait;
use dirs::home_dir;
Expand Down Expand Up @@ -162,6 +163,7 @@ impl SchemaProvider for DynamicFileSchemaProvider {
.ok_or_else(|| plan_datafusion_err!("locking error"))?
.read()
.clone();
let mut builder = SessionStateBuilder::new_from_existing(&state);
let optimized_name = substitute_tilde(name.to_owned());
let table_url = ListingTableUrl::parse(optimized_name.as_str())?;
let scheme = table_url.scheme();
Expand All @@ -178,13 +180,16 @@ 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());
builder =
builder.with_table_options_extension(AwsOptions::default());
}
"gs" | "gcs" => {
state = state.add_table_options_extension(GcpOptions::default())
builder =
builder.with_table_options_extension(GcpOptions::default())
}
_ => {}
};
state = builder.build();
let store = get_object_store(
&state,
table_url.scheme(),
Expand Down
9 changes: 6 additions & 3 deletions datafusion/core/src/execution/context/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,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")]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

pub fn with_state(state: SessionState) -> Self {
Self::new_with_state(state)
}
Expand Down Expand Up @@ -1574,6 +1574,7 @@ mod tests {
use datafusion_common_runtime::SpawnedTask;

use crate::catalog::schema::SchemaProvider;
use crate::execution::session_state::SessionStateBuilder;
use crate::physical_planner::PhysicalPlanner;
use async_trait::async_trait;
use tempfile::TempDir;
Expand Down Expand Up @@ -1734,8 +1735,10 @@ mod tests {
async fn custom_query_planner() -> Result<()> {
let runtime = Arc::new(RuntimeEnv::default());
let session_state =
SessionState::new_with_config_rt(SessionConfig::new(), runtime)
.with_query_planner(Arc::new(MyQueryPlanner {}));
SessionStateBuilder::new_with_config_rt(SessionConfig::new(), runtime)
.with_defaults(true)
.with_query_planner(Arc::new(MyQueryPlanner {}))
.build();
let ctx = SessionContext::new_with_state(session_state);

let df = ctx.sql("SELECT 1").await?;
Expand Down
Loading