ext_algo/
lib.rs

1//! ext-algo: Graph algorithm extension providing PageRank, WCC, and Betweenness Centrality.
2//!
3//! Registers procedures under the `algo` namespace:
4//! - `algo.pageRank(damping, max_iterations, tolerance)` → `{node_id, rank}`
5//! - `algo.wcc()` → `{node_id, component}`
6//! - `algo.betweenness()` → `{node_id, centrality}`
7
8pub mod betweenness;
9pub mod pagerank;
10pub mod wcc;
11
12use std::collections::HashMap;
13
14use kyu_extension::{Extension, ProcColumn, ProcParam, ProcRow, ProcedureSignature};
15use kyu_types::{LogicalType, TypedValue};
16
17/// The graph algorithm extension.
18pub struct AlgoExtension;
19
20impl Extension for AlgoExtension {
21    fn name(&self) -> &str {
22        "algo"
23    }
24
25    fn needs_graph(&self) -> bool {
26        true
27    }
28
29    fn procedures(&self) -> Vec<ProcedureSignature> {
30        vec![
31            ProcedureSignature {
32                name: "pageRank".into(),
33                params: vec![
34                    ProcParam {
35                        name: "damping".into(),
36                        type_desc: "DOUBLE".into(),
37                    },
38                    ProcParam {
39                        name: "max_iterations".into(),
40                        type_desc: "INT64".into(),
41                    },
42                    ProcParam {
43                        name: "tolerance".into(),
44                        type_desc: "DOUBLE".into(),
45                    },
46                ],
47                columns: vec![
48                    ProcColumn {
49                        name: "node_id".into(),
50                        data_type: LogicalType::Int64,
51                    },
52                    ProcColumn {
53                        name: "rank".into(),
54                        data_type: LogicalType::Double,
55                    },
56                ],
57            },
58            ProcedureSignature {
59                name: "wcc".into(),
60                params: vec![],
61                columns: vec![
62                    ProcColumn {
63                        name: "node_id".into(),
64                        data_type: LogicalType::Int64,
65                    },
66                    ProcColumn {
67                        name: "component".into(),
68                        data_type: LogicalType::Int64,
69                    },
70                ],
71            },
72            ProcedureSignature {
73                name: "betweenness".into(),
74                params: vec![],
75                columns: vec![
76                    ProcColumn {
77                        name: "node_id".into(),
78                        data_type: LogicalType::Int64,
79                    },
80                    ProcColumn {
81                        name: "centrality".into(),
82                        data_type: LogicalType::Double,
83                    },
84                ],
85            },
86        ]
87    }
88
89    fn execute(
90        &self,
91        procedure: &str,
92        args: &[String],
93        adjacency: &HashMap<i64, Vec<(i64, f64)>>,
94    ) -> Result<Vec<ProcRow>, String> {
95        match procedure {
96            "pageRank" => {
97                let damping = args
98                    .first()
99                    .and_then(|s| s.parse::<f64>().ok())
100                    .unwrap_or(0.85);
101                let max_iter = args
102                    .get(1)
103                    .and_then(|s| s.parse::<u32>().ok())
104                    .unwrap_or(20);
105                let tolerance = args
106                    .get(2)
107                    .and_then(|s| s.parse::<f64>().ok())
108                    .unwrap_or(1e-6);
109
110                let ranks = pagerank::pagerank(adjacency, damping, max_iter, tolerance);
111                let mut rows: Vec<ProcRow> = ranks
112                    .into_iter()
113                    .map(|(node_id, rank)| {
114                        vec![TypedValue::Int64(node_id), TypedValue::Double(rank)]
115                    })
116                    .collect();
117                rows.sort_by(|a, b| match (&a[1], &b[1]) {
118                    (TypedValue::Double(ra), TypedValue::Double(rb)) => {
119                        rb.partial_cmp(ra).unwrap_or(std::cmp::Ordering::Equal)
120                    }
121                    _ => std::cmp::Ordering::Equal,
122                });
123                Ok(rows)
124            }
125            "wcc" => {
126                let components = wcc::wcc(adjacency);
127                let mut rows: Vec<ProcRow> = components
128                    .into_iter()
129                    .map(|(node_id, component)| {
130                        vec![TypedValue::Int64(node_id), TypedValue::Int64(component)]
131                    })
132                    .collect();
133                rows.sort_by_key(|r| match r[0] {
134                    TypedValue::Int64(id) => id,
135                    _ => 0,
136                });
137                Ok(rows)
138            }
139            "betweenness" => {
140                let scores = betweenness::betweenness_centrality(adjacency);
141                let mut rows: Vec<ProcRow> = scores
142                    .into_iter()
143                    .map(|(node_id, centrality)| {
144                        vec![TypedValue::Int64(node_id), TypedValue::Double(centrality)]
145                    })
146                    .collect();
147                rows.sort_by(|a, b| match (&a[1], &b[1]) {
148                    (TypedValue::Double(ca), TypedValue::Double(cb)) => {
149                        cb.partial_cmp(ca).unwrap_or(std::cmp::Ordering::Equal)
150                    }
151                    _ => std::cmp::Ordering::Equal,
152                });
153                Ok(rows)
154            }
155            _ => Err(format!("unknown procedure: {procedure}")),
156        }
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    fn sample_graph() -> HashMap<i64, Vec<(i64, f64)>> {
165        let mut adj = HashMap::new();
166        // 1 -> 2, 2 -> 3, 3 -> 1 (cycle)
167        adj.insert(1, vec![(2, 1.0)]);
168        adj.insert(2, vec![(3, 1.0)]);
169        adj.insert(3, vec![(1, 1.0)]);
170        adj
171    }
172
173    #[test]
174    fn extension_metadata() {
175        let ext = AlgoExtension;
176        assert_eq!(ext.name(), "algo");
177        assert_eq!(ext.procedures().len(), 3);
178    }
179
180    #[test]
181    fn execute_pagerank() {
182        let ext = AlgoExtension;
183        let adj = sample_graph();
184        let rows = ext
185            .execute(
186                "pageRank",
187                &["0.85".into(), "20".into(), "1e-6".into()],
188                &adj,
189            )
190            .unwrap();
191        assert_eq!(rows.len(), 3);
192        // All rows should have 2 columns: node_id (Int64) and rank (Double).
193        for row in &rows {
194            assert_eq!(row.len(), 2);
195            assert!(matches!(row[0], TypedValue::Int64(_)));
196            assert!(matches!(row[1], TypedValue::Double(_)));
197        }
198    }
199
200    #[test]
201    fn execute_pagerank_defaults() {
202        let ext = AlgoExtension;
203        let adj = sample_graph();
204        let rows = ext.execute("pageRank", &[], &adj).unwrap();
205        assert_eq!(rows.len(), 3);
206    }
207
208    #[test]
209    fn execute_wcc() {
210        let ext = AlgoExtension;
211        let adj = sample_graph();
212        let rows = ext.execute("wcc", &[], &adj).unwrap();
213        assert_eq!(rows.len(), 3);
214        // All nodes in same component.
215        let comp = &rows[0][1];
216        assert!(rows.iter().all(|r| &r[1] == comp));
217    }
218
219    #[test]
220    fn execute_betweenness() {
221        let ext = AlgoExtension;
222        let adj = sample_graph();
223        let rows = ext.execute("betweenness", &[], &adj).unwrap();
224        assert_eq!(rows.len(), 3);
225        for row in &rows {
226            assert!(matches!(row[1], TypedValue::Double(_)));
227        }
228    }
229
230    #[test]
231    fn execute_unknown_procedure() {
232        let ext = AlgoExtension;
233        let adj = sample_graph();
234        let result = ext.execute("nonexistent", &[], &adj);
235        assert!(result.is_err());
236    }
237
238    #[test]
239    fn wcc_two_components() {
240        let ext = AlgoExtension;
241        let mut adj = HashMap::new();
242        adj.insert(1, vec![(2, 1.0)]);
243        adj.insert(2, vec![(1, 1.0)]);
244        adj.insert(10, vec![(11, 1.0)]);
245        adj.insert(11, vec![(10, 1.0)]);
246        let rows = ext.execute("wcc", &[], &adj).unwrap();
247        assert_eq!(rows.len(), 4);
248        // Build a map of node_id -> component.
249        let comp_map: HashMap<i64, i64> = rows
250            .iter()
251            .map(|r| match (&r[0], &r[1]) {
252                (TypedValue::Int64(id), TypedValue::Int64(comp)) => (*id, *comp),
253                _ => panic!("unexpected types"),
254            })
255            .collect();
256        assert_eq!(comp_map[&1], comp_map[&2]);
257        assert_eq!(comp_map[&10], comp_map[&11]);
258        assert_ne!(comp_map[&1], comp_map[&10]);
259    }
260}
ext_algo/lib.rs

ext_algo/
lib.rs
