From 23ccca9cd9bdf1791985f2b4abb6d5d8a4fc5ece Mon Sep 17 00:00:00 2001
From: Andrew Lamb <andrew@nerdnetworks.org>
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<SessionState> 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;