datafusion_session/

session.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use async_trait::async_trait;
use datafusion_common::config::{ConfigOptions, TableOptions};
use datafusion_common::{DFSchema, Result};
use datafusion_execution::TaskContext;
use datafusion_execution::config::SessionConfig;
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::registry::ExtensionTypeRegistryRef;
use datafusion_expr::{
    AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
};
use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr};

use crate::CatalogProviderList;
use parking_lot::{Mutex, RwLock};
use std::any::Any;
use std::collections::HashMap;
use std::sync::{Arc, Weak};

use crate::{PhysicalOptimizerRule, QueryPlanner, UnsupportedQueryPlanner};

/// Interface for accessing [`SessionState`] from the catalog and data source.
///
/// This trait provides access to the information needed to plan and execute
/// queries, such as configuration, functions, and runtime environment. See the
/// documentation on [`SessionState`] for more information.
///
/// Historically, the `SessionState` struct was passed directly to catalog
/// traits such as [`TableProvider`], which required a direct dependency on the
/// DataFusion core. The interface required is now defined by this trait. See
/// [#10782] for more details.
///
/// [#10782]: https://github.com/apache/datafusion/issues/10782
///
/// # Migration from `SessionState`
///
/// Using trait methods is preferred, as the implementation may change in future
/// versions. However, you can downcast a `Session` to a `SessionState` as shown
/// in the example below. If you find yourself needing to do this, please open
/// an issue on the DataFusion repository so we can extend the trait to provide
/// the required information.
///
/// ```
/// # use datafusion_session::Session;
/// # use datafusion_common::{Result, exec_datafusion_err};
/// # struct SessionState {}
/// // Given a `Session` reference, get the concrete `SessionState` reference
/// // Note: this may stop working in future versions,
/// fn session_state_from_session(session: &dyn Session) -> Result<&SessionState> {
///     session
///         .as_any()
///         .downcast_ref::<SessionState>()
///         .ok_or_else(|| {
///             exec_datafusion_err!("Failed to downcast Session to SessionState")
///         })
/// }
/// ```
///
/// [`SessionState`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html
/// [`TableProvider`]: https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProvider.html
#[async_trait]
pub trait Session: Send + Sync {
    /// Return the session ID
    fn session_id(&self) -> &str;

    /// Return the [`SessionConfig`]
    fn config(&self) -> &SessionConfig;

    /// Return the catalogs registered with this session.
    fn catalog_list(&self) -> Arc<dyn CatalogProviderList>;

    /// return the [`ConfigOptions`]
    fn config_options(&self) -> &ConfigOptions {
        self.config().options()
    }

    /// Return the query planner for this session.
    ///
    /// # Warning
    ///
    /// The default implementation returns an [`UnsupportedQueryPlanner`], so
    /// [`Session::create_physical_plan`] will fail. Sessions that support
    /// physical planning should override this method (for example by returning
    /// `SessionState::query_planner`).
    fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync> {
        Arc::new(UnsupportedQueryPlanner)
    }

    /// Optimize a logical plan.
    ///
    /// # Warning
    ///
    /// The default implementation returns the plan **unchanged**, applying no
    /// logical optimizations whatsoever. This is almost never what you want:
    /// without optimization, queries execute in their naive, unoptimized form
    /// and may be dramatically slower or fail to run at all. The default exists
    /// only so this crate need not depend on the optimizer; any real session
    /// should override this method (for example by delegating to
    /// `SessionState::optimize`).
    fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan> {
        Ok(plan.clone())
    }

    /// Return the physical optimizer rules for this session.
    ///
    /// # Warning
    ///
    /// The default implementation returns **no rules**. This is almost never
    /// what you want: DataFusion relies on physical optimizer rules for
    /// correctness-critical rewrites (such as inserting the repartitioning and
    /// coalescing needed for parallel and multi-partition execution), so a
    /// session with no rules will produce plans that are inefficient or that
    /// fail to execute. The default exists only so this crate need not depend
    /// on the optimizer; any real session should override this method (for
    /// example by returning `SessionState::physical_optimizers`).
    fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] {
        &[]
    }

    /// Return the optional statistics registry used during physical optimization.
    fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
        None
    }

    /// Creates a physical [`ExecutionPlan`] plan from a [`LogicalPlan`].
    ///
    /// Note: this will optimize the provided plan first.
    ///
    /// This function will error for [`LogicalPlan`]s such as catalog DDL like
    /// `CREATE TABLE`, which do not have corresponding physical plans and must
    /// be handled by another layer, typically the `SessionContext`.
    async fn create_physical_plan(
        &self,
        logical_plan: &LogicalPlan,
    ) -> Result<Arc<dyn ExecutionPlan>>;

    /// Create a [`PhysicalExpr`] from an [`Expr`] after applying type
    /// coercion, and function rewrites.
    ///
    /// Note: The expression is not simplified or otherwise optimized:  `a = 1
    /// + 2` will not be simplified to `a = 3` as this is a more involved process.
    /// See the [expr_api] example for how to simplify expressions.
    ///
    /// [expr_api]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/query_planning/expr_api.rs
    fn create_physical_expr(
        &self,
        expr: Expr,
        df_schema: &DFSchema,
    ) -> Result<Arc<dyn PhysicalExpr>>;

    /// Return reference to scalar_functions
    fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>>;

    /// Return reference to higher_order_functions
    fn higher_order_functions(&self) -> &HashMap<String, Arc<HigherOrderUDF>>;

    /// Return reference to aggregate_functions
    fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>>;

    /// Return reference to window functions
    fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>>;

    /// Return a reference to the extension type registry
    fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef;

    /// Return the runtime env
    fn runtime_env(&self) -> &Arc<RuntimeEnv>;

    /// Return the execution properties
    fn execution_props(&self) -> &ExecutionProps;

    fn as_any(&self) -> &dyn Any;

    /// Return the table options
    fn table_options(&self) -> &TableOptions;

    /// return the TableOptions options with its extensions
    fn default_table_options(&self) -> TableOptions {
        self.table_options()
            .combine_with_session_config(self.config_options())
    }

    /// Returns a mutable reference to [`TableOptions`]
    fn table_options_mut(&mut self) -> &mut TableOptions;

    /// Get a new TaskContext to run in this session
    fn task_ctx(&self) -> Arc<TaskContext>;
}

/// Create a new task context instance from Session
impl From<&dyn Session> for TaskContext {
    fn from(state: &dyn Session) -> Self {
        let task_id = None;
        TaskContext::new(
            task_id,
            state.session_id().to_string(),
            state.config().clone(),
            state.scalar_functions().clone(),
            state.higher_order_functions().clone(),
            state.aggregate_functions().clone(),
            state.window_functions().clone(),
            Arc::clone(state.runtime_env()),
        )
    }
}
type SessionRefLock = Arc<Mutex<Option<Weak<RwLock<dyn Session>>>>>;
/// The state store that stores the reference of the runtime session state.
#[derive(Debug)]
pub struct SessionStore {
    session: SessionRefLock,
}

impl SessionStore {
    /// Create a new [SessionStore]
    pub fn new() -> Self {
        Self {
            session: Arc::new(Mutex::new(None)),
        }
    }

    /// Set the session state of the store
    pub fn with_state(&self, state: Weak<RwLock<dyn Session>>) {
        let mut lock = self.session.lock();
        *lock = Some(state);
    }

    /// Get the current session of the store
    pub fn get_session(&self) -> Weak<RwLock<dyn Session>> {
        self.session.lock().clone().unwrap()
    }
}

impl Default for SessionStore {
    fn default() -> Self {
        Self::new()
    }
}