linkrs_query/executor/streaming/operators/
recursive_fragment_operator.rs

1use std::sync::Arc;
2
3use parking_lot::RwLock;
4
5use super::algorithms::{
6    bidir_bfs_shortest_path, enumerate_all_paths, materialize_path_endpoint, path_endpoint_pairs,
7    path_endpoint_tags, weighted_shortest_path, AllPathsConfig, BidirBfsConfig,
8};
9use crate::executor::expression::evaluator::traits::ExpressionContext;
10use crate::executor::streaming::chunk::{ColumnInfo, DataChunk, Schema};
11use crate::executor::streaming::context::ValueRowContext;
12use crate::executor::streaming::executor::StreamingExecutor;
13use crate::executor::streaming::operators::source_operator::OperatorConfig;
14use crate::executor::streaming::runtime::ExecutionRuntime;
15use crate::executor::streaming::slot::SlotLayout;
16use crate::planning::plan::core::nodes::traversal::path_algorithms::PathUniqueness;
17use crate::storage::QueryStorage;
18use linkrs_core::error::QueryError;
19use linkrs_core::{EdgeDirection, Value};
20
21use super::spec::RecursiveFragmentSpec;
22
23#[derive(Debug)]
24pub enum RecursiveFragmentOperatorKind {
25    ShortestPath {
26        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
27        space_name: String,
28        edge_types: Vec<String>,
29        direction: EdgeDirection,
30        max_depth: usize,
31        start_vertices: Vec<Value>,
32        target_vertices: Vec<Value>,
33        vertex_tag: String,
34        weight: Option<String>,
35    },
36    MultiShortestPath {
37        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
38        space_name: String,
39        edge_types: Vec<String>,
40        direction: EdgeDirection,
41        max_depth: usize,
42        left_vertex_column: String,
43        right_vertex_column: String,
44        single_shortest: bool,
45        vertex_tag: String,
46    },
47    BFSShortest {
48        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
49        space_name: String,
50        edge_types: Vec<String>,
51        direction: EdgeDirection,
52        max_depth: usize,
53        allow_loops: bool,
54        vertex_tag: String,
55    },
56    AllPaths {
57        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
58        space_name: String,
59        edge_types: Vec<String>,
60        direction: EdgeDirection,
61        min_depth: usize,
62        max_depth: usize,
63        path_uniqueness: PathUniqueness,
64        limit: Option<usize>,
65        offset: usize,
66        start_vertices: Vec<Value>,
67        target_vertices: Vec<Value>,
68        vertex_tag: String,
69    },
70    Fixpoint {
71        /// Mangled CTE tag identifying the working table on the runtime.
72        cte_name: String,
73        /// Anchor sub-plan (runs once, seeds the working table).
74        anchor: Arc<crate::executor::streaming::plan::types::PhysicalPlan>,
75        /// Step sub-plan (re-runs per iteration over the delta).
76        step: Option<Arc<crate::executor::streaming::plan::types::PhysicalPlan>>,
77        /// Fixpoint iteration cap (errors when exceeded).
78        max_iterations: u64,
79        /// Output column names (single column in V1).
80        col_names: Vec<String>,
81        /// Lazily materialized child executors, reused across iterations.
82        state: FixpointState,
83    },
84}
85
86/// Mutable fixpoint execution state, owned by the operator.
87///
88/// Sub-plan executors are built and dropped inside a single `run_fixpoint`
89/// call; only the converged result is retained here.
90#[derive(Debug, Default)]
91pub struct FixpointState {
92    /// Final converged rows once computed; served chunk by chunk.
93    output: Option<Vec<Vec<Value>>>,
94    output_pos: usize,
95}
96
97/// Recursive fragment operator.
98///
99/// Wraps [`RecursiveFragmentOperatorKind`] with the runtime context injected
100/// at `open()`. Lifecycle state is owned exclusively by the executor;
101/// operators never write it.
102#[derive(Debug)]
103pub struct RecursiveFragmentOperator {
104    pub kind: RecursiveFragmentOperatorKind,
105    pub runtime: Option<Arc<ExecutionRuntime>>,
106    pub output_layout: Arc<SlotLayout>,
107    pub config: OperatorConfig,
108}
109
110impl RecursiveFragmentOperator {
111    pub fn from_spec(
112        spec: &RecursiveFragmentSpec,
113        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
114        space_name: String,
115        output_layout: Arc<SlotLayout>,
116    ) -> Self {
117        let kind = match spec {
118            RecursiveFragmentSpec::ShortestPath {
119                edge_types,
120                direction,
121                max_depth,
122                start_vertices,
123                target_vertices,
124                vertex_tag,
125                weight,
126            } => RecursiveFragmentOperatorKind::ShortestPath {
127                storage,
128                space_name,
129                edge_types: edge_types.clone(),
130                direction: *direction,
131                max_depth: *max_depth,
132                start_vertices: start_vertices.clone(),
133                target_vertices: target_vertices.clone(),
134                vertex_tag: vertex_tag.clone(),
135                weight: weight.clone(),
136            },
137            RecursiveFragmentSpec::MultiShortestPath {
138                edge_types,
139                direction,
140                max_depth,
141                left_vertex_column,
142                right_vertex_column,
143                single_shortest,
144                vertex_tag,
145            } => RecursiveFragmentOperatorKind::MultiShortestPath {
146                storage,
147                space_name,
148                edge_types: edge_types.clone(),
149                direction: *direction,
150                max_depth: *max_depth,
151                left_vertex_column: left_vertex_column.clone(),
152                right_vertex_column: right_vertex_column.clone(),
153                single_shortest: *single_shortest,
154                vertex_tag: vertex_tag.clone(),
155            },
156            RecursiveFragmentSpec::BFSShortest {
157                edge_types,
158                direction,
159                max_depth,
160                allow_loops,
161                vertex_tag,
162            } => RecursiveFragmentOperatorKind::BFSShortest {
163                storage,
164                space_name,
165                edge_types: edge_types.clone(),
166                direction: *direction,
167                max_depth: *max_depth,
168                allow_loops: *allow_loops,
169                vertex_tag: vertex_tag.clone(),
170            },
171            RecursiveFragmentSpec::AllPaths {
172                edge_types,
173                direction,
174                min_depth,
175                max_depth,
176                path_uniqueness,
177                limit,
178                offset,
179                start_vertices,
180                target_vertices,
181                vertex_tag,
182            } => RecursiveFragmentOperatorKind::AllPaths {
183                storage,
184                space_name,
185                edge_types: edge_types.clone(),
186                direction: *direction,
187                min_depth: *min_depth,
188                max_depth: *max_depth,
189                path_uniqueness: *path_uniqueness,
190                limit: *limit,
191                offset: *offset,
192                start_vertices: start_vertices.clone(),
193                target_vertices: target_vertices.clone(),
194                vertex_tag: vertex_tag.clone(),
195            },
196            RecursiveFragmentSpec::Fixpoint {
197                cte_name,
198                anchor,
199                step,
200                max_iterations,
201                col_names,
202            } => RecursiveFragmentOperatorKind::Fixpoint {
203                cte_name: cte_name.clone(),
204                anchor: anchor.clone(),
205                step: step.clone(),
206                max_iterations: *max_iterations,
207                col_names: col_names.clone(),
208                state: FixpointState::default(),
209            },
210        };
211        Self::new(kind, output_layout)
212    }
213
214    pub fn new(kind: RecursiveFragmentOperatorKind, output_layout: Arc<SlotLayout>) -> Self {
215        Self {
216            kind,
217            runtime: None,
218            output_layout,
219            config: OperatorConfig::default(),
220        }
221    }
222
223    /// Inject the runtime and execution config (called once by the executor
224    /// before this operator produces any data).
225    pub fn inject_context(
226        &mut self,
227        runtime: Option<&Arc<ExecutionRuntime>>,
228        config: OperatorConfig,
229    ) {
230        if let Some(rt) = runtime {
231            self.runtime = Some(rt.clone());
232        }
233        self.config = config;
234    }
235
236    pub fn open(&mut self, input: &mut StreamingExecutor) -> Result<(), QueryError> {
237        input.open()?;
238        Ok(())
239    }
240
241    pub fn next(&mut self, input: &mut StreamingExecutor) -> Result<Option<DataChunk>, QueryError> {
242        match &mut self.kind {
243            RecursiveFragmentOperatorKind::ShortestPath {
244                storage,
245                space_name,
246                edge_types,
247                direction,
248                max_depth,
249                start_vertices,
250                target_vertices,
251                vertex_tag,
252                weight,
253            } => loop {
254                let Some(mut chunk) = input.advance()? else {
255                    return Ok(None);
256                };
257                chunk.normalize_for_opaque("RecursiveFragment");
258                if let Some(storage_lock) = storage {
259                    let reader = storage_lock.read();
260                    let col_names = chunk.col_names();
261                    let layout = chunk.get_layout();
262                    let endpoint_tags = path_endpoint_tags(
263                        &*reader, space_name, edge_types, *direction, vertex_tag,
264                    );
265                    let (start_tag, end_tag) = endpoint_tags
266                        .as_ref()
267                        .map(|(start, end)| (start.as_str(), end.as_str()))
268                        .unzip();
269                    let mut out_rows = Vec::new();
270                    for row in &chunk.rows {
271                        if let Some(rt) = self.runtime.as_ref() {
272                            rt.ensure_not_cancelled()?;
273                        }
274                        let pairs = path_endpoint_pairs(
275                            row,
276                            layout.clone(),
277                            start_vertices,
278                            target_vertices,
279                            None,
280                        )?;
281                        let et_ref: Option<&[String]> = if edge_types.is_empty() {
282                            None
283                        } else {
284                            Some(edge_types.as_slice())
285                        };
286                        for (src_val, dst_val) in pairs {
287                            let Some(start_vertex) = materialize_path_endpoint(
288                                &*reader, space_name, &src_val, vertex_tag, start_tag,
289                            )?
290                            else {
291                                continue;
292                            };
293                            let Some(end_vertex) = materialize_path_endpoint(
294                                &*reader, space_name, &dst_val, vertex_tag, end_tag,
295                            )?
296                            else {
297                                continue;
298                            };
299                            if let Some(rt) = self.runtime.as_ref() {
300                                rt.ensure_not_cancelled()?;
301                            }
302                            let cancel_token = self.runtime.as_ref().map(|rt| rt.cancel_token());
303                            let paths = if let Some(weight_prop) = weight {
304                                // Weighted variant: Dijkstra on the named edge
305                                // property; one cheapest path per (src, dst).
306                                weighted_shortest_path(
307                                    &*reader,
308                                    &start_vertex,
309                                    &end_vertex,
310                                    space_name,
311                                    vertex_tag,
312                                    et_ref,
313                                    *direction,
314                                    *max_depth,
315                                    weight_prop,
316                                    cancel_token.as_ref(),
317                                )?
318                                .into_iter()
319                                .collect()
320                            } else {
321                                bidir_bfs_shortest_path(
322                                    &*reader,
323                                    &start_vertex,
324                                    &end_vertex,
325                                    BidirBfsConfig {
326                                        space_name,
327                                        vertex_tag,
328                                        edge_type_filter: et_ref,
329                                        max_depth: *max_depth,
330                                        single_shortest: false,
331                                        limit: 1000,
332                                        direction: *direction,
333                                    },
334                                    cancel_token.as_ref(),
335                                )?
336                            };
337                            for path in &paths {
338                                if let Some(rt) = self.runtime.as_ref() {
339                                    rt.ensure_not_cancelled()?;
340                                }
341                                let mut out_row = row.clone();
342                                out_row.push(Value::Path(Box::new(path.clone())));
343                                out_rows.push(out_row);
344                            }
345                        }
346                    }
347                    if out_rows.is_empty() {
348                        continue;
349                    }
350                    let mut new_cols: Vec<ColumnInfo> = col_names
351                        .iter()
352                        .map(|n| ColumnInfo {
353                            name: n.clone(),
354                            data_type: "string".to_string(),
355                        })
356                        .collect();
357                    new_cols.push(ColumnInfo {
358                        name: "path".to_string(),
359                        data_type: "path".to_string(),
360                    });
361                    let _schema = Arc::new(Schema::new(new_cols));
362                    return Ok(Some(DataChunk::new_with_layout(
363                        out_rows,
364                        Arc::clone(&self.output_layout),
365                    )));
366                } else {
367                    return Ok(Some(chunk));
368                }
369            },
370
371            RecursiveFragmentOperatorKind::MultiShortestPath {
372                storage,
373                space_name,
374                edge_types,
375                direction,
376                max_depth,
377                left_vertex_column,
378                right_vertex_column,
379                single_shortest,
380                vertex_tag,
381                ..
382            } => loop {
383                let Some(mut chunk) = input.advance()? else {
384                    return Ok(None);
385                };
386                chunk.normalize_for_opaque("RecursiveFragment");
387                if let Some(storage_lock) = storage {
388                    let reader = storage_lock.read();
389                    let col_names = chunk.col_names();
390                    let layout = chunk.get_layout();
391                    let endpoint_tags = path_endpoint_tags(
392                        &*reader, space_name, edge_types, *direction, vertex_tag,
393                    );
394                    let (start_tag, end_tag) = endpoint_tags
395                        .as_ref()
396                        .map(|(start, end)| (start.as_str(), end.as_str()))
397                        .unzip();
398                    let mut out_rows = Vec::new();
399                    for row in &chunk.rows {
400                        if let Some(rt) = self.runtime.as_ref() {
401                            rt.ensure_not_cancelled()?;
402                        }
403                        let ctx = ValueRowContext::new(row.clone(), layout.clone());
404                        let left_val = ctx
405                            .get_variable(left_vertex_column)
406                            .or_else(|| ctx.get_variable("vid"))
407                            .or_else(|| row.first().cloned())
408                            .unwrap_or(Value::Null(linkrs_core::NullType::Null));
409                        let right_val = ctx
410                            .get_variable(right_vertex_column)
411                            .or_else(|| ctx.get_variable("dst_vid"))
412                            .or_else(|| row.get(1).cloned())
413                            .unwrap_or(Value::Null(linkrs_core::NullType::Null));
414                        let Some(start_vertex) = materialize_path_endpoint(
415                            &*reader, space_name, &left_val, vertex_tag, start_tag,
416                        )?
417                        else {
418                            continue;
419                        };
420                        let Some(end_vertex) = materialize_path_endpoint(
421                            &*reader, space_name, &right_val, vertex_tag, end_tag,
422                        )?
423                        else {
424                            continue;
425                        };
426                        let et_ref: Option<&[String]> = if edge_types.is_empty() {
427                            None
428                        } else {
429                            Some(edge_types.as_slice())
430                        };
431                        if let Some(rt) = self.runtime.as_ref() {
432                            rt.ensure_not_cancelled()?;
433                        }
434                        let cancel_token = self.runtime.as_ref().map(|rt| rt.cancel_token());
435                        let paths = bidir_bfs_shortest_path(
436                            &*reader,
437                            &start_vertex,
438                            &end_vertex,
439                            BidirBfsConfig {
440                                space_name,
441                                vertex_tag,
442                                edge_type_filter: et_ref,
443                                max_depth: *max_depth,
444                                single_shortest: *single_shortest,
445                                limit: if *single_shortest { 1 } else { 10 },
446                                direction: *direction,
447                            },
448                            cancel_token.as_ref(),
449                        )?;
450                        for path in &paths {
451                            if let Some(rt) = self.runtime.as_ref() {
452                                rt.ensure_not_cancelled()?;
453                            }
454                            let mut out_row = row.clone();
455                            out_row.push(Value::Path(Box::new(path.clone())));
456                            out_rows.push(out_row);
457                        }
458                    }
459                    if out_rows.is_empty() {
460                        continue;
461                    }
462                    let mut new_cols: Vec<ColumnInfo> = col_names
463                        .iter()
464                        .map(|n| ColumnInfo {
465                            name: n.clone(),
466                            data_type: "string".to_string(),
467                        })
468                        .collect();
469                    new_cols.push(ColumnInfo {
470                        name: "_multi_shortest_path".to_string(),
471                        data_type: "path".to_string(),
472                    });
473                    let _schema = Arc::new(Schema::new(new_cols));
474                    return Ok(Some(DataChunk::new_with_layout(
475                        out_rows,
476                        Arc::clone(&self.output_layout),
477                    )));
478                } else {
479                    return Ok(Some(chunk));
480                }
481            },
482
483            RecursiveFragmentOperatorKind::BFSShortest {
484                storage,
485                space_name,
486                edge_types,
487                direction,
488                max_depth,
489                allow_loops,
490                vertex_tag,
491            } => loop {
492                let Some(mut chunk) = input.advance()? else {
493                    return Ok(None);
494                };
495                chunk.normalize_for_opaque("RecursiveFragment");
496                if let Some(storage_lock) = storage {
497                    let reader = storage_lock.read();
498                    let col_names = chunk.col_names();
499                    let layout = chunk.get_layout();
500                    let endpoint_tags = path_endpoint_tags(
501                        &*reader, space_name, edge_types, *direction, vertex_tag,
502                    );
503                    let (start_tag, end_tag) = endpoint_tags
504                        .as_ref()
505                        .map(|(start, end)| (start.as_str(), end.as_str()))
506                        .unzip();
507                    let mut out_rows = Vec::new();
508                    for row in &chunk.rows {
509                        if let Some(rt) = self.runtime.as_ref() {
510                            rt.ensure_not_cancelled()?;
511                        }
512                        let ctx = ValueRowContext::new(row.clone(), layout.clone());
513                        let vid_val = ctx
514                            .get_variable("vid")
515                            .or_else(|| row.first().cloned())
516                            .unwrap_or(Value::Null(linkrs_core::NullType::Null));
517                        let end_val = ctx
518                            .get_variable("dst_vid")
519                            .or_else(|| ctx.get_variable("target"))
520                            .or_else(|| row.get(1).cloned())
521                            .unwrap_or(Value::Null(linkrs_core::NullType::Null));
522                        let Some(start_vertex) = materialize_path_endpoint(
523                            &*reader, space_name, &vid_val, vertex_tag, start_tag,
524                        )?
525                        else {
526                            continue;
527                        };
528                        let Some(end_vertex) = materialize_path_endpoint(
529                            &*reader, space_name, &end_val, vertex_tag, end_tag,
530                        )?
531                        else {
532                            continue;
533                        };
534                        let et_ref: Option<&[String]> = if edge_types.is_empty() {
535                            None
536                        } else {
537                            Some(edge_types.as_slice())
538                        };
539                        let cancel_token = self.runtime.as_ref().map(|rt| rt.cancel_token());
540                        let paths = bidir_bfs_shortest_path(
541                            &*reader,
542                            &start_vertex,
543                            &end_vertex,
544                            BidirBfsConfig {
545                                space_name,
546                                vertex_tag,
547                                edge_type_filter: et_ref,
548                                max_depth: *max_depth,
549                                single_shortest: !*allow_loops,
550                                limit: if *allow_loops { 1000 } else { 1 },
551                                direction: *direction,
552                            },
553                            cancel_token.as_ref(),
554                        )?;
555                        for path in &paths {
556                            if let Some(rt) = self.runtime.as_ref() {
557                                rt.ensure_not_cancelled()?;
558                            }
559                            let mut out_row = row.clone();
560                            out_row.push(Value::Path(Box::new(path.clone())));
561                            out_rows.push(out_row);
562                        }
563                    }
564                    if out_rows.is_empty() {
565                        continue;
566                    }
567                    let mut new_cols: Vec<ColumnInfo> = col_names
568                        .iter()
569                        .map(|n| ColumnInfo {
570                            name: n.clone(),
571                            data_type: "string".to_string(),
572                        })
573                        .collect();
574                    new_cols.push(ColumnInfo {
575                        name: "_bfs_path".to_string(),
576                        data_type: "path".to_string(),
577                    });
578                    let _schema = Arc::new(Schema::new(new_cols));
579                    return Ok(Some(DataChunk::new_with_layout(
580                        out_rows,
581                        Arc::clone(&self.output_layout),
582                    )));
583                } else {
584                    return Ok(Some(chunk));
585                }
586            },
587
588            RecursiveFragmentOperatorKind::AllPaths {
589                storage,
590                space_name,
591                edge_types,
592                direction,
593                min_depth,
594                max_depth,
595                path_uniqueness,
596                limit,
597                offset,
598                start_vertices,
599                target_vertices,
600                vertex_tag,
601            } => loop {
602                let Some(mut chunk) = input.advance()? else {
603                    return Ok(None);
604                };
605                chunk.normalize_for_opaque("RecursiveFragment");
606                if let Some(storage_lock) = storage {
607                    let reader = storage_lock.read();
608                    let col_names = chunk.col_names();
609                    let layout = chunk.get_layout();
610                    let endpoint_tags = path_endpoint_tags(
611                        &*reader, space_name, edge_types, *direction, vertex_tag,
612                    );
613                    let (start_tag, end_tag) = endpoint_tags
614                        .as_ref()
615                        .map(|(start, end)| (start.as_str(), end.as_str()))
616                        .unzip();
617                    let mut out_rows = Vec::new();
618                    for row in &chunk.rows {
619                        if let Some(rt) = self.runtime.as_ref() {
620                            rt.ensure_not_cancelled()?;
621                        }
622                        let pairs = path_endpoint_pairs(
623                            row,
624                            layout.clone(),
625                            start_vertices,
626                            target_vertices,
627                            None,
628                        )?;
629                        for (src_val, dst_val) in pairs {
630                            let Some(start_vertex) = materialize_path_endpoint(
631                                &*reader, space_name, &src_val, vertex_tag, start_tag,
632                            )?
633                            else {
634                                continue;
635                            };
636                            let Some(end_vertex) = materialize_path_endpoint(
637                                &*reader, space_name, &dst_val, vertex_tag, end_tag,
638                            )?
639                            else {
640                                continue;
641                            };
642                            let cancel_token = self.runtime.as_ref().map(|rt| rt.cancel_token());
643                            let paths = enumerate_all_paths(
644                                &*reader,
645                                &start_vertex,
646                                &end_vertex,
647                                AllPathsConfig {
648                                    space_name,
649                                    vertex_tag,
650                                    edge_types,
651                                    direction: *direction,
652                                    min_depth: *min_depth,
653                                    max_depth: *max_depth,
654                                    path_uniqueness: *path_uniqueness,
655                                    result_cap: limit.unwrap_or(usize::MAX),
656                                },
657                                cancel_token.as_ref(),
658                            )?;
659                            for path in paths.iter().skip(*offset) {
660                                if let Some(rt) = self.runtime.as_ref() {
661                                    rt.ensure_not_cancelled()?;
662                                }
663                                let mut out_row = row.clone();
664                                out_row.push(Value::Path(Box::new(path.clone())));
665                                out_rows.push(out_row);
666                            }
667                        }
668                    }
669                    if out_rows.is_empty() {
670                        continue;
671                    }
672                    let mut new_cols: Vec<ColumnInfo> = col_names
673                        .iter()
674                        .map(|n| ColumnInfo {
675                            name: n.clone(),
676                            data_type: "string".to_string(),
677                        })
678                        .collect();
679                    new_cols.push(ColumnInfo {
680                        name: "path".to_string(),
681                        data_type: "path".to_string(),
682                    });
683                    let _schema = Arc::new(Schema::new(new_cols));
684                    return Ok(Some(DataChunk::new_with_layout(
685                        out_rows,
686                        Arc::clone(&self.output_layout),
687                    )));
688                } else {
689                    return Ok(Some(chunk));
690                }
691            },
692            RecursiveFragmentOperatorKind::Fixpoint { .. } => self.next_fixpoint(),
693        }
694    }
695
696    pub fn stop(&mut self) -> Result<(), QueryError> {
697        Ok(())
698    }
699
700    pub fn bind_runtime(&mut self, runtime: &Arc<ExecutionRuntime>) {
701        let storage = runtime.storage.clone();
702        let space_name = runtime.query_id().space_name.clone().unwrap_or_default();
703        self.runtime = Some(Arc::clone(runtime));
704        match &mut self.kind {
705            RecursiveFragmentOperatorKind::ShortestPath {
706                storage: target_storage,
707                space_name: target_space,
708                ..
709            }
710            | RecursiveFragmentOperatorKind::MultiShortestPath {
711                storage: target_storage,
712                space_name: target_space,
713                ..
714            }
715            | RecursiveFragmentOperatorKind::BFSShortest {
716                storage: target_storage,
717                space_name: target_space,
718                ..
719            }
720            | RecursiveFragmentOperatorKind::AllPaths {
721                storage: target_storage,
722                space_name: target_space,
723                ..
724            } => {
725                *target_storage = storage;
726                *target_space = space_name;
727            }
728            RecursiveFragmentOperatorKind::Fixpoint { .. } => {
729                // Child sub-plans inherit storage from the shared runtime at
730                // materialization time; nothing to rebind here.
731            }
732        }
733    }
734
735    pub fn close(&mut self) -> Result<(), QueryError> {
736        Ok(())
737    }
738
739    /// Serve the converged fixpoint rows chunk by chunk, computing them on
740    /// first call. The (ignored) fragment input is never advanced: fixpoint
741    /// fragments are always built over a `Start` source.
742    fn next_fixpoint(&mut self) -> Result<Option<DataChunk>, QueryError> {
743        let (cte_name, anchor, step, max_iterations, state, output_layout, chunk_size) =
744            match &mut self.kind {
745                RecursiveFragmentOperatorKind::Fixpoint {
746                    cte_name,
747                    anchor,
748                    step,
749                    max_iterations,
750                    state,
751                    ..
752                } => (
753                    cte_name.clone(),
754                    anchor.clone(),
755                    step.clone(),
756                    *max_iterations,
757                    state,
758                    Arc::clone(&self.output_layout),
759                    self.config.chunk_size,
760                ),
761                _ => {
762                    return Err(QueryError::execution(
763                        "next_fixpoint called for a non-fixpoint fragment".to_string(),
764                    ));
765                }
766            };
767        let runtime = self.runtime.clone().ok_or_else(|| {
768            QueryError::execution("Fixpoint execution requires a runtime".to_string())
769        })?;
770
771        if state.output.is_none() {
772            let rows =
773                Self::run_fixpoint(&runtime, &cte_name, &anchor, step.as_ref(), max_iterations)?;
774            state.output = Some(rows);
775        }
776        let output = state.output.as_ref().ok_or_else(|| {
777            QueryError::execution("Fixpoint output missing after computation".to_string())
778        })?;
779        if state.output_pos >= output.len() {
780            return Ok(None);
781        }
782        let end = (state.output_pos + chunk_size).min(output.len());
783        let rows = output[state.output_pos..end].to_vec();
784        state.output_pos = end;
785        Ok(Some(DataChunk::new_with_layout(rows, output_layout)))
786    }
787
788    /// Run anchor once, then iterate the step over successive deltas until
789    /// no new rows appear or `max_iterations` is exceeded.
790    fn run_fixpoint(
791        runtime: &Arc<ExecutionRuntime>,
792        cte_name: &str,
793        anchor: &Arc<crate::executor::streaming::plan::types::PhysicalPlan>,
794        step: Option<&Arc<crate::executor::streaming::plan::types::PhysicalPlan>>,
795        max_iterations: u64,
796    ) -> Result<Vec<Vec<Value>>, QueryError> {
797        let bindings = Arc::new(Self::child_bindings(runtime));
798
799        let mut anchor_exec = Self::materialize_child(runtime, &bindings, anchor)?;
800        let anchor_rows = Self::collect_all(&mut anchor_exec)?;
801
802        let mut seen_keys = std::collections::HashSet::new();
803        let mut seen_rows: Vec<Vec<Value>> = Vec::new();
804        for row in anchor_rows {
805            let key = postcard::to_allocvec(&row).map_err(|e| {
806                QueryError::execution(format!("Fixpoint row serialization failed: {e}"))
807            })?;
808            if seen_keys.insert(key) {
809                seen_rows.push(row);
810            }
811        }
812
813        let Some(step_plan) = step else {
814            // Non-recursive CTE: the anchor output is the whole result.
815            return Ok(seen_rows);
816        };
817
818        let mut step_exec = Self::materialize_child(runtime, &bindings, step_plan)?;
819        let mut delta: Vec<Vec<Value>> = seen_rows.clone();
820        for _ in 0..max_iterations {
821            runtime.ensure_not_cancelled()?;
822            if delta.is_empty() {
823                runtime.clear_cte_table(cte_name);
824                return Ok(seen_rows);
825            }
826            runtime.set_cte_table(cte_name, delta);
827            step_exec.reset()?;
828            let step_rows = Self::collect_all(&mut step_exec)?;
829            let mut next_delta = Vec::new();
830            for row in step_rows {
831                let key = postcard::to_allocvec(&row).map_err(|e| {
832                    QueryError::execution(format!("Fixpoint row serialization failed: {e}"))
833                })?;
834                if seen_keys.insert(key) {
835                    next_delta.push(row.clone());
836                    seen_rows.push(row);
837                }
838            }
839            delta = next_delta;
840        }
841        runtime.clear_cte_table(cte_name);
842        if delta.is_empty() {
843            Ok(seen_rows)
844        } else {
845            Err(QueryError::execution(format!(
846                "Recursive CTE did not converge within {max_iterations} iterations"
847            )))
848        }
849    }
850
851    /// Materialize a child sub-plan, inheriting the parent runtime (storage,
852    /// cancellation, parameters, CTE tables) like expression subqueries do.
853    fn materialize_child(
854        runtime: &Arc<ExecutionRuntime>,
855        bindings: &Arc<crate::executor::streaming::instance::QueryBindings>,
856        plan: &Arc<crate::executor::streaming::plan::types::PhysicalPlan>,
857    ) -> Result<StreamingExecutor, QueryError> {
858        use crate::executor::streaming::plan::materializer::PhysicalPlanMaterializer;
859
860        let (mut exec, _) = PhysicalPlanMaterializer::materialize(plan, bindings)?;
861        exec.set_chunk_size(bindings.chunk_size);
862        exec.set_runtime(Some(runtime.clone()));
863        exec.open()?;
864        Ok(exec)
865    }
866
867    /// Run an executor to completion, collecting every output row.
868    fn collect_all(exec: &mut StreamingExecutor) -> Result<Vec<Vec<Value>>, QueryError> {
869        let mut rows = Vec::new();
870        exec.reset()?;
871        while let Some(mut chunk) = exec.advance()? {
872            chunk.normalize_for_opaque("RecursiveFixpoint");
873            rows.append(&mut chunk.rows);
874        }
875        Ok(rows)
876    }
877
878    /// Child bindings for fixpoint sub-plans, derived from the parent
879    /// runtime. Serial execution (`max_workers = 1`): fixpoint fragments are
880    /// never partitioned.
881    fn child_bindings(
882        runtime: &Arc<ExecutionRuntime>,
883    ) -> crate::executor::streaming::instance::QueryBindings {
884        use crate::executor::streaming::instance::QueryBindings;
885        use crate::executor::streaming::transaction_scope::TransactionScope;
886
887        let identity = runtime.query_id();
888        QueryBindings {
889            parameters: runtime.parameter_values.clone().unwrap_or_default(),
890            session_variables: runtime.session_variable_values.clone().unwrap_or_default(),
891            parameter_frame: None,
892            space_name: identity.space_name.clone(),
893            storage: runtime.storage.clone(),
894            bound_snapshot: None,
895            memory_budget: runtime.memory_budget.clone(),
896            max_workers: 1,
897            chunk_size: crate::executor::base::ExecutionContext::DEFAULT_CHUNK_SIZE,
898            max_buffered_chunks:
899                crate::executor::base::ExecutionContext::DEFAULT_MAX_BUFFERED_CHUNKS,
900            query_id: identity.query_id,
901            cancel_token: Some(runtime.cancel_token()),
902            query_text: None,
903            session_id: identity.session_id.clone(),
904            user_name: None,
905            transaction: TransactionScope::None,
906            shared_scheduler: None,
907            partition_count: 0,
908            arena: runtime.arena.clone(),
909            feedback_history: None,
910            columnar_policy: None,
911            macro_manager: runtime.macro_manager.clone(),
912            type_alias_manager: runtime.type_alias_manager.clone(),
913            sequence_manager: runtime.sequence_manager.clone(),
914            search: runtime.search.clone(),
915        }
916    }
917
918    /// Reset per-run graph-algorithm state and rewind the input. All state
919    /// is derived per input row, so only the input needs rewinding.
920    pub fn reset(&mut self, input: &mut StreamingExecutor) -> Result<bool, QueryError> {
921        input.reset()?;
922        Ok(false)
923    }
924}
linkrs_query/executor/streaming/operators/recursive_fragment_operator.rs

linkrs_query/executor/streaming/operators/
recursive_fragment_operator.rs
