1use std::sync::Arc;
2
3use parking_lot::RwLock;
4
5use crate::executor::streaming::chunk::DataChunk;
6use crate::executor::streaming::executor::StreamingExecutor;
7use crate::executor::streaming::operators::source_operator::OperatorConfig;
8use crate::executor::streaming::query_registry::CancelToken;
9use crate::executor::streaming::slot::SlotLayout;
10use crate::parser::ast::pattern::PathSemantic;
11use crate::storage::QueryStorage;
12use linkrs_core::error::QueryError;
13use linkrs_core::types::expr::Expression;
14use linkrs_core::{EdgeDirection, Value};
15
16use super::super::runtime::ExecutionRuntime;
17use super::spec::GraphSpec;
18use super::visited_set::VisitedSet;
19
20pub(crate) mod expand;
21mod expand_buffer;
22mod expand_columnar;
23mod expand_dispatch;
24mod expand_frontier;
25mod expand_row_path;
26mod expand_seeds;
27mod subgraph;
28mod traverse;
29
30pub(super) struct ExpandCtx<'a> {
31 pub(super) space_name: &'a str,
32 pub(super) dst_tag: &'a str,
33 pub(super) edge_types: &'a [String],
34 pub(super) direction: EdgeDirection,
35 pub(super) filter_expr: &'a Option<Expression>,
36 pub(super) col_names_template: Vec<String>,
37 pub(super) cancel_token: Option<CancelToken>,
38 pub(super) path_semantic: Option<PathSemantic>,
39 pub(super) edge_required_props: Option<Vec<String>>,
40 pub(super) dst_required_props: Option<Vec<String>>,
41 pub(super) closed_loop: bool,
42 pub(super) skip_rows: bool,
43 pub(super) emit_raw_ids: bool,
44 pub(super) lightweight_source: bool,
45 pub(super) include_empty_paths: bool,
46 pub(super) ttl_active: Option<bool>,
50}
51
52impl ExpandCtx<'_> {
53 pub(super) fn edge_ttl_active(
61 &mut self,
62 reader: &dyn QueryStorage,
63 ) -> Result<bool, QueryError> {
64 if let Some(active) = self.ttl_active {
65 return Ok(active);
66 }
67 let active = reader.edge_types_ttl_active(self.space_name, self.edge_types)?;
68 self.ttl_active = Some(active);
69 Ok(active)
70 }
71}
72
73#[derive(Debug)]
74pub enum GraphOperatorKind {
75 Expand {
76 storage: Option<Arc<RwLock<dyn QueryStorage>>>,
77 space_name: String,
78 dst_tag: String,
79 edge_types: Vec<String>,
80 direction: EdgeDirection,
81 filter_expr: Option<Expression>,
82 },
83 ExpandAll {
84 storage: Option<Arc<RwLock<dyn QueryStorage>>>,
85 space_name: String,
86 dst_tag: String,
87 edge_types: Vec<String>,
88 direction: EdgeDirection,
89 filter_expr: Option<Expression>,
90 col_names: Vec<String>,
91 src_vids: Vec<Value>,
92 step_limit: u32,
93 step_limits: Option<Vec<u32>>,
94 count_only: bool,
95 emit_raw_ids: bool,
96 lightweight_source: bool,
97 path_semantic: Option<PathSemantic>,
98 edge_required_props: Option<Vec<String>>,
99 dst_required_props: Option<Vec<String>>,
100 closed_loop: bool,
101 skip_rows: bool,
102 include_empty_paths: bool,
103 },
104 Traverse {
105 storage: Option<Arc<RwLock<dyn QueryStorage>>>,
106 space_name: String,
107 dst_tag: String,
108 edge_types: Vec<String>,
109 direction: EdgeDirection,
110 min_depth: u32,
111 max_depth: u32,
112 filter_expr: Option<Expression>,
113 visited: VisitedSet,
114 path_semantic: Option<PathSemantic>,
115 },
116 TraverseAll {
117 storage: Option<Arc<RwLock<dyn QueryStorage>>>,
118 space_name: String,
119 dst_tag: String,
120 edge_types: Vec<String>,
121 direction: EdgeDirection,
122 min_depth: u32,
123 max_depth: u32,
124 filter_expr: Option<Expression>,
125 visited: VisitedSet,
126 path_semantic: Option<PathSemantic>,
127 },
128 BiExpand {
129 storage: Option<Arc<RwLock<dyn QueryStorage>>>,
130 space_name: String,
131 dst_tag: String,
132 edge_types: Vec<String>,
133 direction: EdgeDirection,
134 },
135 BiTraverse {
136 storage: Option<Arc<RwLock<dyn QueryStorage>>>,
137 space_name: String,
138 dst_tag: String,
139 edge_types: Vec<String>,
140 direction: EdgeDirection,
141 min_depth: u32,
142 max_depth: u32,
143 visited: VisitedSet,
144 },
145 Subgraph {
146 storage: Option<Arc<RwLock<dyn QueryStorage>>>,
147 space_name: String,
148 dst_tag: String,
149 steps: u32,
150 direction: EdgeDirection,
151 edge_types: Vec<String>,
152 },
153}
154
155#[derive(Debug)]
161pub struct GraphOperator {
162 pub kind: GraphOperatorKind,
163 pub runtime: Option<Arc<ExecutionRuntime>>,
164 pub output_layout: Arc<SlotLayout>,
165 pub config: OperatorConfig,
166}
167
168impl GraphOperator {
169 pub fn bind_runtime(&mut self, runtime: &Arc<ExecutionRuntime>) {
170 let storage = runtime.storage.clone();
171 let space_name = runtime.query_id().space_name.unwrap_or_default();
172 self.runtime = Some(Arc::clone(runtime));
173 match &mut self.kind {
174 GraphOperatorKind::Expand {
175 storage: target_storage,
176 space_name: target_space,
177 ..
178 }
179 | GraphOperatorKind::ExpandAll {
180 storage: target_storage,
181 space_name: target_space,
182 ..
183 }
184 | GraphOperatorKind::Traverse {
185 storage: target_storage,
186 space_name: target_space,
187 ..
188 }
189 | GraphOperatorKind::TraverseAll {
190 storage: target_storage,
191 space_name: target_space,
192 ..
193 }
194 | GraphOperatorKind::BiExpand {
195 storage: target_storage,
196 space_name: target_space,
197 ..
198 }
199 | GraphOperatorKind::BiTraverse {
200 storage: target_storage,
201 space_name: target_space,
202 ..
203 }
204 | GraphOperatorKind::Subgraph {
205 storage: target_storage,
206 space_name: target_space,
207 ..
208 } => {
209 *target_storage = storage;
210 *target_space = space_name;
211 }
212 }
213 }
214
215 pub fn from_spec(
216 spec: &GraphSpec,
217 storage: Option<Arc<RwLock<dyn QueryStorage>>>,
218 space_name: String,
219 output_layout: Arc<SlotLayout>,
220 ) -> Self {
221 let kind = match spec {
222 GraphSpec::Expand {
223 edge_types,
224 direction,
225 filter_expr,
226 dst_tag,
227 ..
228 } => GraphOperatorKind::Expand {
229 storage: storage.clone(),
230 space_name: space_name.clone(),
231 dst_tag: dst_tag.clone(),
232 edge_types: edge_types.clone(),
233 direction: *direction,
234 filter_expr: filter_expr.clone(),
235 },
236 GraphSpec::ExpandAll {
237 edge_types,
238 direction,
239 filter_expr,
240 col_names,
241 src_vids,
242 step_limit,
243 step_limits,
244 count_only,
245 emit_raw_ids,
246 lightweight_source,
247 path_semantic,
248 dst_tag,
249 edge_required_props,
250 dst_required_props,
251 closed_loop,
252 skip_rows,
253 include_empty_paths,
254 } => GraphOperatorKind::ExpandAll {
255 storage: storage.clone(),
256 space_name: space_name.clone(),
257 dst_tag: dst_tag.clone(),
258 edge_types: edge_types.clone(),
259 direction: *direction,
260 filter_expr: filter_expr.clone(),
261 col_names: col_names.clone(),
262 src_vids: src_vids.clone(),
263 step_limit: *step_limit,
264 step_limits: step_limits.clone(),
265 count_only: *count_only,
266 emit_raw_ids: *emit_raw_ids,
267 lightweight_source: *lightweight_source,
268 path_semantic: path_semantic.clone(),
269 edge_required_props: edge_required_props.clone(),
270 dst_required_props: dst_required_props.clone(),
271 closed_loop: *closed_loop,
272 skip_rows: *skip_rows,
273 include_empty_paths: *include_empty_paths,
274 },
275 GraphSpec::Traverse {
276 edge_types,
277 direction,
278 min_depth,
279 max_depth,
280 filter_expr,
281 path_semantic,
282 dst_tag,
283 } => GraphOperatorKind::Traverse {
284 storage: storage.clone(),
285 space_name: space_name.clone(),
286 dst_tag: dst_tag.clone(),
287 edge_types: edge_types.clone(),
288 direction: *direction,
289 min_depth: *min_depth,
290 max_depth: *max_depth,
291 filter_expr: filter_expr.clone(),
292 visited: VisitedSet::new(),
293 path_semantic: path_semantic.clone(),
294 },
295 GraphSpec::BiExpand {
296 edge_types,
297 direction,
298 dst_tag,
299 } => GraphOperatorKind::BiExpand {
300 storage: storage.clone(),
301 space_name: space_name.clone(),
302 dst_tag: dst_tag.clone(),
303 edge_types: edge_types.clone(),
304 direction: *direction,
305 },
306 GraphSpec::BiTraverse {
307 edge_types,
308 direction,
309 min_depth,
310 max_depth,
311 dst_tag,
312 } => GraphOperatorKind::BiTraverse {
313 storage: storage.clone(),
314 space_name: space_name.clone(),
315 dst_tag: dst_tag.clone(),
316 edge_types: edge_types.clone(),
317 direction: *direction,
318 min_depth: *min_depth,
319 max_depth: *max_depth,
320 visited: VisitedSet::new(),
321 },
322 };
323 Self::new(kind, output_layout)
324 }
325
326 pub fn new(kind: GraphOperatorKind, output_layout: Arc<SlotLayout>) -> Self {
327 Self {
328 kind,
329 runtime: None,
330 output_layout,
331 config: OperatorConfig::default(),
332 }
333 }
334
335 pub fn inject_context(
338 &mut self,
339 runtime: Option<&Arc<ExecutionRuntime>>,
340 config: OperatorConfig,
341 ) {
342 if let Some(rt) = runtime {
343 self.runtime = Some(rt.clone());
344 }
345 self.config = config;
346 }
347
348 pub fn open(&mut self, input: &mut StreamingExecutor) -> Result<(), QueryError> {
349 match &mut self.kind {
350 GraphOperatorKind::Expand { .. }
351 | GraphOperatorKind::ExpandAll { .. }
352 | GraphOperatorKind::Traverse { .. }
353 | GraphOperatorKind::TraverseAll { .. }
354 | GraphOperatorKind::BiExpand { .. }
355 | GraphOperatorKind::BiTraverse { .. }
356 | GraphOperatorKind::Subgraph { .. } => {
357 input.open()?;
358 Ok(())
359 }
360 }
361 }
362
363 pub fn next(&mut self, input: &mut StreamingExecutor) -> Result<Option<DataChunk>, QueryError> {
364 let op: &mut GraphOperator = self;
365 if matches!(&op.kind, GraphOperatorKind::Expand { .. }) {
366 return expand::handle(op, input);
367 }
368 if matches!(&op.kind, GraphOperatorKind::ExpandAll { .. }) {
369 return expand::handle_all(op, input);
370 }
371 if matches!(&op.kind, GraphOperatorKind::Traverse { .. }) {
372 return traverse::handle_traverse(op, input);
373 }
374 if matches!(&op.kind, GraphOperatorKind::TraverseAll { .. }) {
375 return traverse::handle_traverse_all(op, input);
376 }
377 if matches!(&op.kind, GraphOperatorKind::BiExpand { .. }) {
378 return traverse::handle_bi_expand(op, input);
379 }
380 if matches!(&op.kind, GraphOperatorKind::BiTraverse { .. }) {
381 return traverse::handle_bi_traverse(op, input);
382 }
383 if matches!(&op.kind, GraphOperatorKind::Subgraph { .. }) {
384 return subgraph::handle(op, input);
385 }
386 unreachable!("graph_operator::next called for an unknown kind")
387 }
388
389 pub fn stop(&mut self) -> Result<(), QueryError> {
390 Ok(())
391 }
392
393 pub fn close(&mut self) -> Result<(), QueryError> {
394 Ok(())
395 }
396
397 pub fn reset(&mut self, input: &mut StreamingExecutor) -> Result<bool, QueryError> {
400 match &mut self.kind {
401 GraphOperatorKind::Traverse { visited, .. }
402 | GraphOperatorKind::TraverseAll { visited, .. }
403 | GraphOperatorKind::BiTraverse { visited, .. } => {
404 *visited = VisitedSet::new();
405 }
406 GraphOperatorKind::Expand { .. }
407 | GraphOperatorKind::ExpandAll { .. }
408 | GraphOperatorKind::BiExpand { .. }
409 | GraphOperatorKind::Subgraph { .. } => {}
410 }
411 input.reset()?;
412 Ok(false)
413 }
414}