uqa_graph/
message_passing.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! K-layer message-passing aggregation over graph vertices.
8//!
9//! Each round each vertex pulls feature values from its neighbors
10//! (out-edges and in-edges merged) and combines them with its own
11//! prior feature via a residual sum. Final per-vertex feature is
12//! squashed through the sigmoid to land in `[0, 1]`.
13
14use std::collections::BTreeMap;
15
16use uqa_core::{DocId, Payload, PostingEntry, PostingList, Value, VertexId};
17
18use crate::posting_list::{GraphPayload, GraphPostingList};
19use crate::store::{GraphStore, GraphStoreError, GraphStoreResult};
20
21/// Protects the public graph API from accidentally scheduling billions of
22/// full-graph propagation rounds from an unchecked `u32` input.
23pub const MAX_MESSAGE_PASSING_LAYERS: u32 = 256;
24const MAX_EXACT_F64_INTEGER: u64 = 1_u64 << 53;
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum AggregationKind {
28    Mean,
29    Sum,
30    Max,
31}
32
33pub struct MessagePassing<'a> {
34    pub graph: &'a str,
35    pub k_layers: u32,
36    pub aggregation: AggregationKind,
37    /// Initial feature: `Some(name)` reads `vertex.properties[name]` and
38    /// uses 0.0 only when the property is absent. Present values must be a
39    /// finite float, an exactly representable integer, or a boolean. `None`
40    /// initializes every vertex to 1.0.
41    pub property_name: Option<String>,
42}
43
44impl<'a> MessagePassing<'a> {
45    pub fn new(graph: &'a str) -> Self {
46        Self {
47            graph,
48            k_layers: 2,
49            aggregation: AggregationKind::Mean,
50            property_name: None,
51        }
52    }
53
54    pub fn k_layers(mut self, k: u32) -> Self {
55        self.k_layers = k;
56        self
57    }
58
59    pub fn aggregation(mut self, kind: AggregationKind) -> Self {
60        self.aggregation = kind;
61        self
62    }
63
64    pub fn property_name(mut self, name: impl Into<String>) -> Self {
65        self.property_name = Some(name.into());
66        self
67    }
68
69    pub fn execute<G: GraphStore>(&self, store: &G) -> GraphStoreResult<GraphPostingList> {
70        if self.k_layers > MAX_MESSAGE_PASSING_LAYERS {
71            return Err(GraphStoreError::InvalidQuery(format!(
72                "message-passing layer count {} exceeds limit {MAX_MESSAGE_PASSING_LAYERS}",
73                self.k_layers
74            )));
75        }
76        let vertices: Vec<VertexId> = store.vertex_ids_in_graph(self.graph)?.into_iter().collect();
77        if vertices.is_empty() {
78            return Ok(GraphPostingList::new());
79        }
80
81        let mut features = self.initial_features(store, &vertices)?;
82        for _ in 0..self.k_layers {
83            features = self.propagate_layer(store, &vertices, &features)?;
84        }
85        self.build_result(vertices, &features)
86    }
87
88    fn initial_features<G: GraphStore>(
89        &self,
90        store: &G,
91        vertices: &[VertexId],
92    ) -> GraphStoreResult<BTreeMap<VertexId, f64>> {
93        let mut features = BTreeMap::new();
94        for vid in vertices {
95            let vertex = store.get_vertex(*vid).ok_or_else(|| {
96                GraphStoreError::CorruptGraph(format!(
97                    "message-passing graph {:?} references missing vertex {vid}",
98                    self.graph
99                ))
100            })?;
101            let value = match &self.property_name {
102                Some(key) => match vertex.properties.get(key) {
103                    Some(value) => numeric_feature(value, key, *vid)?,
104                    None => 0.0,
105                },
106                None => 1.0,
107            };
108            features.insert(*vid, value);
109        }
110        Ok(features)
111    }
112
113    fn propagate_layer<G: GraphStore>(
114        &self,
115        store: &G,
116        vertices: &[VertexId],
117        features: &BTreeMap<VertexId, f64>,
118    ) -> GraphStoreResult<BTreeMap<VertexId, f64>> {
119        let mut next = BTreeMap::new();
120        for vid in vertices {
121            let out_edges = store.out_edge_ids(*vid, self.graph)?;
122            let in_edges = store.in_edge_ids(*vid, self.graph)?;
123            let neighbor_count = out_edges.len().checked_add(in_edges.len()).ok_or_else(|| {
124                GraphStoreError::InvalidQuery(format!(
125                    "message-passing neighbor count overflows usize for vertex {vid}"
126                ))
127            })?;
128            let mut neighbor_values: Vec<f64> = Vec::new();
129            neighbor_values
130                .try_reserve_exact(neighbor_count)
131                .map_err(|error| {
132                    GraphStoreError::InvalidQuery(format!(
133                        "cannot allocate {neighbor_count} message-passing neighbors for vertex {vid}: {error}"
134                    ))
135                })?;
136            for eid in out_edges {
137                let edge = store.get_edge(eid).ok_or_else(|| {
138                    GraphStoreError::CorruptGraph(format!("missing message-passing edge {eid}"))
139                })?;
140                let value = features.get(&edge.target_id).ok_or_else(|| {
141                    GraphStoreError::CorruptGraph(format!(
142                        "edge {eid} references vertex {} outside graph {:?}",
143                        edge.target_id, self.graph
144                    ))
145                })?;
146                neighbor_values.push(*value);
147            }
148            for eid in in_edges {
149                let edge = store.get_edge(eid).ok_or_else(|| {
150                    GraphStoreError::CorruptGraph(format!("missing message-passing edge {eid}"))
151                })?;
152                let value = features.get(&edge.source_id).ok_or_else(|| {
153                    GraphStoreError::CorruptGraph(format!(
154                        "edge {eid} references vertex {} outside graph {:?}",
155                        edge.source_id, self.graph
156                    ))
157                })?;
158                neighbor_values.push(*value);
159            }
160            let own_feature = features.get(vid).copied().ok_or_else(|| {
161                GraphStoreError::CorruptGraph(format!(
162                    "message-passing feature state is missing vertex {vid}"
163                ))
164            })?;
165            let combined = if neighbor_values.is_empty() {
166                own_feature
167            } else {
168                let agg = aggregate_features(&neighbor_values, self.aggregation, *vid)?;
169                finite_add(own_feature, agg, *vid)?
170            };
171            next.insert(*vid, combined);
172        }
173        Ok(next)
174    }
175
176    fn build_result(
177        &self,
178        mut vertices: Vec<VertexId>,
179        features: &BTreeMap<VertexId, f64>,
180    ) -> GraphStoreResult<GraphPostingList> {
181        vertices.sort_unstable();
182        let mut entries = Vec::new();
183        entries.try_reserve_exact(vertices.len()).map_err(|error| {
184            GraphStoreError::InvalidQuery(format!(
185                "cannot allocate {} message-passing result entries: {error}",
186                vertices.len()
187            ))
188        })?;
189        let mut graph_payloads: BTreeMap<DocId, GraphPayload> = BTreeMap::new();
190        for vid in &vertices {
191            let feature = features.get(vid).copied().ok_or_else(|| {
192                GraphStoreError::CorruptGraph(format!(
193                    "message-passing final state is missing vertex {vid}"
194                ))
195            })?;
196            let calibrated = sigmoid(feature);
197            entries.push(PostingEntry::new(*vid, Payload::with_score(calibrated)));
198            graph_payloads.insert(
199                *vid,
200                GraphPayload {
201                    subgraph_vertices: vec![*vid],
202                    subgraph_edges: Vec::new(),
203                    graph_name: self.graph.to_string(),
204                    score_override: Some(calibrated),
205                },
206            );
207        }
208        GraphPostingList::try_from_parts(
209            PostingList::from_sorted_unchecked(entries),
210            graph_payloads,
211        )
212        .map_err(Into::into)
213    }
214}
215
216fn numeric_feature(value: &Value, property: &str, vertex_id: VertexId) -> GraphStoreResult<f64> {
217    match value {
218        Value::Float(value) if value.is_finite() => Ok(*value),
219        Value::Float(value) => Err(GraphStoreError::InvalidQuery(format!(
220            "message-passing property {property:?} on vertex {vertex_id} must be finite, got {value}"
221        ))),
222        Value::Int(value) if value.unsigned_abs() <= MAX_EXACT_F64_INTEGER => Ok(*value as f64),
223        Value::Int(value) => Err(GraphStoreError::InvalidQuery(format!(
224            "message-passing property {property:?} integer {value} on vertex {vertex_id} cannot be represented exactly as f64"
225        ))),
226        Value::Bool(value) => Ok(if *value { 1.0 } else { 0.0 }),
227        other => Err(GraphStoreError::InvalidQuery(format!(
228            "message-passing property {property:?} on vertex {vertex_id} must be numeric or boolean, got {other:?}"
229        ))),
230    }
231}
232
233fn aggregate_features(
234    values: &[f64],
235    aggregation: AggregationKind,
236    vertex_id: VertexId,
237) -> GraphStoreResult<f64> {
238    match aggregation {
239        AggregationKind::Max => values.iter().copied().reduce(f64::max).ok_or_else(|| {
240            GraphStoreError::CorruptGraph(format!(
241                "message-passing aggregation for vertex {vertex_id} has no values"
242            ))
243        }),
244        AggregationKind::Sum | AggregationKind::Mean => {
245            let mut total = 0.0;
246            for value in values {
247                total = finite_add(total, *value, vertex_id)?;
248            }
249            if aggregation == AggregationKind::Mean {
250                let count = u64::try_from(values.len()).map_err(|_| {
251                    GraphStoreError::InvalidQuery(format!(
252                        "message-passing neighbor count exceeds u64 for vertex {vertex_id}"
253                    ))
254                })?;
255                if count > MAX_EXACT_F64_INTEGER {
256                    return Err(GraphStoreError::InvalidQuery(format!(
257                        "message-passing neighbor count {count} for vertex {vertex_id} cannot be represented exactly as f64"
258                    )));
259                }
260                total /= count as f64;
261            }
262            Ok(total)
263        }
264    }
265}
266
267fn finite_add(left: f64, right: f64, vertex_id: VertexId) -> GraphStoreResult<f64> {
268    let result = left + right;
269    if !result.is_finite() {
270        return Err(GraphStoreError::InvalidQuery(format!(
271            "message-passing feature accumulation overflowed for vertex {vertex_id}"
272        )));
273    }
274    Ok(result)
275}
276
277fn sigmoid(x: f64) -> f64 {
278    1.0 / (1.0 + (-x).exp())
279}
uqa_graph/message_passing.rs

uqa_graph/
message_passing.rs
