linkrs_query/executor/streaming/subquery/runner.rs
1//! Compiled subquery runners: specs, runner state and the reset protocol.
2
3use std::collections::{HashMap, HashSet};
4use std::sync::Arc;
5
6use crate::executor::expression::ExpressionError;
7use crate::executor::streaming::executor::StreamingExecutor;
8use crate::executor::streaming::instance::QueryBindings;
9use crate::executor::streaming::plan::materializer::PhysicalPlanMaterializer;
10use crate::executor::streaming::plan::types::PhysicalPlan;
11use crate::executor::streaming::runtime::ExecutionRuntime;
12use crate::executor::streaming::slot::SlotLayout;
13use linkrs_core::error::QueryError;
14use linkrs_core::types::operators::AggregateFunction;
15use linkrs_core::{Expression, Value};
16use parking_lot::Mutex;
17
18/// Immutable subquery config carried by the arena spec.
19///
20/// Safe to share across parallel partitions: it holds only the compiled
21/// sub-plan and the planning-time routing flag. The materializer derives a
22/// fresh, per-operator [`SubqueryRunner`] from each spec so partitions never
23/// share mutable runner state.
24#[derive(Debug, Clone)]
25pub struct SubqueryRunnerSpec {
26 /// Stable identity assigned at planning time (`SubqueryBody.id`).
27 pub id: u64,
28 /// Standalone sub-plan (correlated = `Argument`-rooted subtree; the
29 /// Group-Join form = `Aggregate -> Project` build side).
30 pub plan: Arc<PhysicalPlan>,
31 /// Whether the subquery references outer columns.
32 pub correlated: bool,
33 /// Group-Join decorrelation info; when set the subquery is answered
34 /// from a materialized key→aggregate table instead of re-execution.
35 pub group_join: Option<GroupJoinSpec>,
36}
37
38/// Group-Join decorrelation info for a scalar aggregate subquery.
39#[derive(Debug, Clone)]
40pub struct GroupJoinSpec {
41 /// Outer-side key expressions, evaluated against the hosting row's
42 /// layout.
43 pub hash_keys: Vec<Expression>,
44 /// Leading group-key column count of each materialized sub-plan row;
45 /// the aggregated value follows at that index.
46 pub key_columns: usize,
47 /// The single aggregate function computed by the sub-plan.
48 pub function: AggregateFunction,
49 /// DISTINCT flag of the aggregate.
50 pub distinct: bool,
51}
52
53/// Cached result of a non-correlated subquery, evaluated exactly once.
54#[derive(Debug)]
55pub(crate) enum SubqueryCache {
56 /// EXISTS result: whether the subquery produced any row.
57 Exists(bool),
58 /// IN result set (NULL values removed at collection time).
59 Contains(HashSet<Value>),
60 /// Scalar result: first result value, or NULL for empty results.
61 Scalar(Value),
62}
63
64/// A compiled subquery with a reusable, resettable executor instance.
65///
66/// The executor is materialized once on first use and re-run per row via
67/// `reset()`. Interior mutability via `parking_lot::Mutex` keeps the runner
68/// `Sync` (executor trees are shipped to parallel partition workers); the
69/// lock is never contended because each runner is owned by exactly one
70/// operator instance and executed serially inside its `next()` call stack.
71#[derive(Debug)]
72pub struct SubqueryRunner {
73 /// Stable identity assigned at planning time (`SubqueryBody.id`).
74 pub id: u64,
75 /// Standalone sub-plan (correlated = `Argument`-rooted subtree).
76 pub plan: Arc<PhysicalPlan>,
77 /// Whether the subquery references outer columns.
78 pub correlated: bool,
79 /// Group-Join decorrelation info (see [`GroupJoinSpec`]).
80 pub group_join: Option<GroupJoinSpec>,
81 /// Materialized once on first use; re-run per row via `reset()`.
82 pub(crate) executor: Mutex<Option<StreamingExecutor>>,
83 /// Non-correlated results, evaluated once.
84 pub(crate) cache: Mutex<Option<SubqueryCache>>,
85 /// Group-Join lookup table: group key → aggregated value, built once
86 /// from a single materialization of the sub-plan.
87 pub(crate) group_table: Mutex<Option<HashMap<Vec<Value>, Value>>>,
88 /// Whether the last reset used the close+open fallback (EXPLAIN
89 /// `reset:fallback` audit).
90 pub reset_fallback: Mutex<bool>,
91}
92
93impl SubqueryRunner {
94 pub(super) fn from_spec(spec: &SubqueryRunnerSpec) -> Self {
95 Self {
96 id: spec.id,
97 plan: spec.plan.clone(),
98 correlated: spec.correlated,
99 group_join: spec.group_join.clone(),
100 executor: Mutex::new(None),
101 cache: Mutex::new(None),
102 group_table: Mutex::new(None),
103 reset_fallback: Mutex::new(false),
104 }
105 }
106
107 /// Materialize the sub-plan once (lazily) and run `f` against the
108 /// reused executor. For correlated subqueries the current row is injected
109 /// as the correlation frame before each run.
110 pub(super) fn with_executor<F, T>(
111 &self,
112 runtime: &Arc<ExecutionRuntime>,
113 bindings: &Arc<QueryBindings>,
114 layout: Arc<SlotLayout>,
115 row: Vec<Value>,
116 f: F,
117 ) -> Result<T, ExpressionError>
118 where
119 F: FnOnce(&mut StreamingExecutor) -> Result<T, QueryError>,
120 {
121 let mut slot = self.executor.lock();
122 if slot.is_none() {
123 let (mut exec, _) = PhysicalPlanMaterializer::materialize(&self.plan, bindings)
124 .map_err(|e| {
125 ExpressionError::type_error(format!(
126 "Subquery plan materialization failed: {}",
127 e
128 ))
129 })?;
130 exec.set_chunk_size(bindings.chunk_size);
131 exec.set_runtime(Some(runtime.clone()));
132 exec.open()
133 .map_err(|e| ExpressionError::type_error(format!("Subquery open failed: {}", e)))?;
134 *slot = Some(exec);
135 }
136 let exec = slot.as_mut().ok_or_else(|| {
137 ExpressionError::type_error("Subquery executor failed to materialize".to_string())
138 })?;
139 if self.correlated {
140 exec.inject_correlation_frame(layout, row);
141 }
142 exec.reset()
143 .map_err(|e| ExpressionError::type_error(format!("Subquery reset failed: {}", e)))?;
144 if exec.base().reset_used_fallback {
145 *self.reset_fallback.lock() = true;
146 }
147 f(exec)
148 .map_err(|e| ExpressionError::type_error(format!("Subquery execution failed: {}", e)))
149 }
150}