From 23ccca9cd9bdf1791985f2b4abb6d5d8a4fc5ece Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sat, 24 Aug 2024 17:09:31 -0400 Subject: [PATCH] Add example for configuring SessionContext (#12139) --- datafusion/core/src/execution/context/mod.rs | 42 +++++++++++++++++++- datafusion/core/src/execution/mod.rs | 2 + 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index 972a6f643733..49cbf9ebde9d 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -128,6 +128,9 @@ where /// the state of the connection between a user and an instance of the /// DataFusion engine. /// +/// See examples below for how to use the `SessionContext` to execute queries +/// and how to configure the session. +/// /// # Overview /// /// [`SessionContext`] provides the following functionality: @@ -200,7 +203,38 @@ where /// # } /// ``` /// -/// # `SessionContext`, `SessionState`, and `TaskContext` +/// # Example: Configuring `SessionContext` +/// +/// The `SessionContext` can be configured by creating a [`SessionState`] using +/// [`SessionStateBuilder`]: +/// +/// ``` +/// # use std::sync::Arc; +/// # use datafusion::prelude::*; +/// # use datafusion::execution::SessionStateBuilder; +/// # use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv}; +/// // Configure a 4k batch size +/// let config = SessionConfig::new() .with_batch_size(4 * 1024); +/// +/// // configure a memory limit of 1GB with 20% slop +/// let runtime_env = RuntimeEnv::new( +/// RuntimeConfig::new() +/// .with_memory_limit(1024 * 1024 * 1024, 0.80) +/// ).unwrap(); +/// +/// // Create a SessionState using the config and runtime_env +/// let state = SessionStateBuilder::new() +/// .with_config(config) +/// .with_runtime_env(Arc::new(runtime_env)) +/// // include support for built in functions and configurations +/// .with_default_features() +/// .build(); +/// +/// // Create a SessionContext +/// let ctx = SessionContext::from(state); +/// ``` +/// +/// # Relationship between `SessionContext`, `SessionState`, and `TaskContext` /// /// The state required to optimize, and evaluate queries is /// broken into three levels to allow tailoring @@ -1427,6 +1461,12 @@ impl From<&SessionContext> for TaskContext { } } +impl From for SessionContext { + fn from(state: SessionState) -> Self { + Self::new_with_state(state) + } +} + /// A planner used to add extensions to DataFusion logical and physical plans. #[async_trait] pub trait QueryPlanner { diff --git a/datafusion/core/src/execution/mod.rs b/datafusion/core/src/execution/mod.rs index a1b3eab25f33..10aa16ffe47a 100644 --- a/datafusion/core/src/execution/mod.rs +++ b/datafusion/core/src/execution/mod.rs @@ -19,6 +19,8 @@ pub mod context; pub mod session_state; +pub use session_state::{SessionState, SessionStateBuilder}; + mod session_state_defaults; pub use session_state_defaults::SessionStateDefaults;