linkrs_query/executor/streaming/
executor.rs

1//! StreamingExecutor: Thin dispatch layer over domain-specific operator enums
2
3use std::sync::Arc;
4use std::time::Instant;
5
6use super::chunk::DataChunk;
7use super::runtime::{ExecutionRuntime, OperatorProfile, OperatorProfileKey};
8use super::slot::SlotLayout;
9use crate::executor::base::{MemoryTracker, Spillable};
10use linkrs_core::error::QueryError;
11use linkrs_core::Value;
12
13pub use super::context::ValueRowContext;
14pub use super::helpers::{comparison, conversion};
15pub use super::operators::base::OperatorBase;
16use super::operators::base::OperatorLifecycle;
17use super::operators::source_operator::OperatorConfig;
18
19use super::operators::apply_operator::ApplyOperator;
20use super::operators::blocking::BlockingOperator;
21use super::operators::blocking::BlockingOperatorKind;
22use super::operators::ddl_operator::DdlOperator;
23use super::operators::exchange_operator::ExchangeOperator;
24use super::operators::fulltext_operator::FulltextOperator;
25use super::operators::gather_operator::GatherOperator;
26use super::operators::graph_operator::GraphOperator;
27use super::operators::graph_operator::GraphOperatorKind;
28use super::operators::join_operator::JoinOperator;
29use super::operators::recursive_fragment_operator::RecursiveFragmentOperator;
30use super::operators::set_operator::SetOperator;
31use super::operators::set_operator::SetOperatorKind;
32use super::operators::shuffle_join_operator::HashShuffleJoinOperator;
33use super::operators::sink_operator::SinkOperator;
34use super::operators::source_operator::SourceOperator;
35use super::operators::source_operator::SourceOperatorKind;
36use super::operators::txn_operator::TxnOperator;
37use super::operators::unary_operator::UnaryOperator;
38use super::operators::unary_operator::UnaryOperatorKind;
39use super::operators::vector_operator::VectorOperator;
40use super::operators::wco_operator::WcoIntersectOperator;
41
42mod operator_name;
43
44/// Sort direction for ORDER BY clause
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum SortDirection {
47    Ascending,
48    Descending,
49}
50
51/// Phase for FullOuterJoin execution
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub enum FullOuterJoinPhase {
54    BuildingRight,
55    ProbeLeft,
56    EmitUnmatchedRight,
57}
58
59/// StreamingExecutor: 17-variant dispatch enum over domain-specific operators.
60///
61/// Each variant holds an OperatorBase (shared lifecycle/base fields), zero
62/// or more child executors, and a domain-specific operator wrapper that
63/// implements the per-operator lifecycle logic with zero-context methods.
64///
65/// Gather/Exchange take N children (Vec) and merge their output via
66/// Concatenate or MergeSort mode. Exchange additionally uses the query-level
67/// `MorselWorkerPool` for morsel-style dynamic partition execution.
68#[derive(Debug)]
69pub enum StreamingExecutor {
70    Source(OperatorBase, SourceOperator),
71    Unary(OperatorBase, Box<StreamingExecutor>, UnaryOperator),
72    Join(
73        OperatorBase,
74        Box<StreamingExecutor>,
75        Box<StreamingExecutor>,
76        Box<JoinOperator>,
77    ),
78    Set(
79        OperatorBase,
80        Box<StreamingExecutor>,
81        Box<StreamingExecutor>,
82        SetOperator,
83    ),
84    Apply(
85        OperatorBase,
86        Box<StreamingExecutor>,
87        Box<StreamingExecutor>,
88        ApplyOperator,
89    ),
90    Blocking(OperatorBase, Box<StreamingExecutor>, BlockingOperator),
91    Graph(OperatorBase, Box<StreamingExecutor>, GraphOperator),
92    /// Recursive fragment for variable-length path traversal,
93    /// BFS, shortest-path, and multi-round graph algorithms.
94    RecursiveFragment(
95        OperatorBase,
96        Box<StreamingExecutor>,
97        RecursiveFragmentOperator,
98    ),
99    Sink(OperatorBase, Box<StreamingExecutor>, SinkOperator),
100    Ddl(OperatorBase, Box<StreamingExecutor>, DdlOperator),
101    Fulltext(OperatorBase, Box<StreamingExecutor>, FulltextOperator),
102    Vector(OperatorBase, Box<StreamingExecutor>, VectorOperator),
103    Txn(OperatorBase, Box<StreamingExecutor>, TxnOperator),
104    Gather(OperatorBase, Vec<StreamingExecutor>, GatherOperator),
105    /// Morsel-style exchange: gathers N partition outputs via Concatenate
106    /// or MergeSort using the query-level `MorselWorkerPool`.
107    Exchange(OperatorBase, Vec<StreamingExecutor>, ExchangeOperator),
108    /// Hash repartition join: takes multiple left/right local trees,
109    /// distributes rows by hash of join keys into buckets, joins per bucket.
110    HashShuffleJoin(
111        OperatorBase,
112        Vec<StreamingExecutor>,
113        Vec<StreamingExecutor>,
114        HashShuffleJoinOperator,
115    ),
116    /// N-way worst-case-optimal intersect: one probe input plus N build
117    /// inputs sharing a single intersect variable.
118    Wco(
119        OperatorBase,
120        Box<StreamingExecutor>,
121        Vec<StreamingExecutor>,
122        WcoIntersectOperator,
123    ),
124}
125
126/// Dispatch `open` (all operators open their children themselves).
127macro_rules! dispatch_open {
128    ($self:expr) => {
129        match $self {
130            Self::Source(_, op) => op.open(),
131            Self::Unary(_, input, op) => op.open(input),
132            Self::Join(_, left, right, op) => op.open(left, right),
133            Self::Set(_, left, right, op) => op.open(left, right),
134            Self::Apply(_, left, right, op) => op.open(left, right),
135            Self::Blocking(_, input, op) => op.open(input),
136            Self::Graph(_, input, op) => op.open(input),
137            Self::RecursiveFragment(_, input, op) => op.open(input),
138            Self::Sink(_, input, op) => op.open(input),
139            Self::Ddl(_, input, op) => op.open(input),
140            Self::Fulltext(_, input, op) => op.open(input),
141            Self::Vector(_, input, op) => op.open(input),
142            Self::Txn(_, input, op) => op.open(input),
143            Self::Gather(_, children, op) => op.open(children),
144            Self::Exchange(_, children, op) => op.open(children),
145            Self::HashShuffleJoin(_, left, right, op) => op.open(left, right),
146            Self::Wco(_, probe, builds, op) => op.open(probe, builds),
147        }
148    };
149}
150
151/// Dispatch `next` (same signature family as `open`).
152macro_rules! dispatch_next {
153    ($self:expr) => {
154        match $self {
155            Self::Source(_, op) => op.next(),
156            Self::Unary(_, input, op) => op.next(input),
157            Self::Join(_, left, right, op) => op.next(left, right),
158            Self::Set(_, left, right, op) => op.next(left, right),
159            Self::Apply(_, left, right, op) => op.next(left, right),
160            Self::Blocking(_, input, op) => op.next(input),
161            Self::Graph(_, input, op) => op.next(input),
162            Self::RecursiveFragment(_, input, op) => op.next(input),
163            Self::Sink(_, input, op) => op.next(input),
164            Self::Ddl(_, input, op) => op.next(input),
165            Self::Fulltext(_, input, op) => op.next(input),
166            Self::Vector(_, input, op) => op.next(input),
167            Self::Txn(_, input, op) => op.next(input),
168            Self::Gather(_, children, op) => op.next(children),
169            Self::Exchange(_, children, op) => op.next(children),
170            Self::HashShuffleJoin(_, left, right, op) => op.next(left, right),
171            Self::Wco(_, probe, builds, op) => op.next(probe, builds),
172        }
173    };
174}
175
176/// Dispatch `reset`. Blocking operators fully re-create their state in
177/// `open()` and drop it in `close()`, so the close+open fallback is exact.
178/// The remaining operators (Sink/Gather/Exchange/…/Txn) do not appear inside
179/// resettable sub-plans; the fallback is a transitional audit point
180/// (EXPLAIN `reset:fallback`).
181macro_rules! dispatch_reset {
182    ($self:expr) => {
183        match $self {
184            Self::Source(_, op) => op.reset(),
185            Self::Unary(_, input, op) => op.reset(input),
186            Self::Join(_, left, right, op) => op.reset(left, right),
187            Self::Set(_, left, right, op) => op.reset(left, right),
188            Self::Apply(_, left, right, op) => op.reset(left, right),
189            Self::Graph(_, input, op) => op.reset(input),
190            Self::RecursiveFragment(_, input, op) => op.reset(input),
191            Self::Wco(_, probe, builds, op) => op.reset(probe, builds),
192            _ => $self.fallback_reset(),
193        }
194    };
195}
196
197/// Dispatch `stop`.
198macro_rules! dispatch_stop {
199    ($self:expr) => {
200        match $self {
201            Self::Source(_, op) => op.stop(),
202            Self::Unary(_, _, op) => op.stop(),
203            Self::Join(_, _, _, op) => op.stop(),
204            Self::Set(_, _, _, op) => op.stop(),
205            Self::Apply(_, _, _, op) => op.stop(),
206            Self::Blocking(_, _, op) => op.stop(),
207            Self::Graph(_, _, op) => op.stop(),
208            Self::RecursiveFragment(_, _, op) => op.stop(),
209            Self::Sink(_, _, op) => op.stop(),
210            Self::Ddl(_, _, op) => op.stop(),
211            Self::Fulltext(_, _, op) => op.stop(),
212            Self::Vector(_, _, op) => op.stop(),
213            Self::Txn(_, _, op) => op.stop(),
214            Self::Gather(_, _, op) => op.stop(),
215            Self::Exchange(_, _, op) => op.stop(),
216            Self::HashShuffleJoin(_, _, _, op) => op.stop(),
217            Self::Wco(_, _, _, op) => op.stop(),
218        }
219    };
220}
221
222/// Dispatch `close`.
223macro_rules! dispatch_close {
224    ($self:expr) => {
225        match $self {
226            Self::Source(_, op) => op.close(),
227            Self::Unary(_, _, op) => op.close(),
228            Self::Join(_, _, _, op) => op.close(),
229            Self::Set(_, _, _, op) => op.close(),
230            Self::Apply(_, _, _, op) => op.close(),
231            Self::Blocking(_, _, op) => op.close(),
232            Self::Graph(_, _, op) => op.close(),
233            Self::RecursiveFragment(_, _, op) => op.close(),
234            Self::Sink(_, _, op) => op.close(),
235            Self::Ddl(_, _, op) => op.close(),
236            Self::Fulltext(_, _, op) => op.close(),
237            Self::Vector(_, _, op) => op.close(),
238            Self::Txn(_, _, op) => op.close(),
239            Self::Gather(_, _, op) => op.close(),
240            Self::Exchange(_, _, op) => op.close(),
241            Self::HashShuffleJoin(_, _, _, op) => op.close(),
242            Self::Wco(_, _, _, op) => op.close(),
243        }
244    };
245}
246
247impl StreamingExecutor {
248    /// Recursively inject the runtime and execution config into every
249    /// operator wrapper. `open()` derives the config from the base fields,
250    /// so callers only need `set_runtime`/`set_chunk_size`/`set_partition_id`
251    /// beforehand. Idempotent; safe to call on already-opened trees.
252    pub fn inject_context(
253        &mut self,
254        runtime: Option<Arc<ExecutionRuntime>>,
255        config: OperatorConfig,
256    ) {
257        let runtime_ref = runtime.as_ref();
258        match self {
259            Self::Source(_, op) => op.inject_context(runtime_ref, config),
260            Self::Unary(_, _, op) => op.inject_context(runtime_ref, config),
261            Self::Join(_, _, _, op) => op.inject_context(runtime_ref, config),
262            Self::Set(_, _, _, op) => op.inject_context(runtime_ref, config),
263            Self::Apply(_, _, _, op) => op.inject_context(runtime_ref, config),
264            Self::Blocking(_, _, op) => op.inject_context(runtime_ref, config),
265            Self::Graph(_, _, op) => op.inject_context(runtime_ref, config),
266            Self::RecursiveFragment(_, _, op) => op.inject_context(runtime_ref, config),
267            Self::Sink(_, _, op) => op.inject_context(runtime_ref, config),
268            Self::Ddl(_, _, op) => op.inject_context(runtime_ref, config),
269            Self::Fulltext(_, _, op) => op.inject_context(runtime_ref, config),
270            Self::Vector(_, _, op) => op.inject_context(runtime_ref, config),
271            Self::Txn(_, _, op) => op.inject_context(runtime_ref, config),
272            Self::Gather(_, _, op) => op.inject_context(runtime_ref, config),
273            Self::Exchange(_, _, op) => op.inject_context(runtime_ref, config),
274            Self::HashShuffleJoin(_, _, _, op) => op.inject_context(runtime_ref, config),
275            Self::Wco(_, _, _, op) => op.inject_context(runtime_ref, config),
276        }
277        for child in self.children_mut() {
278            child.inject_context(runtime.clone(), config);
279        }
280    }
281
282    /// Recursively set the runtime on this operator and all children.
283    pub fn set_runtime(&mut self, rt: Option<Arc<ExecutionRuntime>>) {
284        if let (Self::Graph(_, _, operator), Some(runtime)) = (&mut *self, rt.as_ref()) {
285            operator.bind_runtime(runtime);
286        }
287        if let (Self::RecursiveFragment(_, _, operator), Some(runtime)) = (&mut *self, rt.as_ref())
288        {
289            operator.bind_runtime(runtime);
290        }
291        self.base_mut().runtime = rt.clone();
292        for child in self.children_mut() {
293            child.set_runtime(rt.clone());
294        }
295    }
296
297    /// Recursively set the chunk size on this operator and all children.
298    pub fn set_chunk_size(&mut self, chunk_size: usize) {
299        self.base_mut().chunk_size = chunk_size;
300        for child in self.children_mut() {
301            child.set_chunk_size(chunk_size);
302        }
303    }
304
305    /// Return the plan node ID of this operator.
306    pub fn plan_node_id(&self) -> i64 {
307        self.base().plan_node_id
308    }
309
310    /// Return the unique profile key for this concrete operator instance.
311    pub fn profile_key(&self) -> OperatorProfileKey {
312        self.base().profile_key()
313    }
314
315    /// Recursively mark this executor tree as belonging to a local partition.
316    pub fn set_partition_id(&mut self, partition_id: usize) {
317        self.base_mut().partition_id = Some(partition_id);
318        self.base_mut().is_global = false;
319        for child in self.children_mut() {
320            child.set_partition_id(partition_id);
321        }
322    }
323
324    /// Recursively mark an executor tree as global. This is applied before a
325    /// gather/exchange node is attached, so the local children retain their
326    /// own partition identifiers.
327    pub fn set_global(&mut self) {
328        self.base_mut().partition_id = None;
329        self.base_mut().is_global = true;
330        for child in self.children_mut() {
331            child.set_global();
332        }
333    }
334
335    /// Whether this tree may be executed independently for each partition and
336    /// concatenated without changing query semantics.
337    ///
338    /// Global operators are intentionally excluded. They require a dedicated
339    /// physical split (for example local sort plus merge sort), not a copied
340    /// copy of the original logical tree.
341    pub fn is_partition_local(&self) -> bool {
342        match self {
343            Self::Source(_, op) => matches!(
344                &op.kind,
345                SourceOperatorKind::ScanVertices { .. }
346                    | SourceOperatorKind::StorageScanVertices { .. }
347                    | SourceOperatorKind::ScanEdges { .. }
348                    | SourceOperatorKind::StorageScanEdges { .. }
349            ),
350            Self::Unary(_, input, op) => {
351                matches!(
352                    &op.kind,
353                    UnaryOperatorKind::Filter { .. }
354                        | UnaryOperatorKind::Project { .. }
355                        | UnaryOperatorKind::Assign { .. }
356                        | UnaryOperatorKind::Unwind { .. }
357                        | UnaryOperatorKind::AppendVertices { .. }
358                ) && input.is_partition_local()
359            }
360            Self::Blocking(_, input, op) => {
361                matches!(
362                    &op.kind,
363                    // Morsel parallel: Sort/Agg are partition-local
364                    // (local sort + merge, partial/final aggregate). Hash
365                    // aggregation uses `PartialAggregate` (partition-local) +
366                    // `FinalAggregate` (global gather) two-phase; `Sort` is
367                    // local + `Exchange::MergeSort`. `TopN` is bounded local
368                    // sort, `Distinct` and `GroupBy` are hash-partitioned.
369                    BlockingOperatorKind::PartialAggregate { .. }
370                        | BlockingOperatorKind::Sort { .. }
371                        | BlockingOperatorKind::Distinct { .. }
372                        | BlockingOperatorKind::TopN { .. }
373                        | BlockingOperatorKind::GroupBy { .. }
374                        | BlockingOperatorKind::Aggregate { .. }
375                ) && input.is_partition_local()
376            }
377            Self::Graph(_, input, op) => {
378                // Single-hop Expand/BiExpand are morsel-parallel
379                // (partitioned by anchor vertex-id range, batched neighbor
380                // scan via `neighbor_dst_ids_batch`).
381                matches!(
382                    &op.kind,
383                    GraphOperatorKind::Expand { .. }
384                        | GraphOperatorKind::BiExpand { .. }
385                        | GraphOperatorKind::ExpandAll { .. }
386                ) && input.is_partition_local()
387            }
388            Self::HashShuffleJoin(..) | Self::Exchange(..) | Self::RecursiveFragment(..) => false,
389            _ => false,
390        }
391    }
392
393    /// Access the runtime reference, if attached.
394    pub fn get_runtime(&self) -> Option<&ExecutionRuntime> {
395        self.base().runtime.as_deref()
396    }
397
398    /// Check cancellation via the attached runtime.
399    pub fn ensure_not_cancelled(&self) -> Result<(), QueryError> {
400        self.base().ensure_not_cancelled()
401    }
402
403    /// Profile name for this operator variant.
404    pub fn operator_name(&self) -> &'static str {
405        operator_name::operator_name(self)
406    }
407
408    /// Record profile timing for this operator, using the correct operator name.
409    ///
410    /// Fast path: looks up the pre-registered [`ProfileEntry`] via read-lock,
411    /// then updates the atomic counter without holding any lock.
412    pub fn record_profile_timing(&self, phase: &str, elapsed_us: u64) {
413        let Some(rt) = &self.base().runtime else {
414            return;
415        };
416        let entry = rt.profile().get_entry(&self.profile_key());
417        let entry = entry.unwrap_or_else(|| {
418            // First access — register on demand (rare: tests that bypass open())
419            let name = self.operator_name();
420            let op = OperatorProfile {
421                physical_operator_id: self.base().physical_operator_id,
422                node_id: self.plan_node_id(),
423                partition_id: self.base().partition_id,
424                name: name.to_string(),
425                ..OperatorProfile::default()
426            };
427            rt.register_operator(&op)
428        });
429        let ordering = std::sync::atomic::Ordering::Relaxed;
430        if phase == "open" {
431            entry.open_time_us.fetch_add(elapsed_us, ordering);
432        } else if phase == "next" {
433            entry.next_time_us.fetch_add(elapsed_us, ordering);
434        } else if phase == "close" {
435            entry.close_time_us.fetch_add(elapsed_us, ordering);
436        }
437    }
438
439    /// Get peak memory from the memory_tracker, if this operator has one.
440    pub fn peak_memory_bytes(&self) -> u64 {
441        self.memory_tracker().map_or(0, |mt| mt.peak() as u64)
442    }
443
444    /// Return the parallel fallback reason from any operator in this
445    /// tree that has a recorded reason, or `None`.
446    pub fn parallel_fallback_reason(&self) -> Option<String> {
447        match self {
448            Self::Gather(_, children, _) | Self::Exchange(_, children, _) => {
449                children.iter().find_map(|c| c.parallel_fallback_reason())
450            }
451            Self::Unary(_, child, _)
452            | Self::Blocking(_, child, _)
453            | Self::Graph(_, child, _)
454            | Self::RecursiveFragment(_, child, _)
455            | Self::Sink(_, child, _)
456            | Self::Ddl(_, child, _)
457            | Self::Fulltext(_, child, _)
458            | Self::Vector(_, child, _)
459            | Self::Txn(_, child, _) => child.parallel_fallback_reason(),
460            Self::Join(_, left, right, _) => left
461                .parallel_fallback_reason()
462                .or_else(|| right.parallel_fallback_reason()),
463            Self::Set(_, left, right, _) => left
464                .parallel_fallback_reason()
465                .or_else(|| right.parallel_fallback_reason()),
466            Self::Apply(_, outer, inner, _) => outer
467                .parallel_fallback_reason()
468                .or_else(|| inner.parallel_fallback_reason()),
469            Self::HashShuffleJoin(_, left, right, _) => left
470                .iter()
471                .find_map(|c| c.parallel_fallback_reason())
472                .or_else(|| right.iter().find_map(|c| c.parallel_fallback_reason())),
473            Self::Wco(_, probe, builds, _) => probe
474                .parallel_fallback_reason()
475                .or_else(|| builds.iter().find_map(|c| c.parallel_fallback_reason())),
476            Self::Source(..) => None,
477        }
478    }
479
480    /// Record output row count in profile for this operator.
481    pub fn record_profile_rows(&self, count: u64) {
482        self.base().record_profile_rows(count);
483    }
484
485    /// Record peak memory usage in profile for this operator.
486    pub fn record_profile_peak_memory(&self, bytes: u64) {
487        let Some(rt) = &self.base().runtime else {
488            return;
489        };
490        let Some(entry) = rt.profile().get_entry(&self.profile_key()) else {
491            return;
492        };
493        let prev = entry
494            .peak_memory_bytes
495            .fetch_max(bytes, std::sync::atomic::Ordering::Relaxed);
496        if bytes > prev {
497            entry
498                .peak_memory_bytes
499                .store(bytes, std::sync::atomic::Ordering::Relaxed);
500        }
501    }
502
503    /// Record spilled rows, bytes and spill count in profile for this operator.
504    pub fn record_profile_spill(&self, spilled_bytes: u64, spill_count: u64, spilled_rows: u64) {
505        let Some(rt) = &self.base().runtime else {
506            return;
507        };
508        let Some(entry) = rt.profile().get_entry(&self.profile_key()) else {
509            return;
510        };
511        entry
512            .spilled_bytes
513            .fetch_add(spilled_bytes, std::sync::atomic::Ordering::Relaxed);
514        entry
515            .spill_count
516            .fetch_add(spill_count, std::sync::atomic::Ordering::Relaxed);
517        entry
518            .spilled_rows
519            .fetch_add(spilled_rows, std::sync::atomic::Ordering::Relaxed);
520    }
521
522    /// Register a resource cleanup callback with the attached runtime.
523    pub fn register_resource<F>(&self, f: F)
524    where
525        F: FnOnce() + Send + 'static,
526    {
527        self.base().register_resource(f);
528    }
529
530    /// Return the base fields.
531    pub fn base(&self) -> &OperatorBase {
532        match self {
533            Self::Source(base, _) => base,
534            Self::Unary(base, _, _) => base,
535            Self::Join(base, _, _, _) | Self::Set(base, _, _, _) | Self::Apply(base, _, _, _) => {
536                base
537            }
538            Self::Blocking(base, _, _) => base,
539            Self::Graph(base, _, _) | Self::RecursiveFragment(base, _, _) => base,
540            Self::Sink(base, _, _) => base,
541            Self::Ddl(base, _, _) | Self::Fulltext(base, _, _) | Self::Vector(base, _, _) => base,
542            Self::Txn(base, _, _) => base,
543            Self::Gather(base, _, _) | Self::Exchange(base, _, _) => base,
544            Self::HashShuffleJoin(base, _, _, _) => base,
545            Self::Wco(base, _, _, _) => base,
546        }
547    }
548
549    /// Return the base fields (mutable).
550    pub fn base_mut(&mut self) -> &mut OperatorBase {
551        match self {
552            Self::Source(base, _) => base,
553            Self::Unary(base, _, _) => base,
554            Self::Join(base, _, _, _) | Self::Set(base, _, _, _) | Self::Apply(base, _, _, _) => {
555                base
556            }
557            Self::Blocking(base, _, _) => base,
558            Self::Graph(base, _, _) | Self::RecursiveFragment(base, _, _) => base,
559            Self::Sink(base, _, _) => base,
560            Self::Ddl(base, _, _) | Self::Fulltext(base, _, _) | Self::Vector(base, _, _) => base,
561            Self::Txn(base, _, _) => base,
562            Self::Gather(base, _, _) | Self::Exchange(base, _, _) => base,
563            Self::HashShuffleJoin(base, _, _, _) => base,
564            Self::Wco(base, _, _, _) => base,
565        }
566    }
567
568    /// Return mutable references to all child executors.
569    pub fn children_mut(&mut self) -> Vec<&mut Self> {
570        match self {
571            Self::Source(..) => vec![],
572            Self::Unary(_, input, _)
573            | Self::Blocking(_, input, _)
574            | Self::Graph(_, input, _)
575            | Self::RecursiveFragment(_, input, _)
576            | Self::Sink(_, input, _)
577            | Self::Txn(_, input, _) => vec![input.as_mut()],
578            Self::Join(_, left, right, _)
579            | Self::Set(_, left, right, _)
580            | Self::Apply(_, left, right, _) => vec![left.as_mut(), right.as_mut()],
581            Self::Ddl(_, input, _) | Self::Fulltext(_, input, _) | Self::Vector(_, input, _) => {
582                vec![input.as_mut()]
583            }
584            Self::Gather(_, children, _) | Self::Exchange(_, children, _) => {
585                children.iter_mut().collect()
586            }
587            Self::HashShuffleJoin(_, left, right, _) => {
588                let mut all: Vec<&mut Self> = left.iter_mut().collect();
589                all.extend(right.iter_mut());
590                all
591            }
592            Self::Wco(_, probe, builds, _) => {
593                let mut all = vec![probe.as_mut()];
594                all.extend(builds.iter_mut());
595                all
596            }
597        }
598    }
599
600    /// Access the MemoryTracker for blocking/binary operators.
601    pub fn memory_tracker(&self) -> Option<&MemoryTracker> {
602        match self {
603            Self::Source(..)
604            | Self::Unary(..)
605            | Self::Graph(..)
606            | Self::RecursiveFragment(..)
607            | Self::Sink(..)
608            | Self::Txn(..)
609            | Self::Ddl(..)
610            | Self::Fulltext(..)
611            | Self::Vector(..)
612            | Self::Gather(..)
613            | Self::Exchange(..) => None,
614            Self::Blocking(_, _, op) => Some(op.memory_tracker()),
615            Self::Join(_, _, _, op) => Some(op.memory_tracker()),
616            Self::Set(_, _, _, op) => {
617                if matches!(&op.kind, SetOperatorKind::UnionAll { .. }) {
618                    None
619                } else {
620                    Some(op.memory_tracker())
621                }
622            }
623            Self::Apply(_, _, _, op) => Some(op.memory_tracker()),
624            Self::HashShuffleJoin(_, _, _, op) => Some(op.memory_tracker()),
625            Self::Wco(_, _, _, op) => Some(op.memory_tracker()),
626        }
627    }
628
629    /// Whether this operator is in an opened state (can produce chunks).
630    pub fn opened(&self) -> bool {
631        self.base().lifecycle.is_opened()
632    }
633
634    // ── Reset protocol ──
635
636    /// Reset the executor so it re-produces the same logical stream as the
637    /// first run (for the same snapshot / correlation frame).
638    ///
639    /// Lifecycle: `open → (advance)* → reset → (advance)* → … → close`.
640    /// The executor instance is reused; no operator structs are rebuilt.
641    /// Operators without a native reset degrade to `close + open`
642    /// (`reset_used_fallback` is set, surfaced in EXPLAIN).
643    pub fn reset(&mut self) -> Result<(), QueryError> {
644        let used_fallback = dispatch_reset!(self)?;
645        self.base_mut().reset_used_fallback |= used_fallback;
646        self.restore_opened_lifecycle();
647        Ok(())
648    }
649
650    /// Transitional fallback: `close_tree + open` on the whole subtree.
651    /// Reuses the existing close/open semantics; never rebuilds operator
652    /// structs. Must not become the default path for commonly reset
653    /// operators — EXPLAIN marks it via `reset_used_fallback`.
654    fn fallback_reset(&mut self) -> Result<bool, QueryError> {
655        self.close_tree()?;
656        self.mark_tree_new();
657        self.open()?;
658        Ok(true)
659    }
660
661    /// Set every base lifecycle in this tree to `New` (before a re-open).
662    fn mark_tree_new(&mut self) {
663        self.base_mut().mark_new();
664        for child in self.children_mut() {
665            child.mark_tree_new();
666        }
667    }
668
669    /// Mark every base lifecycle in this tree as `Opened` (after a
670    /// successful open). Operators no longer touch lifecycle themselves;
671    /// the executor is the sole owner of the lifecycle state machine.
672    fn mark_tree_opened(&mut self) {
673        self.base_mut().lifecycle = OperatorLifecycle::Opened;
674        for child in self.children_mut() {
675            child.mark_tree_opened();
676        }
677    }
678
679    /// Restore `Opened` on operators that were opened/exhausted so they can
680    /// produce again after a reset.
681    fn restore_opened_lifecycle(&mut self) {
682        if matches!(
683            self.base().lifecycle,
684            OperatorLifecycle::Opened | OperatorLifecycle::Exhausted
685        ) {
686            self.base_mut().lifecycle = OperatorLifecycle::Opened;
687        }
688        for child in self.children_mut() {
689            child.restore_opened_lifecycle();
690        }
691    }
692
693    /// Inject a correlation frame into the `Argument` source of this
694    /// executor tree (the root of a correlated sub-plan). Each executor
695    /// instance owns its frame, so parallel partitions and nested
696    /// subqueries never interfere.
697    pub fn inject_correlation_frame(&mut self, layout: Arc<SlotLayout>, row: Vec<Value>) {
698        if let Self::Source(_, op) = self {
699            if matches!(&op.kind, SourceOperatorKind::Argument) {
700                op.frame = Some((layout, row));
701                return;
702            }
703        }
704        for child in self.children_mut() {
705            child.inject_correlation_frame(layout.clone(), row.clone());
706        }
707    }
708
709    // ── Lifecycle dispatch ──
710
711    /// Open the executor.
712    pub fn open(&mut self) -> Result<(), QueryError> {
713        self.ensure_not_cancelled()?;
714        // Pre-register profile entry so the hot path in advance()
715        // can find it without write-locking.
716        if let Some(rt) = &self.base().runtime {
717            let name = self.operator_name();
718            let op = OperatorProfile {
719                physical_operator_id: self.base().physical_operator_id,
720                node_id: self.plan_node_id(),
721                partition_id: self.base().partition_id,
722                name: name.to_string(),
723                ..OperatorProfile::default()
724            };
725            rt.register_operator(&op);
726        }
727        // Inject the runtime and execution config (derived from the base
728        // fields) into every operator wrapper before any operator code runs.
729        let runtime = self.base().runtime.clone();
730        let config = OperatorConfig {
731            chunk_size: self.base().chunk_size,
732            partition_id: self.base().partition_id,
733            physical_operator_id: self.base().physical_operator_id,
734        };
735        self.inject_context(runtime, config);
736        let start = Instant::now();
737        let result = dispatch_open!(self);
738        let elapsed = start.elapsed().as_micros() as u64;
739        self.record_profile_timing("open", elapsed);
740        if let Err(error) = result {
741            if let Err(close_error) = self.close_tree() {
742                log::warn!(
743                    "Failed to close streaming executor tree after open error: {}",
744                    close_error
745                );
746            }
747            return Err(error);
748        }
749        self.mark_tree_opened();
750        Ok(())
751    }
752
753    /// Pull the next chunk.
754    pub fn advance(&mut self) -> Result<Option<DataChunk>, QueryError> {
755        self.ensure_not_cancelled()?;
756        // Operators no longer guard on lifecycle themselves; the executor
757        // enforces that only opened, non-exhausted executors produce.
758        if self.base().lifecycle.is_exhausted() || !self.base().lifecycle.is_opened() {
759            return Ok(None);
760        }
761        let one_shot = matches!(self, Self::Ddl(..) | Self::Fulltext(..) | Self::Vector(..));
762        let start = Instant::now();
763        let result = dispatch_next!(self);
764        let elapsed = start.elapsed().as_micros() as u64;
765        if let Ok(Some(ref chunk)) = result {
766            // Logical rows (visible * multiplicity), not physical length.
767            self.record_profile_rows(chunk.logical_len());
768        }
769        if matches!(&result, Ok(None)) || (one_shot && matches!(&result, Ok(Some(_)))) {
770            self.base_mut().lifecycle.mark_exhausted();
771            // Ran to exhaustion: this instance's row counts are complete
772            // and safe for the feedback loop. Instances that never exhaust
773            // were starved downstream (Limit / early termination).
774            self.base().mark_profile_exhausted();
775        }
776        self.record_profile_timing("next", elapsed);
777        if result.is_err() {
778            self.base_mut().lifecycle.mark_failed();
779            if let Err(close_error) = self.close_tree() {
780                log::warn!(
781                    "Failed to close streaming executor tree after next error: {}",
782                    close_error
783                );
784            }
785        }
786        result
787    }
788
789    /// Stop the executor (signal no more input needed).
790    pub fn stop(&mut self) -> Result<(), QueryError> {
791        if matches!(
792            self.base().lifecycle,
793            OperatorLifecycle::Stopped | OperatorLifecycle::Closed
794        ) {
795            return Ok(());
796        }
797        let start = Instant::now();
798        let result = dispatch_stop!(self);
799        let elapsed = start.elapsed().as_micros() as u64;
800        self.record_profile_timing("stop", elapsed);
801        if result.is_ok() {
802            self.base_mut().lifecycle.mark_stopped();
803        }
804        result
805    }
806
807    /// Close the executor (clean up resources).
808    pub fn close(&mut self) -> Result<(), QueryError> {
809        if matches!(self.base().lifecycle, OperatorLifecycle::Closed) {
810            return Ok(());
811        }
812        // Capture spill metrics before operator close clears state.
813        let peak = self.peak_memory_bytes();
814        let spilled = self.spilled_size();
815        let sc = self.spill_count();
816        let rows = self.spilled_rows();
817        let start = Instant::now();
818        let result = dispatch_close!(self);
819        let elapsed = start.elapsed().as_micros() as u64;
820        self.record_profile_timing("close", elapsed);
821        if peak > 0 {
822            self.record_profile_peak_memory(peak);
823        }
824        if spilled > 0 || sc > 0 || rows > 0 {
825            self.record_profile_spill(spilled, sc, rows);
826        }
827        if result.is_ok() {
828            self.base_mut().lifecycle.mark_closed();
829        }
830        result
831    }
832
833    /// Close the executor tree in post-order. Individual operators only
834    /// release their own state; this is the sole recursive owner.
835    pub fn close_tree(&mut self) -> Result<(), QueryError> {
836        let mut first_error = None;
837        for child in self.children_mut() {
838            if let Err(error) = child.close_tree() {
839                if first_error.is_none() {
840                    first_error = Some(error);
841                }
842            }
843        }
844        if let Err(error) = self.close() {
845            if first_error.is_none() {
846                first_error = Some(error);
847            }
848        }
849        first_error.map_or(Ok(()), Err)
850    }
851
852    /// Stop the executor tree. The root is signalled before recursively
853    /// stopping children so coordinator operators can cancel worker handles.
854    pub fn stop_tree(&mut self) -> Result<(), QueryError> {
855        let mut first_error = self.stop().err();
856        for child in self.children_mut() {
857            if let Err(error) = child.stop_tree() {
858                if first_error.is_none() {
859                    first_error = Some(error);
860                }
861            }
862        }
863        first_error.map_or(Ok(()), Err)
864    }
865}
866
867impl Spillable for StreamingExecutor {
868    fn spill_to_disk(&mut self) -> Result<(), QueryError> {
869        let sm = self
870            .base()
871            .spill_manager()
872            .ok_or_else(|| QueryError::execution("Spill manager not configured"))?;
873        match self {
874            Self::Blocking(_, _, op) => op.spill_with_manager(&sm),
875            Self::Join(_, _, _, op) => op.spill_with_manager(&sm),
876            Self::Set(_, _, _, op) => op.spill_with_manager(&sm),
877            Self::Apply(_, _, _, _op) => {
878                // ApplyOperator doesn't accumulate enough memory;
879                // propagate to children instead.
880                for child in self.children_mut() {
881                    child.spill_to_disk()?;
882                }
883                Ok(())
884            }
885            Self::Wco(_, _, _, op) => op.spill_with_manager(&sm),
886            // Other operators (Source, Unary, Graph, etc.) don't accumulate
887            // enough memory to warrant spill.  Propagate to children so that
888            // deep-nested blocking operators can still be reached.
889            Self::Source(..)
890            | Self::Unary(..)
891            | Self::Graph(..)
892            | Self::RecursiveFragment(..)
893            | Self::Sink(..)
894            | Self::Ddl(..)
895            | Self::Fulltext(..)
896            | Self::Vector(..)
897            | Self::Txn(..)
898            | Self::Gather(..)
899            | Self::Exchange(..)
900            | Self::HashShuffleJoin(..) => {
901                for child in self.children_mut() {
902                    child.spill_to_disk()?;
903                }
904                Ok(())
905            }
906        }
907    }
908
909    fn spilled_size(&self) -> u64 {
910        match self {
911            Self::Blocking(_, _, op) => op.spilled_bytes(),
912            Self::Join(_, _, _, op) => op.spilled_bytes(),
913            Self::Set(_, _, _, op) => op.spilled_bytes(),
914            Self::Wco(_, _, _, op) => op.spilled_bytes(),
915            _ => 0,
916        }
917    }
918
919    fn spill_count(&self) -> u64 {
920        match self {
921            Self::Blocking(_, _, op) => op.spill_count(),
922            Self::Join(_, _, _, op) => op.spill_count(),
923            Self::Set(_, _, _, op) => op.spill_count(),
924            Self::Wco(_, _, _, _op) => 0,
925            _ => 0,
926        }
927    }
928
929    fn spilled_rows(&self) -> u64 {
930        match self {
931            Self::Blocking(_, _, op) => op.spilled_rows(),
932            Self::Join(_, _, _, op) => op.spilled_rows(),
933            Self::Set(_, _, _, op) => op.spilled_rows(),
934            Self::Wco(_, _, _, op) => op.spilled_rows(),
935            _ => 0,
936        }
937    }
938}
939
940#[cfg(test)]
941mod tests;
linkrs_query/executor/streaming/executor.rs

linkrs_query/executor/streaming/
executor.rs
