From 82a6de948d17e927c6ed4bbca8606f98b04111c7 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 16:15:57 -0800 Subject: [PATCH] perf: reuse session context --- rust/lance-datafusion/src/planner.rs | 23 ++++------------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/rust/lance-datafusion/src/planner.rs b/rust/lance-datafusion/src/planner.rs index 4bdfe3c2508..958b839b940 100644 --- a/rust/lance-datafusion/src/planner.rs +++ b/rust/lance-datafusion/src/planner.rs @@ -7,6 +7,7 @@ use std::borrow::Cow; use std::collections::{BTreeSet, VecDeque}; use std::sync::Arc; +use crate::exec::{get_session_context, LanceExecutionOptions}; use crate::expr::safe_coerce_scalar; use crate::logical_expr::{coerce_filter_type_to_boolean, get_as_string_scalar_opt, resolve_expr}; use crate::sql::{parse_sql_expr, parse_sql_filter}; @@ -19,10 +20,7 @@ use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor use datafusion::common::DFSchema; use datafusion::config::ConfigOptions; use datafusion::error::Result as DFResult; -use datafusion::execution::config::SessionConfig; -use datafusion::execution::context::{SessionContext, SessionState}; -use datafusion::execution::runtime_env::RuntimeEnvBuilder; -use datafusion::execution::session_state::SessionStateBuilder; +use datafusion::execution::context::SessionState; use datafusion::logical_expr::expr::ScalarFunction; use datafusion::logical_expr::planner::{ExprPlanner, PlannerResult, RawFieldAccessExpr}; use datafusion::logical_expr::{ @@ -163,22 +161,9 @@ struct LanceContextProvider { impl Default for LanceContextProvider { fn default() -> Self { - let config = SessionConfig::new(); - let runtime = RuntimeEnvBuilder::new().build_arc().unwrap(); - - let ctx = SessionContext::new_with_config_rt(config.clone(), runtime.clone()); - crate::udf::register_functions(&ctx); - + let ctx = get_session_context(&LanceExecutionOptions::default()); let state = ctx.state(); - - // SessionState does not expose expr_planners, so we need to get them separately - let mut state_builder = SessionStateBuilder::new() - .with_config(config) - .with_runtime_env(runtime) - .with_default_features(); - - // unwrap safe because with_default_features sets expr_planners - let expr_planners = state_builder.expr_planners().as_ref().unwrap().clone(); + let expr_planners = state.expr_planners().to_vec(); Self { options: ConfigOptions::default(),