ursa_engine/graph/
node.rs

1//! The graph function plan node.
2//!
3//! [`GraphAlgorithmNode`] takes its edges as a child plan: the graph's edge
4//! table, projected to the source and destination columns. Anything that
5//! produces edge rows can sit under it (a scan, a view's filter, a join), and
6//! DataFusion's optimizer works on that child as on any plan: projections and
7//! filters reach the scan.
8//!
9//! A node whose calls take a `weight =>` has a second child: the same edge
10//! table projected to the endpoints and one column per distinct weight
11//! expression ([`WEIGHTED_SOURCE`], [`WEIGHTED_DESTINATION`],
12//! [`weight_column`]). The first child stays exactly the plan an unweighted
13//! call has, so its fingerprint, which keys the index cache, does not depend on
14//! the weights: a weighted and an unweighted query over one graph share one
15//! index. When the node runs with weights it reads only the second child (see
16//! [`super::exec`]).
17//!
18//! A neighbour aggregation (`neighbors_agg`) reads one value per node. Its
19//! node has one more child, last: the values, a plan of two columns, node ids
20//! then values (see [`GraphAlgorithmNode::values`]). Like a traversal's node
21//! ids, it runs when the node runs.
22//!
23//! The node's output is `id` (for per-node functions) plus one column per
24//! [`Output`]: one per algorithm, two per layout (`x` and `y`), and a
25//! statistic's columns, in one row and with no `id`.
26
27use std::cmp::Ordering;
28use std::fmt;
29use std::hash::{Hash, Hasher};
30use std::sync::Arc;
31
32use datafusion::arrow::datatypes::{DataType, Field, Schema};
33use datafusion::common::{DFSchema, DFSchemaRef};
34use datafusion::error::Result;
35use datafusion::logical_expr::{Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore};
36use ursa_core::catalog::Call;
37
38use super::functions::{display_call, plan};
39use super::traversal::Traversal;
40
41/// The source column of a node's weighted edge plan.
42pub const WEIGHTED_SOURCE: &str = "__ursa_source";
43/// The destination column of a node's weighted edge plan.
44pub const WEIGHTED_DESTINATION: &str = "__ursa_destination";
45
46/// The column of a node's weighted edge plan that holds weight expression `i`.
47pub fn weight_column(i: usize) -> String {
48    format!("__ursa_weight_{i}")
49}
50
51/// One output column of a graph node: output `field` of `call`, named `name`.
52#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
53pub struct Output {
54    /// The column's name.
55    pub name: String,
56    pub call: Call,
57    /// Which of the call's outputs (a layout's `x` is 0 and `y` is 1).
58    pub field: usize,
59    /// The weight expression the call reads, as an index into the node's
60    /// weights; `None` is unweighted.
61    pub weight: Option<usize>,
62}
63
64/// Run graph functions over the edges its input produces.
65#[derive(Debug, Clone)]
66pub struct GraphAlgorithmNode {
67    /// The graph's name, for plans and messages. The index is keyed by the
68    /// edges, not by this.
69    pub graph: String,
70    /// The edge rows, projected to the endpoints.
71    pub input: Arc<LogicalPlan>,
72    /// The edge rows projected to the endpoints and the weights, when an
73    /// output reads a weight (see the module docs).
74    pub weighted: Option<Arc<LogicalPlan>>,
75    /// The node values a neighbour aggregation reads: a plan of two columns,
76    /// node ids then values.
77    pub values: Option<Arc<LogicalPlan>>,
78    /// The source and destination of each edge: column references into
79    /// `input`.
80    pub source: Expr,
81    pub destination: Expr,
82    /// The distinct weight expressions, as written, for `EXPLAIN`.
83    pub weights: Vec<String>,
84    /// What the node computes, one value column each.
85    pub outputs: Vec<Output>,
86    /// One row per node with an `id` column, or one row for the whole graph.
87    pub per_node: bool,
88    /// The node id type: the endpoints' type, with every integer type read as
89    /// `Int64` and every string type as `Utf8`, as the Python library reads
90    /// them.
91    pub id_type: DataType,
92    schema: DFSchemaRef,
93}
94
95impl GraphAlgorithmNode {
96    #[allow(clippy::too_many_arguments)]
97    pub fn try_new(
98        graph: String,
99        input: Arc<LogicalPlan>,
100        weighted: Option<Arc<LogicalPlan>>,
101        values: Option<Arc<LogicalPlan>>,
102        source: Expr,
103        destination: Expr,
104        weights: Vec<String>,
105        outputs: Vec<Output>,
106        per_node: bool,
107    ) -> Result<Self> {
108        let id_type = endpoint_id_type(&graph, &input, &source, &destination)?;
109        if let Some(values) = &values {
110            let columns = values.schema().fields().len();
111            if columns != 2 {
112                return Err(plan(format!(
113                    "the node values of graph '{graph}' must be two columns, node ids then \
114                     values; the query returns {columns}"
115                )));
116            }
117        }
118        let mut fields = Vec::new();
119        if per_node {
120            fields.push(Field::new("id", id_type.clone(), false));
121        }
122        for o in &outputs {
123            let spec = o.call.spec().outputs[o.field];
124            fields.push(Field::new(&o.name, spec.ty.data_type(), spec.nullable));
125        }
126        let schema = Arc::new(DFSchema::try_from(Schema::new(fields))?);
127        Ok(Self {
128            graph,
129            input,
130            weighted,
131            values,
132            source,
133            destination,
134            weights,
135            outputs,
136            per_node,
137            id_type,
138            schema,
139        })
140    }
141
142    /// The positions of the source and destination columns in the input.
143    pub fn endpoint_indices(&self) -> Result<(usize, usize)> {
144        endpoint_positions(&self.graph, &self.input, &self.source, &self.destination)
145    }
146
147    /// The outputs as SQL would spell them, for `EXPLAIN`: `name: call`, or
148    /// just the call when the column is named after its function.
149    pub(crate) fn display_outputs(&self) -> String {
150        let outputs: Vec<String> = self
151            .outputs
152            .iter()
153            .map(|o| {
154                let weight = o.weight.map(|i| self.weights[i].as_str());
155                let call = display_call(&o.call, weight);
156                let spec = o.call.spec();
157                if spec.outputs.len() == 1 && o.name == spec.outputs[0].name {
158                    call
159                } else {
160                    format!("{}: {call}", o.name)
161                }
162            })
163            .collect();
164        outputs.join(", ")
165    }
166}
167
168/// The node id type of a graph whose edges `input` holds, with endpoint
169/// columns `source` and `destination`: both integer or both string.
170pub(crate) fn endpoint_id_type(
171    graph: &str,
172    input: &LogicalPlan,
173    source: &Expr,
174    destination: &Expr,
175) -> Result<DataType> {
176    let source_type = source.get_type(input.schema())?;
177    let destination_type = destination.get_type(input.schema())?;
178    match (id_type(&source_type), id_type(&destination_type)) {
179        (Some(s), Some(d)) if s == d => Ok(s),
180        (Some(_), Some(_)) => Err(plan(format!(
181            "graph '{graph}': the source and destination columns must have one id type \
182             (both integer or both string); got {source_type} and {destination_type}"
183        ))),
184        (None, _) | (_, None) => {
185            let bad = if id_type(&source_type).is_none() {
186                source_type
187            } else {
188                destination_type
189            };
190            Err(plan(format!(
191                "graph '{graph}': node ids must be an integer or string column; got {bad}"
192            )))
193        }
194    }
195}
196
197/// The positions of the endpoint columns `source` and `destination` in
198/// `input`.
199fn endpoint_positions(
200    graph: &str,
201    input: &LogicalPlan,
202    source: &Expr,
203    destination: &Expr,
204) -> Result<(usize, usize)> {
205    let index = |e: &Expr| -> Result<usize> {
206        match e {
207            Expr::Column(c) => input.schema().index_of_column(c),
208            other => Err(plan(format!(
209                "graph '{graph}': an edge endpoint must be a column; got {other}"
210            ))),
211        }
212    };
213    Ok((index(source)?, index(destination)?))
214}
215
216/// The node id type a column of `t` gives: integers as `Int64`, strings as
217/// `Utf8`, or `None` if `t` cannot hold node ids. The rule every source shares
218/// (`ursa_io::ids`), so a graph's ids are the Python library's.
219pub(crate) fn id_type(t: &DataType) -> Option<DataType> {
220    ursa_io::ids::canonical_id_type(t)
221}
222
223/// What identifies a node besides its plans and expressions.
224type Key<'a> = (&'a String, &'a Vec<Output>, &'a Vec<String>, bool);
225
226impl GraphAlgorithmNode {
227    fn key(&self) -> Key<'_> {
228        (&self.graph, &self.outputs, &self.weights, self.per_node)
229    }
230}
231
232// The schema is derived from the other fields, so equality, hashing and order
233// leave it out.
234impl PartialEq for GraphAlgorithmNode {
235    fn eq(&self, other: &Self) -> bool {
236        self.key() == other.key()
237            && self.source == other.source
238            && self.destination == other.destination
239            && self.input == other.input
240            && self.weighted == other.weighted
241            && self.values == other.values
242    }
243}
244
245impl Eq for GraphAlgorithmNode {}
246
247impl Hash for GraphAlgorithmNode {
248    fn hash<H: Hasher>(&self, state: &mut H) {
249        self.key().hash(state);
250        self.source.hash(state);
251        self.destination.hash(state);
252        self.input.hash(state);
253        self.weighted.hash(state);
254        self.values.hash(state);
255    }
256}
257
258impl PartialOrd for GraphAlgorithmNode {
259    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
260        match self.key().cmp(&other.key()) {
261            Ordering::Equal => {}
262            unequal => return Some(unequal),
263        }
264        match self.source.partial_cmp(&other.source)? {
265            Ordering::Equal => {}
266            unequal => return Some(unequal),
267        }
268        match self.destination.partial_cmp(&other.destination)? {
269            Ordering::Equal => {}
270            unequal => return Some(unequal),
271        }
272        match self.input.partial_cmp(&other.input)? {
273            Ordering::Equal => {}
274            unequal => return Some(unequal),
275        }
276        match self.weighted.partial_cmp(&other.weighted)? {
277            Ordering::Equal => {}
278            unequal => return Some(unequal),
279        }
280        self.values.partial_cmp(&other.values)
281    }
282}
283
284impl UserDefinedLogicalNodeCore for GraphAlgorithmNode {
285    fn name(&self) -> &str {
286        "GraphAlgorithm"
287    }
288
289    fn inputs(&self) -> Vec<&LogicalPlan> {
290        let mut inputs = vec![self.input.as_ref()];
291        inputs.extend(self.weighted.as_deref());
292        inputs.extend(self.values.as_deref());
293        inputs
294    }
295
296    fn schema(&self) -> &DFSchemaRef {
297        &self.schema
298    }
299
300    /// The endpoint columns, so DataFusion keeps them in the input.
301    fn expressions(&self) -> Vec<Expr> {
302        vec![self.source.clone(), self.destination.clone()]
303    }
304
305    fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
306        write!(
307            f,
308            "GraphAlgorithm: graph={}, edges=({}, {}), outputs=[{}]",
309            self.graph,
310            self.source,
311            self.destination,
312            self.display_outputs()
313        )
314    }
315
316    fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> Result<Self> {
317        let [source, destination]: [Expr; 2] = exprs
318            .try_into()
319            .map_err(|_| plan("GraphAlgorithm takes two expressions".into()))?;
320        let expected =
321            1 + usize::from(self.weighted.is_some()) + usize::from(self.values.is_some());
322        if inputs.len() != expected {
323            return Err(plan(format!(
324                "GraphAlgorithm takes {expected} inputs; got {}",
325                inputs.len()
326            )));
327        }
328        let mut inputs = inputs.into_iter().map(Arc::new);
329        let input = inputs.next().expect("checked above");
330        let weighted = self.weighted.as_ref().and_then(|_| inputs.next());
331        let values = self.values.as_ref().and_then(|_| inputs.next());
332        Self::try_new(
333            self.graph.clone(),
334            input,
335            weighted,
336            values,
337            source,
338            destination,
339            self.weights.clone(),
340            self.outputs.clone(),
341            self.per_node,
342        )
343    }
344
345    /// Every output needs both endpoints and nothing else of the first input,
346    /// every column of the weighted input, and both columns of the values.
347    fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option<Vec<Vec<usize>>> {
348        let (source, destination) = self.endpoint_indices().ok()?;
349        let mut needed = vec![source, destination];
350        needed.sort_unstable();
351        needed.dedup();
352        let mut children = vec![needed];
353        if let Some(weighted) = &self.weighted {
354            children.push((0..weighted.schema().fields().len()).collect());
355        }
356        if self.values.is_some() {
357            children.push(vec![0, 1]);
358        }
359        Some(children)
360    }
361
362    /// Filters on the output are about nodes, not edges: none moves below.
363    fn prevent_predicate_push_down_columns(&self) -> std::collections::HashSet<String> {
364        self.schema
365            .fields()
366            .iter()
367            .map(|f| f.name().clone())
368            .collect()
369    }
370}
371
372/// Run a traversal (`hop`, `shortest_path`, `random_walk`) over the edges its
373/// first input produces.
374///
375/// Its inputs, in order:
376///
377/// 1. the edge rows projected to the endpoints, exactly as a
378///    [`GraphAlgorithmNode`]'s first input, so the two share one index;
379/// 2. for a weighted `shortest_path`, the weighted edge plan, as a
380///    [`GraphAlgorithmNode`]'s second input (with one weight column);
381/// 3. one plan per node-id parameter (`seeds`; `source` and `target`;
382///    `start`), each of one column: a `VALUES` list for a literal, or the
383///    planned subquery. They run when the node runs, never while the query is
384///    planned, and their values are cast to the graph's id type then.
385///
386/// The output is the traversal's own rows (see [`super::traversal`]).
387#[derive(Debug, Clone)]
388pub struct GraphTraversalNode {
389    /// The graph's name, for plans and messages.
390    pub graph: String,
391    pub traversal: Traversal,
392    /// The edge rows, projected to the endpoints.
393    pub input: Arc<LogicalPlan>,
394    /// The edge rows projected to the endpoints and the weight, for a
395    /// weighted shortest path.
396    pub weighted: Option<Arc<LogicalPlan>>,
397    /// The weight expression, as written, for `EXPLAIN`.
398    pub weight: Option<String>,
399    /// One plan per node-id parameter, in the spec's order.
400    pub nodes: Vec<Arc<LogicalPlan>>,
401    /// Each node-id argument as written, for `EXPLAIN`.
402    pub written: Vec<String>,
403    /// The source and destination of each edge: column references into
404    /// `input`.
405    pub source: Expr,
406    pub destination: Expr,
407    /// The node id type (see [`GraphAlgorithmNode::id_type`]).
408    pub id_type: DataType,
409    schema: DFSchemaRef,
410}
411
412impl GraphTraversalNode {
413    #[allow(clippy::too_many_arguments)]
414    pub fn try_new(
415        graph: String,
416        traversal: Traversal,
417        input: Arc<LogicalPlan>,
418        weighted: Option<Arc<LogicalPlan>>,
419        weight: Option<String>,
420        nodes: Vec<Arc<LogicalPlan>>,
421        written: Vec<String>,
422        source: Expr,
423        destination: Expr,
424    ) -> Result<Self> {
425        let id_type = endpoint_id_type(&graph, &input, &source, &destination)?;
426        let spec = traversal.spec();
427        let expected = spec.node_params().count();
428        if nodes.len() != expected || written.len() != expected {
429            return Err(plan(format!(
430                "{}() takes {expected} node-id inputs; got {}",
431                spec.name,
432                nodes.len()
433            )));
434        }
435        if weighted.is_some() != weight.is_some() {
436            return Err(plan(format!(
437                "{}(): a weighted plan needs its weight expression",
438                spec.name
439            )));
440        }
441        for (p, n) in spec.node_params().zip(&nodes) {
442            let columns = n.schema().fields().len();
443            if columns != 1 {
444                return Err(plan(format!(
445                    "{}({} =>) must be one column of node ids; the query returns {columns} \
446                     columns",
447                    spec.name, p.name
448                )));
449            }
450        }
451        let schema = Arc::new(DFSchema::try_from(spec.schema(&id_type))?);
452        Ok(Self {
453            graph,
454            traversal,
455            input,
456            weighted,
457            weight,
458            nodes,
459            written,
460            source,
461            destination,
462            id_type,
463            schema,
464        })
465    }
466
467    /// The positions of the source and destination columns in the input.
468    pub fn endpoint_indices(&self) -> Result<(usize, usize)> {
469        endpoint_positions(&self.graph, &self.input, &self.source, &self.destination)
470    }
471
472    /// The call as SQL would spell it, for `EXPLAIN`.
473    pub(crate) fn display_call(&self) -> String {
474        self.traversal
475            .display(&self.written, self.weight.as_deref())
476    }
477}
478
479/// What identifies a traversal node besides its plans and expressions.
480type TraversalKey<'a> = (
481    &'a String,
482    &'a Traversal,
483    &'a Option<String>,
484    &'a Vec<String>,
485);
486
487impl GraphTraversalNode {
488    fn key(&self) -> TraversalKey<'_> {
489        (&self.graph, &self.traversal, &self.weight, &self.written)
490    }
491}
492
493// The schema and id type are derived from the other fields, so equality,
494// hashing and order leave them out.
495impl PartialEq for GraphTraversalNode {
496    fn eq(&self, other: &Self) -> bool {
497        self.key() == other.key()
498            && self.source == other.source
499            && self.destination == other.destination
500            && self.input == other.input
501            && self.weighted == other.weighted
502            && self.nodes == other.nodes
503    }
504}
505
506impl Eq for GraphTraversalNode {}
507
508impl Hash for GraphTraversalNode {
509    fn hash<H: Hasher>(&self, state: &mut H) {
510        self.key().hash(state);
511        self.source.hash(state);
512        self.destination.hash(state);
513        self.input.hash(state);
514        self.weighted.hash(state);
515        self.nodes.hash(state);
516    }
517}
518
519impl PartialOrd for GraphTraversalNode {
520    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
521        match self.key().cmp(&other.key()) {
522            Ordering::Equal => {}
523            unequal => return Some(unequal),
524        }
525        match self.source.partial_cmp(&other.source)? {
526            Ordering::Equal => {}
527            unequal => return Some(unequal),
528        }
529        match self.destination.partial_cmp(&other.destination)? {
530            Ordering::Equal => {}
531            unequal => return Some(unequal),
532        }
533        match self.input.partial_cmp(&other.input)? {
534            Ordering::Equal => {}
535            unequal => return Some(unequal),
536        }
537        match self.weighted.partial_cmp(&other.weighted)? {
538            Ordering::Equal => {}
539            unequal => return Some(unequal),
540        }
541        self.nodes.partial_cmp(&other.nodes)
542    }
543}
544
545impl UserDefinedLogicalNodeCore for GraphTraversalNode {
546    fn name(&self) -> &str {
547        "GraphTraversal"
548    }
549
550    fn inputs(&self) -> Vec<&LogicalPlan> {
551        let mut inputs = vec![self.input.as_ref()];
552        inputs.extend(self.weighted.as_deref());
553        inputs.extend(self.nodes.iter().map(Arc::as_ref));
554        inputs
555    }
556
557    fn schema(&self) -> &DFSchemaRef {
558        &self.schema
559    }
560
561    /// The endpoint columns, so DataFusion keeps them in the input.
562    fn expressions(&self) -> Vec<Expr> {
563        vec![self.source.clone(), self.destination.clone()]
564    }
565
566    fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
567        write!(
568            f,
569            "GraphTraversal: graph={}, edges=({}, {}), call={}",
570            self.graph,
571            self.source,
572            self.destination,
573            self.display_call()
574        )
575    }
576
577    fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> Result<Self> {
578        let [source, destination]: [Expr; 2] = exprs
579            .try_into()
580            .map_err(|_| plan("GraphTraversal takes two expressions".into()))?;
581        let weighted = usize::from(self.weighted.is_some());
582        let expected = 1 + weighted + self.nodes.len();
583        if inputs.len() != expected {
584            return Err(plan(format!(
585                "GraphTraversal takes {expected} inputs; got {}",
586                inputs.len()
587            )));
588        }
589        let mut inputs = inputs.into_iter().map(Arc::new);
590        let input = inputs.next().expect("checked above");
591        let weighted = (weighted == 1).then(|| inputs.next().expect("checked above"));
592        Self::try_new(
593            self.graph.clone(),
594            self.traversal,
595            input,
596            weighted,
597            self.weight.clone(),
598            inputs.collect(),
599            self.written.clone(),
600            source,
601            destination,
602        )
603    }
604
605    /// Both endpoints of the first input, every column of the weighted input,
606    /// and the one column of each node-id input.
607    fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option<Vec<Vec<usize>>> {
608        let (source, destination) = self.endpoint_indices().ok()?;
609        let mut needed = vec![source, destination];
610        needed.sort_unstable();
611        needed.dedup();
612        let mut children = vec![needed];
613        if let Some(weighted) = &self.weighted {
614            children.push((0..weighted.schema().fields().len()).collect());
615        }
616        children.extend(self.nodes.iter().map(|_| vec![0]));
617        Some(children)
618    }
619
620    /// Filters on the output are about the traversal's rows, not the edges or
621    /// the node ids: none moves below.
622    fn prevent_predicate_push_down_columns(&self) -> std::collections::HashSet<String> {
623        self.schema
624            .fields()
625            .iter()
626            .map(|f| f.name().clone())
627            .collect()
628    }
629}
ursa_engine/graph/node.rs

ursa_engine/graph/
node.rs
