linkrs_query/executor/base/
execution_context.rs

1use parking_lot::RwLock;
2use std::collections::HashMap;
3use std::sync::Arc;
4
5use super::execution_result::ExecutionResult;
6use super::traits::SearchProvider;
7use super::MemoryBudget;
8use crate::executor::expression::functions::global_registry_ref;
9use crate::executor::expression::functions::OwnedFunctionRef;
10use crate::executor::streaming::pool::SharedScheduler;
11use crate::optimizer::stats::feedback::history::QueryFeedbackHistory;
12use crate::optimizer::JoinAlgorithm;
13use crate::storage::QueryStorage;
14use linkrs_core::metadata::SequenceManager;
15use linkrs_core::metadata::{MacroManager, TypeAliasManager};
16use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
17use linkrs_core::Arena;
18use linkrs_core::Value;
19#[cfg(feature = "fulltext")]
20use linkrs_fulltext::manager::FulltextIndexManager;
21#[cfg(feature = "vector")]
22use linkrs_sync::VectorSyncCoordinator;
23
24/// Search-related runtime state bundled into a single sub-struct.
25///
26/// Extracted from `ExecutionContext` to reduce feature-flag pollution:
27/// adding a new search backend only requires touching this struct and its
28/// `Default` impl instead of every `ExecutionContext` constructor.
29#[derive(Debug, Clone, Default)]
30pub struct SearchContext {
31    #[cfg(feature = "fulltext")]
32    pub fulltext_manager: Option<Arc<FulltextIndexManager>>,
33    #[cfg(feature = "vector")]
34    pub vector_coordinator: Option<Arc<VectorSyncCoordinator>>,
35    /// Unified search providers for discovery and enumeration.
36    pub search_providers: Vec<Arc<dyn SearchProvider>>,
37}
38
39#[derive(Debug, Clone)]
40pub struct ExecutionContext {
41    pub results: Arc<RwLock<HashMap<String, ExecutionResult>>>,
42    pub variables: Arc<RwLock<HashMap<String, linkrs_core::Value>>>,
43    pub expression_context: Arc<ExpressionAnalysisContext>,
44    /// Aggregated search backends (fulltext, vector, etc.).
45    pub search: SearchContext,
46    pub storage: Option<Arc<RwLock<dyn QueryStorage>>>,
47    /// Snapshot handle pinned by the bound storage, when the storage was
48    /// bound to a read/auto-commit operation context (storage boundary).
49    ///
50    /// Populated by the pipeline when it binds the per-query storage; lets
51    /// the execution layer observe which snapshot the query reads at.
52    pub bound_snapshot: Option<crate::storage::SnapshotHandle>,
53    pub space_name: Option<String>,
54    pub parameters: Arc<HashMap<String, linkrs_core::Value>>,
55    /// Session variable snapshot for this query (resolves
56    /// `Expression::SessionVariable`); captured once per statement.
57    pub session_variables: Arc<HashMap<String, linkrs_core::Value>>,
58    /// Per-query memory budget for blocking operators.
59    pub memory_budget: MemoryBudget,
60    /// Maximum intra-query workers. 1 = serial only.
61    pub max_workers: usize,
62    /// Server-assigned query ID for KILL QUERY / cancellation support.
63    pub query_id: u64,
64    /// Request-scoped cancellation token threaded to the runtime/registry.
65    pub cancel_token: Option<crate::executor::streaming::query_registry::CancelToken>,
66    /// Streaming chunk size (rows per chunk).
67    pub chunk_size: usize,
68    /// Maximum buffered chunks before back-pressure.
69    pub max_buffered_chunks: usize,
70    /// Engine-level shared scheduler.
71    pub shared_scheduler: Option<Arc<SharedScheduler>>,
72    /// Optional thread-safe bumpalo arena for executor temporary allocations.
73    pub arena: Option<Arc<parking_lot::Mutex<Arena>>>,
74    /// Shared query feedback history for collecting execution statistics.
75    ///
76    /// Injected from the optimizer engine by the pipeline; the execution
77    /// instance records estimated-vs-actual feedback here after execution.
78    pub feedback_history: Option<Arc<QueryFeedbackHistory>>,
79    /// Shared cross-query policy for the typed columnar chunk layout.
80    ///
81    /// Injected from the optimizer engine by the pipeline; each query merges
82    /// its columnar hit/miss counts back into the policy at completion.
83    pub columnar_policy: Option<Arc<crate::executor::streaming::chunk::ColumnarPolicy>>,
84    /// Cost-based join algorithm decisions keyed by planner node id
85    /// (CBO join-order decision channel).
86    ///
87    /// Populated by the optimizer from the join reorder walker and consumed
88    /// by the arena builder when converting `InnerJoin`/`LeftJoin` nodes.
89    /// Absent keys fall back to the default heuristic (hash join for valid
90    /// equi keys).
91    pub join_algorithms: HashMap<i64, JoinAlgorithm>,
92    /// Transaction isolation level for this execution, when running inside
93    /// an explicit transaction. Execution-time knob threaded from the
94    /// API layer through [`crate::QueryContext`]; `None` = auto-commit.
95    pub isolation_level: Option<linkrs_core::types::TransactionIsolationLevel>,
96    /// Consistency requirement for secondary-index reads.
97    /// `None` = eventual; `Some(cfg)` = read-your-writes with timeout and optional LSN.
98    pub ryw_config: Option<linkrs_core::types::ReadYourWritesConfig>,
99    /// Sequence manager for curr_val/next_val functions.
100    pub sequence_manager: Option<Arc<SequenceManager>>,
101    /// Macro catalog manager for user-defined macros (expansion + DDL).
102    pub macro_manager: Option<Arc<MacroManager>>,
103    /// Type-alias catalog manager for user-defined types.
104    pub type_alias_manager: Option<Arc<TypeAliasManager>>,
105}
106
107/// Internal: build the non-search portion of an `ExecutionContext`.
108fn new_base(expression_context: Arc<ExpressionAnalysisContext>) -> ExecutionContext {
109    ExecutionContext {
110        results: Arc::new(RwLock::new(HashMap::new())),
111        variables: Arc::new(RwLock::new(HashMap::new())),
112        expression_context,
113        search: SearchContext::default(),
114        storage: None,
115        bound_snapshot: None,
116        space_name: None,
117        parameters: Arc::new(HashMap::new()),
118        session_variables: Arc::new(HashMap::new()),
119        memory_budget: MemoryBudget::default_budget(),
120        max_workers: 1,
121        query_id: 0,
122        cancel_token: None,
123        chunk_size: ExecutionContext::DEFAULT_CHUNK_SIZE,
124        max_buffered_chunks: ExecutionContext::DEFAULT_MAX_BUFFERED_CHUNKS,
125        shared_scheduler: None,
126        arena: None,
127        feedback_history: None,
128        columnar_policy: None,
129        join_algorithms: HashMap::new(),
130        isolation_level: None,
131        ryw_config: None,
132        sequence_manager: None,
133        macro_manager: None,
134        type_alias_manager: None,
135    }
136}
137
138impl ExecutionContext {
139    pub const DEFAULT_CHUNK_SIZE: usize = 2048;
140    pub const DEFAULT_MAX_BUFFERED_CHUNKS: usize = 10;
141
142    pub fn new(expression_context: Arc<ExpressionAnalysisContext>) -> Self {
143        new_base(expression_context)
144    }
145
146    pub fn with_parameters(
147        expression_context: Arc<ExpressionAnalysisContext>,
148        parameters: HashMap<String, linkrs_core::Value>,
149    ) -> Self {
150        let mut ctx = new_base(expression_context);
151        ctx.parameters = Arc::new(parameters);
152        ctx
153    }
154
155    /// Create a context with memory limit configured from a config value.
156    ///
157    /// When `max_memory_per_query` is non-zero it overrides the default
158    /// 512 MiB budget; otherwise the default is used.
159    pub fn with_memory_limit(
160        expression_context: Arc<ExpressionAnalysisContext>,
161        max_memory_per_query: u64,
162    ) -> Self {
163        let budget = if max_memory_per_query > 0 {
164            MemoryBudget::new(max_memory_per_query as usize)
165        } else {
166            MemoryBudget::default_budget()
167        };
168        let mut ctx = new_base(expression_context);
169        ctx.memory_budget = budget;
170        ctx
171    }
172
173    pub fn set_result(&self, name: String, result: ExecutionResult) {
174        self.results.write().insert(name, result);
175    }
176
177    pub fn get_result(&self, name: &str) -> Option<ExecutionResult> {
178        self.results.write().get(name).cloned()
179    }
180
181    pub fn set_variable(&self, name: String, value: linkrs_core::Value) {
182        self.variables.write().insert(name, value);
183    }
184
185    pub fn get_variable(&self, name: &str) -> Option<linkrs_core::Value> {
186        self.variables.write().get(name).cloned()
187    }
188
189    pub fn expression_context(&self) -> &Arc<ExpressionAnalysisContext> {
190        &self.expression_context
191    }
192
193    pub fn get_param(&self, name: &str) -> Option<&linkrs_core::Value> {
194        self.parameters.get(name)
195    }
196
197    pub fn current_space_id(&self) -> Option<u64> {
198        self.variables
199            .write()
200            .get("space_id")
201            .and_then(|v| match v {
202                Value::Int(id) => Some(*id as u64),
203                _ => None,
204            })
205    }
206
207    pub fn set_space_id(&self, space_id: u64) {
208        self.variables
209            .write()
210            .insert("space_id".into(), Value::Int(space_id as i32));
211    }
212
213    pub fn sequence_manager(&self) -> Option<&Arc<SequenceManager>> {
214        self.sequence_manager.as_ref()
215    }
216
217    pub fn set_sequence_manager(&mut self, manager: Arc<SequenceManager>) {
218        self.sequence_manager = Some(manager);
219    }
220
221    pub fn macro_manager(&self) -> Option<&Arc<MacroManager>> {
222        self.macro_manager.as_ref()
223    }
224
225    pub fn set_macro_manager(&mut self, manager: Arc<MacroManager>) {
226        self.macro_manager = Some(manager);
227    }
228
229    pub fn type_alias_manager(&self) -> Option<&Arc<TypeAliasManager>> {
230        self.type_alias_manager.as_ref()
231    }
232
233    pub fn set_type_alias_manager(&mut self, manager: Arc<TypeAliasManager>) {
234        self.type_alias_manager = Some(manager);
235    }
236}
237
238impl Default for ExecutionContext {
239    fn default() -> Self {
240        let mut ctx = new_base(Arc::new(ExpressionAnalysisContext::new()));
241        // Default constructors use default chunk size / max buffered chunks.
242        ctx.chunk_size = Self::DEFAULT_CHUNK_SIZE;
243        ctx.max_buffered_chunks = Self::DEFAULT_MAX_BUFFERED_CHUNKS;
244        ctx
245    }
246}
247
248impl crate::executor::expression::evaluator::traits::ExpressionContext for ExecutionContext {
249    fn get_variable(&self, name: &str) -> Option<Value> {
250        self.variables.write().get(name).cloned()
251    }
252
253    fn set_variable(&mut self, name: String, value: Value) {
254        self.variables.write().insert(name, value);
255    }
256
257    fn get_function(&self, name: &str) -> Option<OwnedFunctionRef> {
258        let registry = global_registry_ref();
259        registry
260            .get_builtin(name)
261            .map(|f| OwnedFunctionRef::Builtin(f.clone()))
262            .or_else(|| {
263                registry
264                    .get_custom(name)
265                    .map(|f| OwnedFunctionRef::Custom(f.clone()))
266            })
267    }
268
269    fn get_sequence_manager(&self) -> Option<Arc<SequenceManager>> {
270        self.sequence_manager.clone()
271    }
272}
linkrs_query/executor/base/execution_context.rs

linkrs_query/executor/base/
execution_context.rs
