1use serde::{Deserialize, Serialize};
10
11#[derive(
17 Debug,
18 Clone,
19 Copy,
20 PartialEq,
21 Eq,
22 Hash,
23 Serialize,
24 Deserialize,
25 zerompk::ToMessagePack,
26 zerompk::FromMessagePack,
27)]
28#[msgpack(c_enum)]
29pub enum GraphAlgorithm {
30 PageRank,
32 Wcc,
34 LabelPropagation,
36 Lcc,
38 Sssp,
40 Betweenness,
42 Closeness,
44 Harmonic,
46 Degree,
48 Louvain,
50 Triangles,
52 Diameter,
54 KCore,
56}
57
58impl GraphAlgorithm {
59 pub fn name(&self) -> &'static str {
61 match self {
62 Self::PageRank => "pagerank",
63 Self::Wcc => "wcc",
64 Self::LabelPropagation => "label_propagation",
65 Self::Lcc => "lcc",
66 Self::Sssp => "sssp",
67 Self::Betweenness => "betweenness",
68 Self::Closeness => "closeness",
69 Self::Harmonic => "harmonic",
70 Self::Degree => "degree",
71 Self::Louvain => "louvain",
72 Self::Triangles => "triangles",
73 Self::Diameter => "diameter",
74 Self::KCore => "kcore",
75 }
76 }
77
78 pub fn is_iterative(&self) -> bool {
80 matches!(
81 self,
82 Self::PageRank | Self::LabelPropagation | Self::Louvain
83 )
84 }
85
86 pub fn result_schema(&self) -> &'static [(&'static str, AlgoColumnType)] {
91 use AlgoColumnType::*;
92 match self {
93 Self::PageRank => &[("node_id", Text), ("rank", Float64)],
94 Self::Wcc => &[("node_id", Text), ("component_id", Int64)],
95 Self::LabelPropagation => &[("node_id", Text), ("community_id", Int64)],
96 Self::Lcc => &[("node_id", Text), ("coefficient", Float64)],
97 Self::Sssp => &[("node_id", Text), ("distance", Float64)],
98 Self::Betweenness => &[("node_id", Text), ("centrality", Float64)],
99 Self::Closeness => &[("node_id", Text), ("centrality", Float64)],
100 Self::Harmonic => &[("node_id", Text), ("centrality", Float64)],
101 Self::Degree => &[("node_id", Text), ("centrality", Float64)],
102 Self::Louvain => &[
103 ("node_id", Text),
104 ("community_id", Int64),
105 ("modularity", Float64),
106 ],
107 Self::Triangles => &[("node_id", Text), ("triangles", Int64)],
108 Self::Diameter => &[("diameter", Int64), ("radius", Int64)],
109 Self::KCore => &[("node_id", Text), ("coreness", Int64)],
110 }
111 }
112}
113
114#[derive(Debug, Clone, Copy, PartialEq, Eq)]
116pub enum AlgoColumnType {
117 Text,
118 Float64,
119 Int64,
120}
121
122#[derive(
128 Debug,
129 Clone,
130 Default,
131 PartialEq,
132 Serialize,
133 Deserialize,
134 zerompk::ToMessagePack,
135 zerompk::FromMessagePack,
136)]
137pub struct AlgoParams {
138 pub collection: String,
140
141 pub edge_label: Option<String>,
143
144 pub damping: Option<f64>,
146
147 pub max_iterations: Option<usize>,
149
150 pub tolerance: Option<f64>,
152
153 pub source_node: Option<String>,
155
156 pub sample_size: Option<usize>,
159
160 pub direction: Option<String>,
162
163 pub resolution: Option<f64>,
165
166 pub mode: Option<String>,
168
169 #[serde(default)]
176 pub personalization_vector: Option<std::collections::HashMap<String, f64>>,
177}
178
179impl AlgoParams {
180 pub fn damping_factor(&self) -> f64 {
182 self.damping.unwrap_or(0.85).clamp(0.01, 0.99)
183 }
184
185 pub fn iterations(&self, default: usize) -> usize {
192 const ITERATIONS_HARD_CAP: usize = 1_000;
193 self.max_iterations
194 .unwrap_or(default)
195 .clamp(1, ITERATIONS_HARD_CAP)
196 }
197
198 pub fn convergence_tolerance(&self) -> f64 {
200 let t = self.tolerance.unwrap_or(1e-7);
201 if t > 0.0 { t } else { 1e-7 }
202 }
203
204 pub fn louvain_resolution(&self) -> f64 {
206 let r = self.resolution.unwrap_or(1.0);
207 if r > 0.0 { r } else { 1.0 }
208 }
209
210 pub fn personalization_vector(&self) -> Option<&std::collections::HashMap<String, f64>> {
215 self.personalization_vector.as_ref()
216 }
217}
218
219#[cfg(test)]
220mod tests {
221 use super::*;
222 use sonic_rs;
223
224 #[test]
225 fn algorithm_names() {
226 assert_eq!(GraphAlgorithm::PageRank.name(), "pagerank");
227 assert_eq!(GraphAlgorithm::Wcc.name(), "wcc");
228 assert_eq!(GraphAlgorithm::KCore.name(), "kcore");
229 }
230
231 #[test]
232 fn iterative_algorithms() {
233 assert!(GraphAlgorithm::PageRank.is_iterative());
234 assert!(GraphAlgorithm::LabelPropagation.is_iterative());
235 assert!(GraphAlgorithm::Louvain.is_iterative());
236 assert!(!GraphAlgorithm::Wcc.is_iterative());
237 assert!(!GraphAlgorithm::Sssp.is_iterative());
238 }
239
240 #[test]
241 fn result_schema_columns() {
242 let schema = GraphAlgorithm::PageRank.result_schema();
243 assert_eq!(schema.len(), 2);
244 assert_eq!(schema[0], ("node_id", AlgoColumnType::Text));
245 assert_eq!(schema[1], ("rank", AlgoColumnType::Float64));
246 }
247
248 #[test]
249 fn louvain_schema_has_three_columns() {
250 let schema = GraphAlgorithm::Louvain.result_schema();
251 assert_eq!(schema.len(), 3);
252 }
253
254 #[test]
255 fn params_defaults() {
256 let p = AlgoParams::default();
257 assert_eq!(p.damping_factor(), 0.85);
258 assert_eq!(p.iterations(20), 20);
259 assert_eq!(p.convergence_tolerance(), 1e-7);
260 assert_eq!(p.louvain_resolution(), 1.0);
261 }
262
263 #[test]
264 fn params_clamping() {
265 let p = AlgoParams {
266 damping: Some(2.0),
267 tolerance: Some(-1.0),
268 resolution: Some(0.0),
269 ..Default::default()
270 };
271 assert_eq!(p.damping_factor(), 0.99);
272 assert_eq!(p.convergence_tolerance(), 1e-7);
273 assert_eq!(p.louvain_resolution(), 1.0);
274 }
275
276 #[test]
277 fn params_serde_roundtrip() {
278 let p = AlgoParams {
279 collection: "users".into(),
280 damping: Some(0.9),
281 max_iterations: Some(30),
282 source_node: Some("alice".into()),
283 ..Default::default()
284 };
285 let json = sonic_rs::to_string(&p).unwrap();
286 let p2: AlgoParams = sonic_rs::from_str(&json).unwrap();
287 assert_eq!(p2.collection, "users");
288 assert_eq!(p2.damping, Some(0.9));
289 assert_eq!(p2.max_iterations, Some(30));
290 assert_eq!(p2.source_node, Some("alice".into()));
291 }
292}