Skip to main content

linkrs_query/executor/streaming/operators/
state.rs

1//! OperatorState: Mutable per-execution state for operator nodes.
2//!
3//! Created fresh for each query execution from an immutable
4//! [`OperatorSpec`](super::spec::OperatorSpec).  Holds cursors,
5//! buffers, hash tables, row iterators, and all other mutable data that
6//! must NOT be shared across concurrent executions of the same cached plan.
7
8use std::collections::HashSet;
9
10use crate::executor::base::MemoryTracker;
11use crate::storage::{EdgeCursor, IndexCursor, ShardSlice, VertexCursor};
12use linkrs_core::types::expr::Expression;
13use linkrs_core::Value;
14
15use super::super::chunk::DataChunk;
16use super::super::executor::SortDirection;
17use super::source_operator::NeighborScanState;
18use super::spec::SourceSpec;
19
20// ── Source state ─────────────────────────────────────────────────────────────
21
22/// Mutable execution state for source operators.
23#[derive(Debug)]
24pub enum SourceState {
25    ScanVertices {
26        current_index: usize,
27        col_names: Vec<String>,
28    },
29    StorageScanVertices {
30        partition_id: usize,
31        partition_range: Option<std::ops::Range<i64>>,
32        cursor: Option<Box<dyn VertexCursor>>,
33        buffer: Vec<Vec<Value>>,
34        current_index: usize,
35        col_names: Vec<String>,
36    },
37    ScanEdges {
38        current_index: usize,
39        col_names: Vec<String>,
40    },
41    StorageScanEdges {
42        partition_id: usize,
43        shard_slice: Option<ShardSlice>,
44        cursor: Option<Box<dyn EdgeCursor>>,
45        buffer: Vec<Vec<Value>>,
46        current_index: usize,
47        col_names: Vec<String>,
48    },
49    GetVertices {
50        position: usize,
51    },
52    GetEdges {
53        cursor: Option<Box<dyn EdgeCursor>>,
54    },
55    GetNeighbors {
56        state: NeighborScanState,
57    },
58    IndexScan {
59        cursor: Option<Box<dyn IndexCursor<Row = Value>>>,
60    },
61    Argument,
62    GetProp {
63        entity_slot: usize,
64        prop_names: Vec<String>,
65    },
66    Start {
67        emitted: bool,
68    },
69}
70
71impl SourceState {
72    /// Create a fresh state for the given spec.
73    pub fn from_spec(spec: &SourceSpec) -> Self {
74        match spec {
75            SourceSpec::ScanVertices { rows: _, col_names } => SourceState::ScanVertices {
76                current_index: 0,
77                col_names: col_names.clone(),
78            },
79            SourceSpec::StandaloneValues { col_names, .. } => SourceState::ScanVertices {
80                current_index: 0,
81                col_names: col_names.clone(),
82            },
83            SourceSpec::StorageScanVertices { col_names, .. } => SourceState::StorageScanVertices {
84                partition_id: 0,
85                partition_range: None,
86                cursor: None,
87                buffer: Vec::new(),
88                current_index: 0,
89                col_names: col_names.clone(),
90            },
91            SourceSpec::ScanEdges { rows: _, col_names } => SourceState::ScanEdges {
92                current_index: 0,
93                col_names: col_names.clone(),
94            },
95            SourceSpec::StorageScanEdges { col_names, .. } => SourceState::StorageScanEdges {
96                partition_id: 0,
97                shard_slice: None,
98                cursor: None,
99                buffer: Vec::new(),
100                current_index: 0,
101                col_names: col_names.clone(),
102            },
103            SourceSpec::GetVertices { .. } => SourceState::GetVertices { position: 0 },
104            SourceSpec::GetEdges { .. } => SourceState::GetEdges { cursor: None },
105            SourceSpec::GetNeighbors { .. } => SourceState::GetNeighbors {
106                state: NeighborScanState::Init,
107            },
108            SourceSpec::IndexScan { .. } => SourceState::IndexScan { cursor: None },
109            SourceSpec::Argument { .. } => SourceState::Argument,
110            // CTE scans buffer rows in the operator kind itself; the arena
111            // state only records scan identity for diagnostics.
112            SourceSpec::CteScan { col_names, .. } => SourceState::ScanVertices {
113                current_index: 0,
114                col_names: col_names.clone(),
115            },
116            SourceSpec::GetProp {
117                entity_slot,
118                prop_names,
119                ..
120            } => SourceState::GetProp {
121                entity_slot: *entity_slot,
122                prop_names: prop_names.clone(),
123            },
124            SourceSpec::Start => SourceState::Start { emitted: false },
125        }
126    }
127}
128
129// ── Blocking state ───────────────────────────────────────────────────────────
130pub use super::blocking::{
131    AggregateState, DataCollectState, DistinctState, FinalAggregateState, GroupByState,
132    MaterializeState, PartialAggregateState, RollUpApplyState, SortState, TopNState,
133    WindowFunctionState, WindowState,
134};
135
136/// Mutable execution state for blocking operators.
137#[derive(Debug)]
138pub enum BlockingState {
139    Sort {
140        memory_tracker: MemoryTracker,
141        state: Option<SortState>,
142    },
143    Aggregate {
144        memory_tracker: MemoryTracker,
145        state: Option<AggregateState>,
146    },
147    GroupBy {
148        memory_tracker: MemoryTracker,
149        state: Option<GroupByState>,
150    },
151    WindowFunction {
152        memory_tracker: MemoryTracker,
153        state: Option<WindowFunctionState>,
154    },
155    Window {
156        memory_tracker: MemoryTracker,
157        state: Option<WindowState>,
158    },
159    TopN {
160        memory_tracker: MemoryTracker,
161        state: Option<TopNState>,
162    },
163    Distinct {
164        memory_tracker: MemoryTracker,
165        state: Option<DistinctState>,
166    },
167    Materialize {
168        memory_tracker: MemoryTracker,
169        state: Option<MaterializeState>,
170    },
171    DataCollect {
172        memory_tracker: MemoryTracker,
173        state: Option<DataCollectState>,
174    },
175    RollUpApply {
176        memory_tracker: MemoryTracker,
177        state: Option<RollUpApplyState>,
178    },
179    PartialAggregate {
180        memory_tracker: MemoryTracker,
181        state: Option<PartialAggregateState>,
182    },
183    FinalAggregate {
184        memory_tracker: MemoryTracker,
185        state: Option<FinalAggregateState>,
186    },
187}
188
189// ── Join state ───────────────────────────────────────────────────────────────
190
191/// Mutable execution state for join operators.
192#[derive(Debug)]
193pub enum JoinState {
194    HashJoin {
195        build_side: super::join_operator::HashJoinBuildSide,
196        left_consumed: bool,
197        memory_tracker: MemoryTracker,
198        right_col_names: Vec<String>,
199    },
200    HashLeftJoin {
201        build_side: super::join_operator::HashJoinBuildSide,
202        left_consumed: bool,
203        memory_tracker: MemoryTracker,
204        right_col_names: Vec<String>,
205    },
206    NestedLoopJoin {
207        build_side_tuples: Vec<Vec<Value>>,
208        left_consumed: bool,
209        memory_tracker: MemoryTracker,
210        right_col_names: Vec<String>,
211    },
212    InnerJoin {
213        build_side_tuples: Vec<Vec<Value>>,
214        left_consumed: bool,
215        memory_tracker: MemoryTracker,
216        right_col_names: Vec<String>,
217    },
218    LeftJoin {
219        build_side_tuples: Vec<Vec<Value>>,
220        left_consumed: bool,
221        memory_tracker: MemoryTracker,
222        right_col_names: Vec<String>,
223    },
224    RightJoin {
225        build_side_tuples: Vec<Vec<Value>>,
226        right_consumed: bool,
227        memory_tracker: MemoryTracker,
228        right_col_names: Vec<String>,
229    },
230    FullOuterJoin {
231        left_rows: Vec<Vec<Value>>,
232        right_rows: Vec<Vec<Value>>,
233        matched_right_indices: HashSet<usize>,
234        result_iter: Option<std::vec::IntoIter<Vec<Value>>>,
235        phase: super::super::executor::FullOuterJoinPhase,
236        memory_tracker: MemoryTracker,
237        right_col_names: Vec<String>,
238    },
239    CrossJoin {
240        all_left_rows: Vec<Vec<Value>>,
241        all_right_rows: Vec<Vec<Value>>,
242        left_consumed: bool,
243        right_consumed: bool,
244        memory_tracker: MemoryTracker,
245        right_col_names: Vec<String>,
246    },
247    SemiJoin {
248        right_rows: Vec<Vec<Value>>,
249        right_consumed: bool,
250        memory_tracker: MemoryTracker,
251        right_col_names: Vec<String>,
252    },
253}
254
255// ── Graph state ──────────────────────────────────────────────────────────────
256
257/// Mutable execution state for graph traversal operators.
258#[derive(Debug)]
259pub enum GraphState {
260    Expand,
261    ExpandAll,
262    Traverse {
263        visited: HashSet<String>,
264    },
265    BiExpand,
266    BiTraverse {
267        visited: HashSet<String>,
268    },
269    ShortestPath,
270    BFSShortest {
271        frontier: Vec<Vec<Value>>,
272        visited: HashSet<String>,
273    },
274    AllPaths {
275        all_paths: Vec<Vec<Value>>,
276        result_iter: Option<std::vec::IntoIter<Vec<Value>>>,
277    },
278    MultiShortestPath {
279        all_paths: Vec<Vec<Value>>,
280        result_iter: Option<std::vec::IntoIter<Vec<Value>>>,
281    },
282}
283
284// ── Sink state ───────────────────────────────────────────────────────────────
285
286/// Mutable execution state for sink (data modification) operators.
287#[derive(Debug)]
288pub enum SinkState {
289    InsertVertices { rows_inserted: u64 },
290    InsertEdges { rows_inserted: u64 },
291    UpdateVertices { rows_updated: u64 },
292    UpdateEdges { rows_updated: u64 },
293    DeleteVertices { rows_deleted: u64 },
294    DeleteEdges { rows_deleted: u64 },
295    PipeDeleteVertices { rows_deleted: u64 },
296    PipeDeleteEdges { rows_deleted: u64 },
297}
298
299// ── Exchange state ───────────────────────────────────────────────────────────
300
301/// Mutable execution state for exchange (gather / merge / repartition) operators.
302#[derive(Debug)]
303pub enum ExchangeState {
304    Concatenate {
305        current_index: usize,
306        col_names: Option<Vec<String>>,
307    },
308    MergeSort {
309        sort_expressions: Vec<Expression>,
310        sort_directions: Vec<SortDirection>,
311        inputs: Vec<MergeInputState>,
312        col_names: Option<Vec<String>>,
313        limit: Option<usize>,
314        emitted: usize,
315    },
316    /// Hash-based repartition: buffer input rows, rehash by key, route to
317    /// output buckets.  Each bucket is pulled sequentially.
318    RepartitionHash {
319        /// Number of output buckets / partitions.
320        num_partitions: usize,
321        /// Buffered input rows for each partition bucket.
322        /// Indexed by hash(keys) % num_partitions.
323        buckets: Vec<Vec<Vec<Value>>>,
324        /// Current bucket being drained.
325        current_bucket: usize,
326        /// Current row index within `current_bucket`.
327        current_row: usize,
328        /// Hash expressions (from spec, cloned for state).
329        hash_expressions: Vec<Expression>,
330        /// Column names of the data flowing through.
331        col_names: Option<Vec<String>>,
332    },
333    /// Broadcast: replicates every input chunk across N output channels.
334    Broadcast {
335        /// Number of consumers / output channels.
336        num_consumers: usize,
337        /// Buffered chunks from upstream (all consumed before broadcast).
338        buffered_chunks: Vec<DataChunk>,
339        /// Current position within the broadcast output sequence.
340        /// Cycles through consumers for each row/chunk.
341        current_consumer: usize,
342        /// Current chunk index being broadcast.
343        chunk_index: usize,
344        /// Next row index within the current chunk.
345        row_index: usize,
346    },
347    /// Barrier: collect input-fragment completion signals, then pass through.
348    Barrier {
349        /// Whether the barrier has been passed.
350        passed: bool,
351    },
352    /// Materialize: fully consume child output before producing it.
353    Materialize {
354        /// All rows collected from children.
355        rows: Vec<Vec<Value>>,
356        /// Current position in the materialized output.
357        position: usize,
358        /// Column names.
359        col_names: Option<Vec<String>>,
360    },
361}
362
363/// Internal merge cursor state for Exchange merge-sort.
364#[derive(Debug)]
365pub enum MergeInputState {
366    Pending,
367    Buffered { chunk: DataChunk, row_index: usize },
368    Exhausted,
369}
370
371// ── Set state ────────────────────────────────────────────────────────────────
372
373/// Mutable execution state for set operators.
374#[derive(Debug)]
375pub enum SetState {
376    Union {
377        seen_rows: HashSet<String>,
378        left_consumed: bool,
379        memory_tracker: MemoryTracker,
380    },
381    UnionAll {
382        left_consumed: bool,
383    },
384    Intersect {
385        left_rows: Vec<Vec<Value>>,
386        right_rows: HashSet<String>,
387        left_buffered: bool,
388        right_buffered: bool,
389        memory_tracker: MemoryTracker,
390    },
391    Except {
392        exclude_rows: HashSet<String>,
393        right_buffered: bool,
394        memory_tracker: MemoryTracker,
395    },
396    Minus {
397        exclude_rows: HashSet<String>,
398        right_buffered: bool,
399        memory_tracker: MemoryTracker,
400    },
401}
402
403// ── Apply state ──────────────────────────────────────────────────────────────
404
405/// Mutable execution state for apply operators.
406#[derive(Debug)]
407pub enum ApplyState {
408    Apply,
409    PatternApply {
410        all_rows: Vec<Vec<Value>>,
411        result_iter: Option<std::vec::IntoIter<Vec<Value>>>,
412        memory_tracker: MemoryTracker,
413    },
414    RollUpApply,
415}
416
417// ── DDL state ────────────────────────────────────────────────────────────────
418
419/// Mutable execution state for DDL operators (minimal — DDL has no cursor/accumulator).
420#[derive(Debug)]
421pub enum DdlState {
422    SpaceManage,
423    TagManage,
424    EdgeManage,
425    IndexManage,
426    UserManage,
427    ShowStats,
428    ShowConfigs,
429    ShowQueries,
430    ShowSessions,
431    ShowFunctions,
432    ShowGraphs,
433    ShowMacros,
434    ShowAttachedDatabases,
435    ShowExtensions,
436    LoadFrom,
437    InQueryCall,
438    Analyze,
439    Migrate,
440    MigratePlan,
441    MigrateRun,
442    MigrateRollback,
443    SequenceManage,
444    MacroManage,
445    TypeManage,
446}
447impl DdlState {
448    pub fn from_spec(spec: &super::spec::DdlSpec) -> Self {
449        match spec {
450            super::spec::DdlSpec::SpaceManage { .. } => DdlState::SpaceManage,
451            super::spec::DdlSpec::TagManage { .. } => DdlState::TagManage,
452            super::spec::DdlSpec::EdgeManage { .. } => DdlState::EdgeManage,
453            super::spec::DdlSpec::IndexManage { .. } => DdlState::IndexManage,
454            super::spec::DdlSpec::DeleteIndex { .. } => DdlState::IndexManage,
455            super::spec::DdlSpec::UserManage { .. } => DdlState::UserManage,
456            super::spec::DdlSpec::ShowStats { .. } => DdlState::ShowStats,
457            super::spec::DdlSpec::ShowConfigs { .. } => DdlState::ShowConfigs,
458            super::spec::DdlSpec::ShowQueries { .. } => DdlState::ShowQueries,
459            super::spec::DdlSpec::ShowSessions { .. } => DdlState::ShowSessions,
460            super::spec::DdlSpec::ShowFunctions { .. } => DdlState::ShowFunctions,
461            super::spec::DdlSpec::ShowGraphs { .. } => DdlState::ShowGraphs,
462            super::spec::DdlSpec::ShowMacros { .. } => DdlState::ShowMacros,
463            super::spec::DdlSpec::ShowAttachedDatabases { .. } => DdlState::ShowAttachedDatabases,
464            super::spec::DdlSpec::ShowExtensions { .. } => DdlState::ShowExtensions,
465            super::spec::DdlSpec::LoadFrom { .. } => DdlState::LoadFrom,
466            super::spec::DdlSpec::InQueryCall { .. } => DdlState::InQueryCall,
467            super::spec::DdlSpec::Analyze { .. } => DdlState::Analyze,
468            super::spec::DdlSpec::Migrate { .. } => DdlState::Migrate,
469            super::spec::DdlSpec::MigratePlan { .. } => DdlState::MigratePlan,
470            super::spec::DdlSpec::MigrateRun { .. } => DdlState::MigrateRun,
471            super::spec::DdlSpec::MigrateRollback { .. } => DdlState::MigrateRollback,
472            super::spec::DdlSpec::SequenceManage { .. } => DdlState::SequenceManage,
473            super::spec::DdlSpec::MacroManage { .. } => DdlState::MacroManage,
474            super::spec::DdlSpec::TypeManage { .. } => DdlState::TypeManage,
475        }
476    }
477}
478
479// ── Fulltext state ───────────────────────────────────────────────────────────
480
481/// Mutable execution state for fulltext operators.
482#[derive(Debug)]
483pub enum FulltextState {
484    FulltextManage,
485    FulltextSearch,
486    FulltextLookup,
487    MatchFulltext,
488}
489
490impl FulltextState {
491    pub fn from_spec(spec: &super::spec::FulltextSpec) -> Self {
492        match spec {
493            super::spec::FulltextSpec::FulltextManage { .. } => FulltextState::FulltextManage,
494            super::spec::FulltextSpec::FulltextSearch { .. } => FulltextState::FulltextSearch,
495            super::spec::FulltextSpec::FulltextLookup { .. } => FulltextState::FulltextLookup,
496            super::spec::FulltextSpec::MatchFulltext { .. } => FulltextState::MatchFulltext,
497        }
498    }
499}
500
501// ── Vector state ─────────────────────────────────────────────────────────────
502
503/// Mutable execution state for vector search operators.
504#[derive(Debug)]
505pub enum VectorState {
506    VectorManage,
507    VectorSearch,
508    VectorLookup,
509    VectorMatch,
510}
511
512impl VectorState {
513    pub fn from_spec(spec: &super::spec::VectorSpec) -> Self {
514        match spec {
515            super::spec::VectorSpec::VectorManage { .. } => VectorState::VectorManage,
516            super::spec::VectorSpec::VectorSearch { .. } => VectorState::VectorSearch,
517            super::spec::VectorSpec::VectorLookup { .. } => VectorState::VectorLookup,
518            super::spec::VectorSpec::VectorMatch { .. } => VectorState::VectorMatch,
519        }
520    }
521}
522
523// ── RecursiveFragment state ─────────────────────────────────────────────
524
525/// Mutable execution state for recursive fragment operators.
526#[derive(Debug)]
527pub enum RecursiveFragmentState {
528    ShortestPath,
529    MultiShortestPath,
530    BFSShortest,
531    AllPaths,
532    Fixpoint,
533}
534
535impl RecursiveFragmentState {
536    pub fn from_spec(spec: &super::spec::RecursiveFragmentSpec) -> Self {
537        match spec {
538            super::spec::RecursiveFragmentSpec::ShortestPath { .. } => {
539                RecursiveFragmentState::ShortestPath
540            }
541            super::spec::RecursiveFragmentSpec::MultiShortestPath { .. } => {
542                RecursiveFragmentState::MultiShortestPath
543            }
544            super::spec::RecursiveFragmentSpec::BFSShortest { .. } => {
545                RecursiveFragmentState::BFSShortest
546            }
547            super::spec::RecursiveFragmentSpec::AllPaths { .. } => RecursiveFragmentState::AllPaths,
548            super::spec::RecursiveFragmentSpec::Fixpoint { .. } => RecursiveFragmentState::Fixpoint,
549        }
550    }
551}
552
553// ── Txn state ────────────────────────────────────────────────────────────────
554
555/// Mutable execution state for transaction operators.
556#[derive(Debug)]
557pub enum TxnState {
558    BeginTransaction,
559    Commit,
560    Rollback,
561    RollbackToSavepoint,
562    Savepoint,
563    ReleaseSavepoint,
564}
565
566impl TxnState {
567    pub fn from_spec(spec: &super::spec::TxnSpec) -> Self {
568        match spec {
569            super::spec::TxnSpec::BeginTransaction => TxnState::BeginTransaction,
570            super::spec::TxnSpec::Commit => TxnState::Commit,
571            super::spec::TxnSpec::Rollback => TxnState::Rollback,
572            super::spec::TxnSpec::RollbackToSavepoint { .. } => TxnState::RollbackToSavepoint,
573            super::spec::TxnSpec::Savepoint { .. } => TxnState::Savepoint,
574            super::spec::TxnSpec::ReleaseSavepoint { .. } => TxnState::ReleaseSavepoint,
575        }
576    }
577}