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.as_any()
61///     .downcast_ref::<SessionState>()
62///     .ok_or_else(|| exec_datafusion_err!("Failed to downcast Session to SessionState"))
63/// }
64/// ```
65///
66/// [`SessionState`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html
67/// [`TableProvider`]: https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProvider.html
68#[async_trait]
69pub trait Session: Send + Sync {
70    /// Return the session ID
71    fn session_id(&self) -> &str;
72
73    /// Return the [`SessionConfig`]
74    fn config(&self) -> &SessionConfig;
75
76    /// return the [`ConfigOptions`]
77    fn config_options(&self) -> &ConfigOptions {
78        self.config().options()
79    }
80
81    /// Creates a physical [`ExecutionPlan`] plan from a [`LogicalPlan`].
82    ///
83    /// Note: this will optimize the provided plan first.
84    ///
85    /// This function will error for [`LogicalPlan`]s such as catalog DDL like
86    /// `CREATE TABLE`, which do not have corresponding physical plans and must
87    /// be handled by another layer, typically the `SessionContext`.
88    async fn create_physical_plan(
89        &self,
90        logical_plan: &LogicalPlan,
91    ) -> Result<Arc<dyn ExecutionPlan>>;
92
93    /// Create a [`PhysicalExpr`] from an [`Expr`] after applying type
94    /// coercion, and function rewrites.
95    ///
96    /// Note: The expression is not simplified or otherwise optimized:  `a = 1
97    /// + 2` will not be simplified to `a = 3` as this is a more involved process.
98    /// See the [expr_api] example for how to simplify expressions.
99    ///
100    /// [expr_api]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/expr_api.rs
101    fn create_physical_expr(
102        &self,
103        expr: Expr,
104        df_schema: &DFSchema,
105    ) -> Result<Arc<dyn PhysicalExpr>>;
106
107    /// Return reference to scalar_functions
108    fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>>;
109
110    /// Return reference to aggregate_functions
111    fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>>;
112
113    /// Return reference to window functions
114    fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>>;
115
116    /// Return the runtime env
117    fn runtime_env(&self) -> &Arc<RuntimeEnv>;
118
119    /// Return the execution properties
120    fn execution_props(&self) -> &ExecutionProps;
121
122    fn as_any(&self) -> &dyn Any;
123
124    /// Return the table options
125    fn table_options(&self) -> &TableOptions;
126
127    /// return the TableOptions options with its extensions
128    fn default_table_options(&self) -> TableOptions {
129        self.table_options()
130            .combine_with_session_config(self.config_options())
131    }
132
133    /// Returns a mutable reference to [`TableOptions`]
134    fn table_options_mut(&mut self) -> &mut TableOptions;
135
136    /// Get a new TaskContext to run in this session
137    fn task_ctx(&self) -> Arc<TaskContext>;
138}
139
140/// Create a new task context instance from Session
141impl From<&dyn Session> for TaskContext {
142    fn from(state: &dyn Session) -> Self {
143        let task_id = None;
144        TaskContext::new(
145            task_id,
146            state.session_id().to_string(),
147            state.config().clone(),
148            state.scalar_functions().clone(),
149            state.aggregate_functions().clone(),
150            state.window_functions().clone(),
151            Arc::clone(state.runtime_env()),
152        )
153    }
154}
155type SessionRefLock = Arc<Mutex<Option<Weak<RwLock<dyn Session>>>>>;
156/// The state store that stores the reference of the runtime session state.
157#[derive(Debug)]
158pub struct SessionStore {
159    session: SessionRefLock,
160}
161
162impl SessionStore {
163    /// Create a new [SessionStore]
164    pub fn new() -> Self {
165        Self {
166            session: Arc::new(Mutex::new(None)),
167        }
168    }
169
170    /// Set the session state of the store
171    pub fn with_state(&self, state: Weak<RwLock<dyn Session>>) {
172        let mut lock = self.session.lock();
173        *lock = Some(state);
174    }
175
176    /// Get the current session of the store
177    pub fn get_session(&self) -> Weak<RwLock<dyn Session>> {
178        self.session.lock().clone().unwrap()
179    }
180}
181
182impl Default for SessionStore {
183    fn default() -> Self {
184        Self::new()
185    }
186}
datafusion_session/session.rs

datafusion_session/
session.rs
