linkrs_query/
pipeline.rs

1mod compiler;
2mod diagnostics;
3mod execution;
4mod frontend;
5mod merge_conditional;
6mod prepared;
7
8use crate::cache::{ParameterizedQueryHandler, PlanCacheConfig, QueryPlanCache};
9use crate::executor::streaming::plan::PhysicalPlan;
10use crate::executor::streaming::pool::SharedScheduler;
11use crate::executor::streaming::query_registry::QueryRegistry;
12use crate::executor::streaming::SessionTransactionController;
13use crate::extensions::ExtensionRegistry;
14use crate::optimizer::OptimizerEngine;
15use crate::query_manager::QueryManager;
16use crate::storage::QueryStorage;
17use linkrs_core::metadata::index_manager::IndexMetadataManager;
18use linkrs_core::metadata::SchemaManager;
19#[cfg(feature = "fulltext")]
20use linkrs_fulltext::manager::FulltextIndexManager;
21use linkrs_metrics::StatsManager;
22#[cfg(feature = "vector")]
23use linkrs_sync::vector_sync::VectorSyncCoordinator;
24use linkrs_sync::SyncManager;
25use parking_lot::RwLock;
26use std::collections::HashMap;
27use std::sync::atomic::{AtomicU64, Ordering};
28use std::sync::Arc;
29
30/// Capacity of the DML shape caches below. Single-entry caching thrashed on
31/// alternating shapes (A,B,A,B re-parsed and re-bound every time); a small
32/// bounded map keeps the common bulk-load shapes resident while bounding
33/// memory. Lookup is a linear scan, which is cheap at this size.
34pub(crate) const DML_SHAPE_CACHE_CAP: usize = 16;
35
36pub struct QueryPipelineManager<S: QueryStorage + 'static> {
37    pub(crate) stats_manager: Arc<StatsManager>,
38    pub(crate) optimizer_engine: Arc<OptimizerEngine>,
39    pub(crate) plan_cache: Arc<QueryPlanCache>,
40    pub(crate) param_handler: ParameterizedQueryHandler,
41    pub(crate) schema_manager: Option<Arc<SchemaManager>>,
42    pub(crate) index_manager: Option<Arc<dyn IndexMetadataManager>>,
43    #[cfg(feature = "fulltext")]
44    pub(crate) fulltext_manager: Option<Arc<FulltextIndexManager>>,
45    #[cfg(feature = "vector")]
46    pub(crate) vector_coordinator: Option<Arc<VectorSyncCoordinator>>,
47    pub(crate) storage: Option<Arc<RwLock<S>>>,
48    pub(crate) sync_manager: Option<Arc<SyncManager>>,
49    pub(crate) schema_generation: Arc<AtomicU64>,
50    pub(crate) index_generation: Arc<AtomicU64>,
51    pub(crate) query_registry: Option<Arc<QueryRegistry>>,
52    /// Optional global QueryManager shared with the server so per-query runtimes
53    /// can forward KILL/finish/progress events into one registry. None until the
54    /// assembly installs it after the pipeline is built.
55    pub(crate) query_manager: Option<Arc<QueryManager>>,
56    /// Engine-level shared scheduler, created once and reused across queries.
57    pub(crate) shared_scheduler: Option<Arc<SharedScheduler>>,
58    pub(crate) session_controller: parking_lot::RwLock<Option<Arc<SessionTransactionController>>>,
59    /// Shared macro catalog (user-defined macros), engine-wide.
60    ///
61    /// Created once with the pipeline manager and shared with every query
62    /// through the binder (expansion), parser-adjacent planning and the
63    /// execution runtime (DDL). Persistence/WAL backends can be attached by
64    /// replacing this instance before serving queries.
65    pub macro_manager: Arc<linkrs_core::metadata::MacroManager>,
66    /// Shared type-alias catalog (user-defined types), engine-wide.
67    ///
68    /// Threaded the same way as [`Self::macro_manager`].
69    pub type_alias_manager: Arc<linkrs_core::metadata::TypeAliasManager>,
70    /// Shared sequence catalog (named sequences for `curr_val`/`next_val`),
71    /// engine-wide.
72    ///
73    /// Threaded the same way as [`Self::macro_manager`].
74    pub sequence_manager: Arc<linkrs_core::metadata::SequenceManager>,
75    /// Serializes statistics collection to prevent concurrent re-collection.
76    pub(crate) statistics_collect_lock: Arc<parking_lot::Mutex<()>>,
77    /// Sample cap for per-tag/per-edge-type degree estimation during collection.
78    pub(crate) statistics_sample_limit: usize,
79    /// Minimum data-epoch advance that triggers an automatic refresh.
80    ///
81    /// The epoch bumps on every write, so without a threshold any commit
82    /// invalidates statistics and the next query pays a full re-collection.
83    /// Smaller advances reuse the last collection; explicit ANALYZE still
84    /// forces a refresh.
85    pub(crate) statistics_min_epoch_delta: u64,
86    /// Display label of the global cost profile (e.g. `Ssd`); empty until
87    /// the assembly provides one. Only used for EXPLAIN observability.
88    pub(crate) default_cost_profile_label: String,
89    /// Display labels of per-space cost profile overrides (space -> label).
90    pub(crate) space_cost_profile_labels: HashMap<String, String>,
91    pub(crate) dml_shape_cache_enabled: bool,
92    /// Fast path for repeated same-shape DML: skips bind and serves the last
93    /// compiled physical plans directly. Bounded multi-entry map keyed by
94    /// [`DmlPlanMemoKey`]; see that struct for the key dimensions.
95    pub(crate) dml_plan_memo: parking_lot::Mutex<Vec<DmlPlanMemoEntry>>,
96    pub last_dml_plan_hits: std::sync::atomic::AtomicU64,
97    /// Parsed ASTs of normalized DML templates (`INSERT ... VALUES
98    /// (@__dml_0, ...)`), avoiding a re-parse per execution. Bounded
99    /// multi-entry map keyed by exact normalized text.
100    pub(crate) dml_template_ast:
101        parking_lot::Mutex<Vec<(String, Arc<crate::parser::ast::stmt::Ast>)>>,
102    pub(crate) dml_template_ast_parse_count: std::sync::atomic::AtomicU64,
103    pub(crate) dml_template_ast_hit_count: std::sync::atomic::AtomicU64,
104    pub(crate) dml_bind_skipped_count: std::sync::atomic::AtomicU64,
105    /// Experimental pipeline extensions. Empty by default,
106    /// leaving the builtin pipeline untouched.
107    pub(crate) extensions: Arc<ExtensionRegistry>,
108    /// Retained password-history depth threaded into every query context so
109    /// user-management plans rotate history uniformly across channels.
110    pub(crate) password_history_depth: usize,
111}
112
113/// Lookup key for one [`DmlPlanMemoEntry`].
114///
115/// Mirrors the [`crate::cache::PlanCacheKey`] dimensions that are computable
116/// *before* binding (the whole point of the memo is to skip bind): query
117/// shape, space, schema/index generations, parameter type signature, and the
118/// optimizer/planning configuration. The post-planning partition fingerprint
119/// used by the plan cache cannot be part of this key; instead memo hits are
120/// validated with [`crate::executor::streaming::plan::PhysicalPlanValidator::check_compatibility`]
121/// before being served.
122#[derive(Debug, Clone, PartialEq, Eq)]
123pub(crate) struct DmlPlanMemoKey {
124    pub normalized_text: String,
125    pub space_name: Option<String>,
126    pub schema_version: Option<u64>,
127    pub index_version: Option<u64>,
128    pub param_sig: u64,
129    pub optimizer_version: u64,
130    pub planning_config_hash: u64,
131}
132
133pub(crate) struct DmlPlanMemoEntry {
134    pub key: DmlPlanMemoKey,
135    pub plan: Arc<PhysicalPlan>,
136}
137
138impl<S: QueryStorage + 'static> QueryPipelineManager<S> {
139    /// Fallback sampling window used when the server assembly did not pass
140    /// an explicit `statistics_sample_limit` (mirrors the config default).
141    pub const DEFAULT_STATISTICS_SAMPLE_LIMIT: usize = 10_000;
142    /// Fallback epoch delta used when the server assembly did not pass an
143    /// explicit `statistics_min_epoch_delta` (mirrors the config default).
144    pub const DEFAULT_STATISTICS_MIN_EPOCH_DELTA: u64 = 100;
145
146    pub fn with_optimizer(
147        storage: Arc<RwLock<S>>,
148        stats_manager: Arc<StatsManager>,
149        optimizer_engine: Arc<OptimizerEngine>,
150    ) -> Self {
151        let plan_cache =
152            Arc::new(QueryPlanCache::default().with_stats_manager(stats_manager.clone()));
153        let param_handler = ParameterizedQueryHandler;
154
155        optimizer_engine.set_cte_cache_stats_manager(stats_manager.clone());
156
157        log::info!("Query pipeline manager created, using optimizer engine and query plan cache");
158
159        Self {
160            stats_manager,
161            optimizer_engine,
162            plan_cache,
163            param_handler,
164            schema_manager: None,
165            index_manager: None,
166            #[cfg(feature = "fulltext")]
167            fulltext_manager: None,
168            #[cfg(feature = "vector")]
169            vector_coordinator: None,
170            storage: Some(storage),
171            sync_manager: None,
172            schema_generation: Arc::new(AtomicU64::new(1)),
173            index_generation: Arc::new(AtomicU64::new(1)),
174            query_registry: None,
175            shared_scheduler: None,
176            query_manager: None,
177            macro_manager: Arc::new(linkrs_core::metadata::MacroManager::new()),
178            type_alias_manager: Arc::new(linkrs_core::metadata::TypeAliasManager::new()),
179            sequence_manager: Arc::new(linkrs_core::metadata::SequenceManager::new()),
180            session_controller: parking_lot::RwLock::new(None),
181            statistics_collect_lock: Arc::new(parking_lot::Mutex::new(())),
182            statistics_sample_limit: Self::DEFAULT_STATISTICS_SAMPLE_LIMIT,
183            statistics_min_epoch_delta: Self::DEFAULT_STATISTICS_MIN_EPOCH_DELTA,
184            default_cost_profile_label: String::new(),
185            space_cost_profile_labels: HashMap::new(),
186            dml_shape_cache_enabled: true,
187            dml_plan_memo: parking_lot::Mutex::new(Vec::new()),
188            last_dml_plan_hits: std::sync::atomic::AtomicU64::new(0),
189            dml_template_ast: parking_lot::Mutex::new(Vec::new()),
190            dml_template_ast_parse_count: std::sync::atomic::AtomicU64::new(0),
191            dml_template_ast_hit_count: std::sync::atomic::AtomicU64::new(0),
192            dml_bind_skipped_count: std::sync::atomic::AtomicU64::new(0),
193            extensions: Arc::new(ExtensionRegistry::new()),
194            password_history_depth: 0,
195        }
196    }
197
198    pub fn with_statistics_sample_limit(mut self, sample_limit: usize) -> Self {
199        self.statistics_sample_limit = sample_limit.max(1);
200        self
201    }
202
203    pub fn with_statistics_min_epoch_delta(mut self, delta: u64) -> Self {
204        self.statistics_min_epoch_delta = delta;
205        self
206    }
207
208    /// Set the retained password-history depth for user-management
209    /// operations (assembly-time only).
210    pub fn with_password_history_depth(mut self, depth: usize) -> Self {
211        self.password_history_depth = depth;
212        self
213    }
214
215    /// Replace the optimizer engine (assembly-time only).
216    ///
217    /// Re-runs the post-construction injection so the replacement behaves
218    /// like an engine passed to [`Self::with_optimizer`]. Callers must swap
219    /// before the first query; execution-time replacement is unsupported.
220    pub fn set_optimizer_engine(&mut self, optimizer_engine: Arc<OptimizerEngine>) {
221        optimizer_engine.set_cte_cache_stats_manager(self.stats_manager.clone());
222        self.optimizer_engine = optimizer_engine;
223    }
224
225    /// Effective sampling window currently in force for statistics
226    /// collection (server config value or the fallback default).
227    pub fn statistics_sample_limit(&self) -> usize {
228        self.statistics_sample_limit
229    }
230
231    /// Effective epoch delta currently in force for automatic refresh.
232    pub fn statistics_min_epoch_delta(&self) -> u64 {
233        self.statistics_min_epoch_delta
234    }
235
236    /// One-line summary of the effective statistics settings, surfaced in
237    /// EXPLAIN diagnostics so operators can verify which sampling window
238    /// and refresh threshold produced the statistics behind a plan.
239    pub fn statistics_settings_summary(&self) -> String {
240        format!(
241            "sample_limit={}, min_epoch_delta={}",
242            self.statistics_sample_limit, self.statistics_min_epoch_delta
243        )
244    }
245
246    /// Install per-space cost model overrides on the engine, with display
247    /// labels for EXPLAIN observability (assembly-time only).
248    pub fn set_space_cost_configs(
249        &mut self,
250        configs: HashMap<String, crate::optimizer::CostModelConfig>,
251        labels: HashMap<String, String>,
252    ) {
253        self.optimizer_engine.set_space_cost_configs(configs);
254        self.space_cost_profile_labels = labels;
255    }
256
257    /// Display label of the global cost profile (assembly-time only).
258    pub fn set_default_cost_profile_label(&mut self, label: String) {
259        self.default_cost_profile_label = label;
260    }
261
262    /// Effective cost profile label for a space, if the assembly provided
263    /// one: the space override when present, else the global label.
264    /// Empty global label means unassembled: no label at all.
265    pub fn cost_profile_label_for(&self, space: Option<&str>) -> Option<String> {
266        if self.default_cost_profile_label.is_empty() {
267            return None;
268        }
269        let label = space
270            .and_then(|name| self.space_cost_profile_labels.get(name))
271            .unwrap_or(&self.default_cost_profile_label);
272        Some(label.clone())
273    }
274
275    /// Collect (or serve cached) statistics for a space into the optimizer's
276    /// `StatisticsManager`.
277    ///
278    /// `force` bypasses the version gate so an explicit ANALYZE always
279    /// refreshes the statistics. Collection is serialized internally.
280    pub fn collect_statistics(
281        &self,
282        space: &str,
283        force: bool,
284    ) -> Result<crate::optimizer::stats::CollectedSummary, String> {
285        let _guard = self.statistics_collect_lock.lock();
286        let stats_manager = self.optimizer_engine.stats_manager().clone();
287        let schema_version = self
288            .schema_generation
289            .load(std::sync::atomic::Ordering::Relaxed);
290        if force {
291            stats_manager.mark_space_dirty(space);
292            self.optimizer_engine.invalidate_space_feedback(Some(space));
293        }
294        let storage: Arc<RwLock<dyn QueryStorage>> = self
295            .storage
296            .as_ref()
297            .ok_or_else(|| "No storage binding available for statistics collection".to_string())?
298            .clone();
299        let data_epoch = storage.read().stats_epoch();
300        let result = crate::optimizer::stats::StatisticsCollector::collect_space(
301            &stats_manager,
302            &storage,
303            space,
304            schema_version,
305            data_epoch,
306            self.statistics_sample_limit,
307        );
308        match &result {
309            Ok(summary) => {
310                log::info!(
311                    "Statistics collection for space '{}': {} tags, {} edge types{}",
312                    space,
313                    summary.tags,
314                    summary.edge_types,
315                    if summary.cached { " (cached)" } else { "" }
316                );
317            }
318            Err(error) => {
319                log::warn!(
320                    "Statistics collection failed for space '{}': {}",
321                    space,
322                    error
323                );
324            }
325        }
326        result
327    }
328
329    /// Ensure statistics for `space` are fresh before planning.
330    ///
331    /// Missing stamps trigger one best-effort collection; small data-epoch
332    /// advances below `statistics_min_epoch_delta` reuse the last collection
333    /// so write-heavy workloads do not re-sample on every commit. Failures
334    /// degrade to no-statistics fallbacks and never block the query. Explicit
335    /// ANALYZE remains the manual forced-refresh entry via
336    /// `collect_statistics(space, true)`.
337    pub(crate) fn ensure_statistics(&self, space: &str) {
338        let stats_manager = self.optimizer_engine.stats_manager();
339        let schema_version = self
340            .schema_generation
341            .load(std::sync::atomic::Ordering::Relaxed);
342        let data_epoch = match self.storage.as_ref() {
343            Some(storage) => storage.read().stats_epoch(),
344            None => return,
345        };
346        if schema_version != 0 && data_epoch != 0 {
347            match stats_manager.space_stamp(space) {
348                None => {}
349                Some((stamped_schema, stamped_epoch))
350                    if stamped_schema == schema_version && stamped_epoch == data_epoch =>
351                {
352                    return;
353                }
354                Some((stamped_schema, stamped_epoch))
355                    if stamped_schema == schema_version
356                        && data_epoch > stamped_epoch
357                        && data_epoch - stamped_epoch < self.statistics_min_epoch_delta =>
358                {
359                    return;
360                }
361                _ => {}
362            }
363        }
364        if let Err(error) = self.collect_statistics(space, false) {
365            log::warn!(
366                "Statistics refresh for space '{}' failed, degrading to no-stats: {}",
367                space,
368                error
369            );
370        }
371    }
372
373    pub fn with_query_registry(mut self, registry: Arc<QueryRegistry>) -> Self {
374        self.query_registry = Some(registry);
375        self
376    }
377
378    pub fn with_shared_scheduler(mut self, scheduler: Arc<SharedScheduler>) -> Self {
379        self.shared_scheduler = Some(scheduler);
380        self
381    }
382
383    /// Set the shared scheduler (used when the pipeline is built before the
384    /// scheduler is created, e.g. the vector-enabled server path).
385    pub fn set_shared_scheduler(&mut self, scheduler: Option<Arc<SharedScheduler>>) {
386        self.shared_scheduler = scheduler;
387    }
388
389    /// Set the query registry (used when the pipeline is built before the
390    /// registry is created, e.g. the vector-enabled server path).
391    pub fn set_query_registry(&mut self, registry: Option<Arc<QueryRegistry>>) {
392        self.query_registry = registry;
393    }
394
395    /// Access the shared scheduler instance, if any.
396    pub fn shared_scheduler(&self) -> Option<Arc<SharedScheduler>> {
397        self.shared_scheduler.clone()
398    }
399
400    /// Access the query registry instance, if any.
401    pub fn query_registry(&self) -> Option<Arc<QueryRegistry>> {
402        self.query_registry.clone()
403    }
404
405    /// Set the global QueryManager (used when the pipeline is built before the
406    /// manager is created, e.g. the server startup path).
407    pub fn set_query_manager(&mut self, manager: Option<Arc<QueryManager>>) {
408        self.query_manager = manager;
409    }
410
411    /// Access the global QueryManager instance, if installed.
412    pub fn query_manager(&self) -> Option<Arc<QueryManager>> {
413        self.query_manager.clone()
414    }
415
416    pub fn with_optimizer_and_cache(
417        storage: Arc<RwLock<S>>,
418        stats_manager: Arc<StatsManager>,
419        optimizer_engine: Arc<OptimizerEngine>,
420        plan_cache_config: PlanCacheConfig,
421    ) -> Self {
422        let plan_cache = Arc::new(
423            QueryPlanCache::new(plan_cache_config).with_stats_manager(stats_manager.clone()),
424        );
425        let param_handler = ParameterizedQueryHandler;
426
427        optimizer_engine.set_cte_cache_stats_manager(stats_manager.clone());
428
429        log::info!(
430            "Query pipeline manager created, using optimizer engine and custom query plan cache"
431        );
432
433        Self {
434            stats_manager,
435            optimizer_engine,
436            plan_cache,
437            param_handler,
438            schema_manager: None,
439            index_manager: None,
440            #[cfg(feature = "fulltext")]
441            fulltext_manager: None,
442            #[cfg(feature = "vector")]
443            vector_coordinator: None,
444            storage: Some(storage),
445            sync_manager: None,
446            schema_generation: Arc::new(AtomicU64::new(1)),
447            index_generation: Arc::new(AtomicU64::new(1)),
448            query_registry: None,
449            shared_scheduler: None,
450            query_manager: None,
451            macro_manager: Arc::new(linkrs_core::metadata::MacroManager::new()),
452            type_alias_manager: Arc::new(linkrs_core::metadata::TypeAliasManager::new()),
453            sequence_manager: Arc::new(linkrs_core::metadata::SequenceManager::new()),
454            session_controller: parking_lot::RwLock::new(None),
455            statistics_collect_lock: Arc::new(parking_lot::Mutex::new(())),
456            statistics_sample_limit: Self::DEFAULT_STATISTICS_SAMPLE_LIMIT,
457            statistics_min_epoch_delta: Self::DEFAULT_STATISTICS_MIN_EPOCH_DELTA,
458            default_cost_profile_label: String::new(),
459            space_cost_profile_labels: HashMap::new(),
460            dml_shape_cache_enabled: true,
461            dml_plan_memo: parking_lot::Mutex::new(Vec::new()),
462            last_dml_plan_hits: std::sync::atomic::AtomicU64::new(0),
463            dml_template_ast: parking_lot::Mutex::new(Vec::new()),
464            dml_template_ast_parse_count: std::sync::atomic::AtomicU64::new(0),
465            dml_template_ast_hit_count: std::sync::atomic::AtomicU64::new(0),
466            dml_bind_skipped_count: std::sync::atomic::AtomicU64::new(0),
467            extensions: Arc::new(ExtensionRegistry::new()),
468            password_history_depth: 0,
469        }
470    }
471
472    /// Attach a shared pipeline-extension registry (experimental).
473    pub fn with_extensions(mut self, extensions: Arc<ExtensionRegistry>) -> Self {
474        self.extensions = extensions;
475        self
476    }
477
478    /// Replace the pipeline-extension registry on an existing manager
479    /// (experimental). Used when the manager was built by another
480    /// constructor and the caller only has `&mut` access.
481    pub fn set_extensions(&mut self, extensions: Arc<ExtensionRegistry>) {
482        self.extensions = extensions;
483    }
484
485    /// Pipeline-extension registry backing this manager.
486    pub fn extensions(&self) -> Arc<ExtensionRegistry> {
487        Arc::clone(&self.extensions)
488    }
489
490    pub fn optimizer_engine(&self) -> &OptimizerEngine {
491        &self.optimizer_engine
492    }
493
494    pub fn with_schema_manager(mut self, schema_manager: Arc<SchemaManager>) -> Self {
495        self.schema_manager = Some(schema_manager);
496        self
497    }
498
499    pub fn with_index_manager(mut self, index_manager: Arc<dyn IndexMetadataManager>) -> Self {
500        self.index_manager = Some(index_manager);
501        self
502    }
503
504    #[cfg(feature = "fulltext")]
505    pub fn with_fulltext_manager(mut self, fulltext_manager: Arc<FulltextIndexManager>) -> Self {
506        self.fulltext_manager = Some(fulltext_manager);
507        self
508    }
509
510    #[cfg(feature = "vector")]
511    pub fn with_vector_coordinator(
512        mut self,
513        vector_coordinator: Arc<VectorSyncCoordinator>,
514    ) -> Self {
515        self.vector_coordinator = Some(vector_coordinator);
516        self
517    }
518
519    pub fn plan_cache(&self) -> &QueryPlanCache {
520        &self.plan_cache
521    }
522
523    pub fn plan_cache_metrics(&self) -> Arc<crate::cache::PlanCacheStats> {
524        self.plan_cache.stats()
525    }
526
527    pub fn clear_plan_cache(&self) {
528        self.plan_cache.clear();
529        // The DML shape caches feed plan construction, so a full clear must
530        // cover them too; otherwise `clear_plan_cache` would leave the bind
531        // fast path serving plans the plan cache itself just dropped.
532        self.dml_plan_memo.lock().clear();
533        self.dml_template_ast.lock().clear();
534        log::info!("Query plan cache cleared");
535    }
536
537    pub fn with_sync_manager(mut self, sync_manager: Arc<SyncManager>) -> Self {
538        self.sync_manager = Some(sync_manager);
539        self
540    }
541
542    /// Toggle DML shape plan caching (independent switch / rollback path).
543    pub fn with_dml_shape_cache(mut self, enabled: bool) -> Self {
544        self.dml_shape_cache_enabled = enabled;
545        self
546    }
547
548    /// Whether DML shape plan caching is enabled.
549    pub fn dml_shape_cache_enabled(&self) -> bool {
550        self.dml_shape_cache_enabled
551    }
552
553    pub fn dml_template_ast_parse_count(&self) -> u64 {
554        self.dml_template_ast_parse_count
555            .load(std::sync::atomic::Ordering::Relaxed)
556    }
557
558    pub fn dml_template_ast_hit_count(&self) -> u64 {
559        self.dml_template_ast_hit_count
560            .load(std::sync::atomic::Ordering::Relaxed)
561    }
562
563    pub fn dml_bind_skipped_count(&self) -> u64 {
564        self.dml_bind_skipped_count
565            .load(std::sync::atomic::Ordering::Relaxed)
566    }
567
568    /// The underlying storage binding, if any.
569    pub fn storage(&self) -> Option<Arc<RwLock<S>>> {
570        self.storage.clone()
571    }
572
573    /// Planning configuration snapshot used for DML memo keys. Must stay
574    /// identical to the config built in `compile_or_get_cached`; both sides
575    /// of the memo (prepare-time lookup, compile-time store) derive their
576    /// `optimizer_version` / `planning_config_hash` from here so the keys
577    /// stay comparable.
578    pub(crate) fn dml_planning_config(
579        &self,
580    ) -> crate::executor::streaming::plan::context::PlanningConfig {
581        crate::executor::streaming::plan::context::PlanningConfig {
582            max_partitions: self
583                .optimizer_engine
584                .partitioning_config()
585                .max_workers
586                .max(1),
587            config_hash: Self::partitioning_config_hash(
588                self.optimizer_engine.partitioning_config(),
589            ),
590            // Any cost configuration change (global or per-space) retires
591            // cached plans priced under the old configuration.
592            optimizer_version: self.optimizer_engine.cost_epoch(),
593            ..Default::default()
594        }
595    }
596
597    /// Look up a parsed normalized-DML template AST. Hits are moved to the
598    /// front (cheap LRU approximation for the linear-scan map).
599    pub(crate) fn lookup_dml_template(
600        &self,
601        normalized_text: &str,
602    ) -> Option<Arc<crate::parser::ast::stmt::Ast>> {
603        let mut guard = self.dml_template_ast.lock();
604        let pos = guard.iter().position(|(text, _)| text == normalized_text)?;
605        if pos > 0 {
606            let entry = guard.remove(pos);
607            guard.insert(0, entry);
608        }
609        self.dml_template_ast_hit_count
610            .fetch_add(1, Ordering::Relaxed);
611        Some(guard[0].1.clone())
612    }
613
614    /// Store a parsed normalized-DML template AST, evicting the least
615    /// recently used entry beyond [`DML_SHAPE_CACHE_CAP`].
616    pub(crate) fn store_dml_template(
617        &self,
618        normalized_text: String,
619        ast: Arc<crate::parser::ast::stmt::Ast>,
620    ) {
621        let mut guard = self.dml_template_ast.lock();
622        guard.retain(|(text, _)| text != &normalized_text);
623        guard.insert(0, (normalized_text, ast));
624        guard.truncate(DML_SHAPE_CACHE_CAP);
625    }
626
627    /// Look up a memoized DML physical plan by full key. Serves only as a
628    /// candidate: callers must still validate the plan with
629    /// `PhysicalPlanValidator::check_compatibility` before serving it, since
630    /// the memo key cannot capture post-planning layout changes.
631    pub(crate) fn find_dml_plan(&self, key: &DmlPlanMemoKey) -> Option<Arc<PhysicalPlan>> {
632        let mut guard = self.dml_plan_memo.lock();
633        let pos = guard.iter().position(|entry| entry.key == *key)?;
634        if pos > 0 {
635            let entry = guard.remove(pos);
636            guard.insert(0, entry);
637        }
638        Some(guard[0].plan.clone())
639    }
640
641    /// Store a memoized DML physical plan, replacing any same-key entry and
642    /// evicting the least recently used entry beyond [`DML_SHAPE_CACHE_CAP`].
643    /// Stale keys (e.g. from before a schema-generation bump) never match
644    /// new lookups and age out through this cap.
645    pub(crate) fn store_dml_plan(&self, key: DmlPlanMemoKey, plan: Arc<PhysicalPlan>) {
646        let mut guard = self.dml_plan_memo.lock();
647        guard.retain(|entry| entry.key != key);
648        guard.insert(0, DmlPlanMemoEntry { key, plan });
649        guard.truncate(DML_SHAPE_CACHE_CAP);
650    }
651}
linkrs_query/pipeline.rs

linkrs_query/
pipeline.rs
