datafusion_session/
session.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use async_trait::async_trait;
19use datafusion_common::config::{ConfigOptions, TableOptions};
20use datafusion_common::{DFSchema, Result};
21use datafusion_execution::config::SessionConfig;
22use datafusion_execution::runtime_env::RuntimeEnv;
23use datafusion_execution::TaskContext;
24use datafusion_expr::execution_props::ExecutionProps;
25use datafusion_expr::{AggregateUDF, Expr, LogicalPlan, ScalarUDF, WindowUDF};
26use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr};
27use parking_lot::{Mutex, RwLock};
28use std::any::Any;
29use std::collections::HashMap;
30use std::sync::{Arc, Weak};
31
32/// Interface for accessing [`SessionState`] from the catalog and data source.
33///
34/// This trait provides access to the information needed to plan and execute
35/// queries, such as configuration, functions, and runtime environment. See the
36/// documentation on [`SessionState`] for more information.
37///
38/// Historically, the `SessionState` struct was passed directly to catalog
39/// traits such as [`TableProvider`], which required a direct dependency on the
40/// DataFusion core. The interface required is now defined by this trait. See
41/// [#10782] for more details.
42///
43/// [#10782]: https://github.com/apache/datafusion/issues/10782
44///
45/// # Migration from `SessionState`
46///
47/// Using trait methods is preferred, as the implementation may change in future
48/// versions. However, you can downcast a `Session` to a `SessionState` as shown
49/// in the example below. If you find yourself needing to do this, please open
50/// an issue on the DataFusion repository so we can extend the trait to provide
51/// the required information.
52///
53/// ```
54/// # use datafusion_session::Session;
55/// # use datafusion_common::{Result, exec_datafusion_err};
56/// # struct SessionState {}
57/// // Given a `Session` reference, get the concrete `SessionState` reference
58/// // Note: this may stop working in future versions,
59/// fn session_state_from_session(session: &dyn Session) -> Result<&SessionState> {
60///     session
61///         .as_any()
62///         .downcast_ref::<SessionState>()
63///         .ok_or_else(|| {
64///             exec_datafusion_err!("Failed to downcast Session to SessionState")
65///         })
66/// }
67/// ```
68///
69/// [`SessionState`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html
70/// [`TableProvider`]: https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProvider.html
71#[async_trait]
72pub trait Session: Send + Sync {
73    /// Return the session ID
74    fn session_id(&self) -> &str;
75
76    /// Return the [`SessionConfig`]
77    fn config(&self) -> &SessionConfig;
78
79    /// return the [`ConfigOptions`]
80    fn config_options(&self) -> &ConfigOptions {
81        self.config().options()
82    }
83
84    /// Creates a physical [`ExecutionPlan`] plan from a [`LogicalPlan`].
85    ///
86    /// Note: this will optimize the provided plan first.
87    ///
88    /// This function will error for [`LogicalPlan`]s such as catalog DDL like
89    /// `CREATE TABLE`, which do not have corresponding physical plans and must
90    /// be handled by another layer, typically the `SessionContext`.
91    async fn create_physical_plan(
92        &self,
93        logical_plan: &LogicalPlan,
94    ) -> Result<Arc<dyn ExecutionPlan>>;
95
96    /// Create a [`PhysicalExpr`] from an [`Expr`] after applying type
97    /// coercion, and function rewrites.
98    ///
99    /// Note: The expression is not simplified or otherwise optimized:  `a = 1
100    /// + 2` will not be simplified to `a = 3` as this is a more involved process.
101    /// See the [expr_api] example for how to simplify expressions.
102    ///
103    /// [expr_api]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/expr_api.rs
104    fn create_physical_expr(
105        &self,
106        expr: Expr,
107        df_schema: &DFSchema,
108    ) -> Result<Arc<dyn PhysicalExpr>>;
109
110    /// Return reference to scalar_functions
111    fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>>;
112
113    /// Return reference to aggregate_functions
114    fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>>;
115
116    /// Return reference to window functions
117    fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>>;
118
119    /// Return the runtime env
120    fn runtime_env(&self) -> &Arc<RuntimeEnv>;
121
122    /// Return the execution properties
123    fn execution_props(&self) -> &ExecutionProps;
124
125    fn as_any(&self) -> &dyn Any;
126
127    /// Return the table options
128    fn table_options(&self) -> &TableOptions;
129
130    /// return the TableOptions options with its extensions
131    fn default_table_options(&self) -> TableOptions {
132        self.table_options()
133            .combine_with_session_config(self.config_options())
134    }
135
136    /// Returns a mutable reference to [`TableOptions`]
137    fn table_options_mut(&mut self) -> &mut TableOptions;
138
139    /// Get a new TaskContext to run in this session
140    fn task_ctx(&self) -> Arc<TaskContext>;
141}
142
143/// Create a new task context instance from Session
144impl From<&dyn Session> for TaskContext {
145    fn from(state: &dyn Session) -> Self {
146        let task_id = None;
147        TaskContext::new(
148            task_id,
149            state.session_id().to_string(),
150            state.config().clone(),
151            state.scalar_functions().clone(),
152            state.aggregate_functions().clone(),
153            state.window_functions().clone(),
154            Arc::clone(state.runtime_env()),
155        )
156    }
157}
158type SessionRefLock = Arc<Mutex<Option<Weak<RwLock<dyn Session>>>>>;
159/// The state store that stores the reference of the runtime session state.
160#[derive(Debug)]
161pub struct SessionStore {
162    session: SessionRefLock,
163}
164
165impl SessionStore {
166    /// Create a new [SessionStore]
167    pub fn new() -> Self {
168        Self {
169            session: Arc::new(Mutex::new(None)),
170        }
171    }
172
173    /// Set the session state of the store
174    pub fn with_state(&self, state: Weak<RwLock<dyn Session>>) {
175        let mut lock = self.session.lock();
176        *lock = Some(state);
177    }
178
179    /// Get the current session of the store
180    pub fn get_session(&self) -> Weak<RwLock<dyn Session>> {
181        self.session.lock().clone().unwrap()
182    }
183}
184
185impl Default for SessionStore {
186    fn default() -> Self {
187        Self::new()
188    }
189}
datafusion_session/session.rs

datafusion_session/
session.rs
