1use std::cmp::Ordering;
33use std::fmt;
34use std::hash::{Hash, Hasher};
35use std::sync::Arc;
36
37use datafusion::arrow::datatypes::{DataType, Field, Schema};
38use datafusion::common::{DFSchema, DFSchemaRef};
39use datafusion::error::Result;
40use datafusion::logical_expr::{Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore};
41use ursa_core::catalog::Call;
42
43use super::functions::{display_call, plan};
44use super::traversal::Traversal;
45
46pub const WEIGHTED_SOURCE: &str = "__ursa_source";
48pub const WEIGHTED_DESTINATION: &str = "__ursa_destination";
50
51pub fn weight_column(i: usize) -> String {
53 format!("__ursa_weight_{i}")
54}
55
56#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
58pub struct Output {
59 pub name: String,
61 pub call: Call,
62 pub field: usize,
64 pub weight: Option<usize>,
67}
68
69#[derive(Debug, Clone)]
71pub struct GraphAlgorithmNode {
72 pub graph: String,
75 pub input: Arc<LogicalPlan>,
77 pub weighted: Option<Arc<LogicalPlan>>,
80 pub vertices: Option<Arc<LogicalPlan>>,
83 pub values: Option<Arc<LogicalPlan>>,
86 pub source: Expr,
89 pub destination: Expr,
90 pub weights: Vec<String>,
92 pub outputs: Vec<Output>,
94 pub per_node: bool,
96 pub id_type: DataType,
100 schema: DFSchemaRef,
101}
102
103impl GraphAlgorithmNode {
104 #[allow(clippy::too_many_arguments)]
105 pub fn try_new(
106 graph: String,
107 input: Arc<LogicalPlan>,
108 weighted: Option<Arc<LogicalPlan>>,
109 vertices: Option<Arc<LogicalPlan>>,
110 values: Option<Arc<LogicalPlan>>,
111 source: Expr,
112 destination: Expr,
113 weights: Vec<String>,
114 outputs: Vec<Output>,
115 per_node: bool,
116 ) -> Result<Self> {
117 let id_type = endpoint_id_type(&graph, &input, &source, &destination)?;
118 check_vertices(&graph, vertices.as_deref())?;
119 if let Some(values) = &values {
120 let columns = values.schema().fields().len();
121 if columns != 2 {
122 return Err(plan(format!(
123 "the node values of graph '{graph}' must be two columns, node ids then \
124 values; the query returns {columns}"
125 )));
126 }
127 }
128 let mut fields = Vec::new();
129 if per_node {
130 fields.push(Field::new("id", id_type.clone(), false));
131 }
132 for o in &outputs {
133 let spec = super::functions::outputs(o.call.spec())[o.field];
134 let ty = match (&o.call, &values) {
136 (Call::NeighborAgg { agg, .. }, Some(values)) => {
137 super::kernels::neighbors_agg_type(*agg, values.schema().field(1).data_type())
138 }
139 _ => spec.ty.data_type(),
140 };
141 fields.push(Field::new(&o.name, ty, spec.nullable));
142 }
143 let schema = Arc::new(DFSchema::try_from(Schema::new(fields))?);
144 Ok(Self {
145 graph,
146 input,
147 weighted,
148 vertices,
149 values,
150 source,
151 destination,
152 weights,
153 outputs,
154 per_node,
155 id_type,
156 schema,
157 })
158 }
159
160 pub fn endpoint_indices(&self) -> Result<(usize, usize)> {
162 endpoint_positions(&self.graph, &self.input, &self.source, &self.destination)
163 }
164
165 pub(crate) fn display_outputs(&self) -> String {
168 let outputs: Vec<String> = self
169 .outputs
170 .iter()
171 .map(|o| {
172 let weight = o.weight.map(|i| self.weights[i].as_str());
173 let call = display_call(&o.call, weight);
174 let spec = o.call.spec();
175 if spec.outputs.len() == 1 && o.name == spec.outputs[0].name {
176 call
177 } else {
178 format!("{}: {call}", o.name)
179 }
180 })
181 .collect();
182 outputs.join(", ")
183 }
184}
185
186pub(crate) fn endpoint_id_type(
189 graph: &str,
190 input: &LogicalPlan,
191 source: &Expr,
192 destination: &Expr,
193) -> Result<DataType> {
194 let source_type = source.get_type(input.schema())?;
195 let destination_type = destination.get_type(input.schema())?;
196 match (id_type(&source_type), id_type(&destination_type)) {
197 (Some(s), Some(d)) if s == d => Ok(s),
198 (Some(_), Some(_)) => Err(plan(format!(
199 "graph '{graph}': the source and destination columns must have one id type \
200 (both integer or both string); got {source_type} and {destination_type}"
201 ))),
202 (None, _) | (_, None) => {
203 let bad = if id_type(&source_type).is_none() {
204 source_type
205 } else {
206 destination_type
207 };
208 Err(plan(format!(
209 "graph '{graph}': node ids must be an integer or string column; got {bad}"
210 )))
211 }
212 }
213}
214
215pub(crate) fn check_vertices(graph: &str, vertices: Option<&LogicalPlan>) -> Result<()> {
217 match vertices.map(|v| v.schema().fields().len()) {
218 None | Some(1) => Ok(()),
219 Some(n) => Err(plan(format!(
220 "graph '{graph}': the vertex keys must be one column; got {n}"
221 ))),
222 }
223}
224
225fn endpoint_positions(
228 graph: &str,
229 input: &LogicalPlan,
230 source: &Expr,
231 destination: &Expr,
232) -> Result<(usize, usize)> {
233 let index = |e: &Expr| -> Result<usize> {
234 match e {
235 Expr::Column(c) => input.schema().index_of_column(c),
236 other => Err(plan(format!(
237 "graph '{graph}': an edge endpoint must be a column; got {other}"
238 ))),
239 }
240 };
241 Ok((index(source)?, index(destination)?))
242}
243
244pub(crate) fn id_type(t: &DataType) -> Option<DataType> {
248 ursa_io::ids::canonical_id_type(t)
249}
250
251type Key<'a> = (&'a String, &'a Vec<Output>, &'a Vec<String>, bool);
253
254impl GraphAlgorithmNode {
255 fn key(&self) -> Key<'_> {
256 (&self.graph, &self.outputs, &self.weights, self.per_node)
257 }
258}
259
260impl PartialEq for GraphAlgorithmNode {
263 fn eq(&self, other: &Self) -> bool {
264 self.key() == other.key()
265 && self.source == other.source
266 && self.destination == other.destination
267 && self.input == other.input
268 && self.weighted == other.weighted
269 && self.vertices == other.vertices
270 && self.values == other.values
271 }
272}
273
274impl Eq for GraphAlgorithmNode {}
275
276impl Hash for GraphAlgorithmNode {
277 fn hash<H: Hasher>(&self, state: &mut H) {
278 self.key().hash(state);
279 self.source.hash(state);
280 self.destination.hash(state);
281 self.input.hash(state);
282 self.weighted.hash(state);
283 self.vertices.hash(state);
284 self.values.hash(state);
285 }
286}
287
288impl PartialOrd for GraphAlgorithmNode {
289 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
290 match self.key().cmp(&other.key()) {
291 Ordering::Equal => {}
292 unequal => return Some(unequal),
293 }
294 match self.source.partial_cmp(&other.source)? {
295 Ordering::Equal => {}
296 unequal => return Some(unequal),
297 }
298 match self.destination.partial_cmp(&other.destination)? {
299 Ordering::Equal => {}
300 unequal => return Some(unequal),
301 }
302 match self.input.partial_cmp(&other.input)? {
303 Ordering::Equal => {}
304 unequal => return Some(unequal),
305 }
306 match self.weighted.partial_cmp(&other.weighted)? {
307 Ordering::Equal => {}
308 unequal => return Some(unequal),
309 }
310 match self.vertices.partial_cmp(&other.vertices)? {
311 Ordering::Equal => {}
312 unequal => return Some(unequal),
313 }
314 self.values.partial_cmp(&other.values)
315 }
316}
317
318impl UserDefinedLogicalNodeCore for GraphAlgorithmNode {
319 fn name(&self) -> &str {
320 "GraphAlgorithm"
321 }
322
323 fn inputs(&self) -> Vec<&LogicalPlan> {
324 let mut inputs = vec![self.input.as_ref()];
325 inputs.extend(self.weighted.as_deref());
326 inputs.extend(self.vertices.as_deref());
327 inputs.extend(self.values.as_deref());
328 inputs
329 }
330
331 fn schema(&self) -> &DFSchemaRef {
332 &self.schema
333 }
334
335 fn expressions(&self) -> Vec<Expr> {
337 vec![self.source.clone(), self.destination.clone()]
338 }
339
340 fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
341 write!(
342 f,
343 "GraphAlgorithm: graph={}, edges=({}, {}), outputs=[{}]",
344 self.graph,
345 self.source,
346 self.destination,
347 self.display_outputs()
348 )
349 }
350
351 fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> Result<Self> {
352 let [source, destination]: [Expr; 2] = exprs
353 .try_into()
354 .map_err(|_| plan("GraphAlgorithm takes two expressions".into()))?;
355 let expected = 1
356 + usize::from(self.weighted.is_some())
357 + usize::from(self.vertices.is_some())
358 + usize::from(self.values.is_some());
359 if inputs.len() != expected {
360 return Err(plan(format!(
361 "GraphAlgorithm takes {expected} inputs; got {}",
362 inputs.len()
363 )));
364 }
365 let mut inputs = inputs.into_iter().map(Arc::new);
366 let input = inputs.next().expect("checked above");
367 let weighted = self.weighted.as_ref().and_then(|_| inputs.next());
368 let vertices = self.vertices.as_ref().and_then(|_| inputs.next());
369 let values = self.values.as_ref().and_then(|_| inputs.next());
370 Self::try_new(
371 self.graph.clone(),
372 input,
373 weighted,
374 vertices,
375 values,
376 source,
377 destination,
378 self.weights.clone(),
379 self.outputs.clone(),
380 self.per_node,
381 )
382 }
383
384 fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option<Vec<Vec<usize>>> {
388 let (source, destination) = self.endpoint_indices().ok()?;
389 let mut needed = vec![source, destination];
390 needed.sort_unstable();
391 needed.dedup();
392 let mut children = vec![needed];
393 if let Some(weighted) = &self.weighted {
394 children.push((0..weighted.schema().fields().len()).collect());
395 }
396 if self.vertices.is_some() {
397 children.push(vec![0]);
398 }
399 if self.values.is_some() {
400 children.push(vec![0, 1]);
401 }
402 Some(children)
403 }
404
405 fn prevent_predicate_push_down_columns(&self) -> std::collections::HashSet<String> {
407 self.schema
408 .fields()
409 .iter()
410 .map(|f| f.name().clone())
411 .collect()
412 }
413}
414
415#[derive(Debug, Clone)]
433pub struct GraphTraversalNode {
434 pub graph: String,
436 pub traversal: Traversal,
437 pub input: Arc<LogicalPlan>,
439 pub weighted: Option<Arc<LogicalPlan>>,
442 pub vertices: Option<Arc<LogicalPlan>>,
444 pub weight: Option<String>,
446 pub nodes: Vec<Arc<LogicalPlan>>,
448 pub written: Vec<String>,
450 pub source: Expr,
453 pub destination: Expr,
454 pub id_type: DataType,
456 schema: DFSchemaRef,
457}
458
459impl GraphTraversalNode {
460 #[allow(clippy::too_many_arguments)]
461 pub fn try_new(
462 graph: String,
463 traversal: Traversal,
464 input: Arc<LogicalPlan>,
465 weighted: Option<Arc<LogicalPlan>>,
466 vertices: Option<Arc<LogicalPlan>>,
467 weight: Option<String>,
468 nodes: Vec<Arc<LogicalPlan>>,
469 written: Vec<String>,
470 source: Expr,
471 destination: Expr,
472 ) -> Result<Self> {
473 let id_type = endpoint_id_type(&graph, &input, &source, &destination)?;
474 check_vertices(&graph, vertices.as_deref())?;
475 let spec = traversal.spec();
476 let expected = spec.node_params().count();
477 if nodes.len() != expected || written.len() != expected {
478 return Err(plan(format!(
479 "{}() takes {expected} node-id inputs; got {}",
480 spec.name,
481 nodes.len()
482 )));
483 }
484 if weighted.is_some() != weight.is_some() {
485 return Err(plan(format!(
486 "{}(): a weighted plan needs its weight expression",
487 spec.name
488 )));
489 }
490 for (p, n) in spec.node_params().zip(&nodes) {
491 let columns = n.schema().fields().len();
492 if columns != 1 {
493 return Err(plan(format!(
494 "{}({} =>) must be one column of node ids; the query returns {columns} \
495 columns",
496 spec.name, p.name
497 )));
498 }
499 }
500 let schema = Arc::new(DFSchema::try_from(spec.schema(&id_type))?);
501 Ok(Self {
502 graph,
503 traversal,
504 input,
505 weighted,
506 vertices,
507 weight,
508 nodes,
509 written,
510 source,
511 destination,
512 id_type,
513 schema,
514 })
515 }
516
517 pub fn endpoint_indices(&self) -> Result<(usize, usize)> {
519 endpoint_positions(&self.graph, &self.input, &self.source, &self.destination)
520 }
521
522 pub(crate) fn display_call(&self) -> String {
524 self.traversal
525 .display(&self.written, self.weight.as_deref())
526 }
527}
528
529type TraversalKey<'a> = (
531 &'a String,
532 &'a Traversal,
533 &'a Option<String>,
534 &'a Vec<String>,
535);
536
537impl GraphTraversalNode {
538 fn key(&self) -> TraversalKey<'_> {
539 (&self.graph, &self.traversal, &self.weight, &self.written)
540 }
541}
542
543impl PartialEq for GraphTraversalNode {
546 fn eq(&self, other: &Self) -> bool {
547 self.key() == other.key()
548 && self.source == other.source
549 && self.destination == other.destination
550 && self.input == other.input
551 && self.weighted == other.weighted
552 && self.vertices == other.vertices
553 && self.nodes == other.nodes
554 }
555}
556
557impl Eq for GraphTraversalNode {}
558
559impl Hash for GraphTraversalNode {
560 fn hash<H: Hasher>(&self, state: &mut H) {
561 self.key().hash(state);
562 self.source.hash(state);
563 self.destination.hash(state);
564 self.input.hash(state);
565 self.weighted.hash(state);
566 self.vertices.hash(state);
567 self.nodes.hash(state);
568 }
569}
570
571impl PartialOrd for GraphTraversalNode {
572 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
573 match self.key().cmp(&other.key()) {
574 Ordering::Equal => {}
575 unequal => return Some(unequal),
576 }
577 match self.source.partial_cmp(&other.source)? {
578 Ordering::Equal => {}
579 unequal => return Some(unequal),
580 }
581 match self.destination.partial_cmp(&other.destination)? {
582 Ordering::Equal => {}
583 unequal => return Some(unequal),
584 }
585 match self.input.partial_cmp(&other.input)? {
586 Ordering::Equal => {}
587 unequal => return Some(unequal),
588 }
589 match self.weighted.partial_cmp(&other.weighted)? {
590 Ordering::Equal => {}
591 unequal => return Some(unequal),
592 }
593 match self.vertices.partial_cmp(&other.vertices)? {
594 Ordering::Equal => {}
595 unequal => return Some(unequal),
596 }
597 self.nodes.partial_cmp(&other.nodes)
598 }
599}
600
601impl UserDefinedLogicalNodeCore for GraphTraversalNode {
602 fn name(&self) -> &str {
603 "GraphTraversal"
604 }
605
606 fn inputs(&self) -> Vec<&LogicalPlan> {
607 let mut inputs = vec![self.input.as_ref()];
608 inputs.extend(self.weighted.as_deref());
609 inputs.extend(self.vertices.as_deref());
610 inputs.extend(self.nodes.iter().map(Arc::as_ref));
611 inputs
612 }
613
614 fn schema(&self) -> &DFSchemaRef {
615 &self.schema
616 }
617
618 fn expressions(&self) -> Vec<Expr> {
620 vec![self.source.clone(), self.destination.clone()]
621 }
622
623 fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
624 write!(
625 f,
626 "GraphTraversal: graph={}, edges=({}, {}), call={}",
627 self.graph,
628 self.source,
629 self.destination,
630 self.display_call()
631 )
632 }
633
634 fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> Result<Self> {
635 let [source, destination]: [Expr; 2] = exprs
636 .try_into()
637 .map_err(|_| plan("GraphTraversal takes two expressions".into()))?;
638 let weighted = usize::from(self.weighted.is_some());
639 let vertices = usize::from(self.vertices.is_some());
640 let expected = 1 + weighted + vertices + self.nodes.len();
641 if inputs.len() != expected {
642 return Err(plan(format!(
643 "GraphTraversal takes {expected} inputs; got {}",
644 inputs.len()
645 )));
646 }
647 let mut inputs = inputs.into_iter().map(Arc::new);
648 let input = inputs.next().expect("checked above");
649 let weighted = (weighted == 1).then(|| inputs.next().expect("checked above"));
650 let vertices = (vertices == 1).then(|| inputs.next().expect("checked above"));
651 Self::try_new(
652 self.graph.clone(),
653 self.traversal,
654 input,
655 weighted,
656 vertices,
657 self.weight.clone(),
658 inputs.collect(),
659 self.written.clone(),
660 source,
661 destination,
662 )
663 }
664
665 fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option<Vec<Vec<usize>>> {
668 let (source, destination) = self.endpoint_indices().ok()?;
669 let mut needed = vec![source, destination];
670 needed.sort_unstable();
671 needed.dedup();
672 let mut children = vec![needed];
673 if let Some(weighted) = &self.weighted {
674 children.push((0..weighted.schema().fields().len()).collect());
675 }
676 if self.vertices.is_some() {
677 children.push(vec![0]);
678 }
679 children.extend(self.nodes.iter().map(|_| vec![0]));
680 Some(children)
681 }
682
683 fn prevent_predicate_push_down_columns(&self) -> std::collections::HashSet<String> {
686 self.schema
687 .fields()
688 .iter()
689 .map(|f| f.name().clone())
690 .collect()
691 }
692}