uni_algo/algo/
procedures.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2024-2026 Dragonscale Team
3
4//! Algorithm procedure interface for Cypher integration.
5//!
6//! Procedures are registered with `AlgorithmRegistry` and can be invoked
7//! via `CALL algo.name(...)` in Cypher queries.
8
9use anyhow::{Result, anyhow};
10use futures::Stream;
11use serde_json::Value;
12use std::pin::Pin;
13
14/// Procedure signature for documentation and validation.
15#[derive(Debug, Clone)]
16pub struct ProcedureSignature {
17    /// Required arguments: (name, type)
18    pub args: Vec<(&'static str, ValueType)>,
19    /// Optional arguments: (name, type, default)
20    pub optional_args: Vec<(&'static str, ValueType, Value)>,
21    /// Output columns: (name, type)
22    pub yields: Vec<(&'static str, ValueType)>,
23}
24
25impl ProcedureSignature {
26    /// Validate arguments against signature and fill defaults for optional args.
27    pub fn validate_args(&self, mut args: Vec<Value>) -> Result<Vec<Value>> {
28        let req_count = self.args.len();
29        let total_count = req_count + self.optional_args.len();
30
31        if args.len() < req_count {
32            return Err(anyhow!(
33                "Too few arguments. Expected at least {}, got {}",
34                req_count,
35                args.len()
36            ));
37        }
38
39        if args.len() > total_count {
40            return Err(anyhow!(
41                "Too many arguments. Expected at most {}, got {}",
42                total_count,
43                args.len()
44            ));
45        }
46
47        // Validate required args
48        for (i, (name, ty)) in self.args.iter().enumerate() {
49            if !ty.matches(&args[i]) {
50                return Err(anyhow!(
51                    "Invalid type for argument '{}'. Expected {:?}, got {:?}",
52                    name,
53                    ty,
54                    args[i]
55                ));
56            }
57        }
58
59        // Validate provided optional args and fill defaults for missing ones
60        for i in 0..self.optional_args.len() {
61            let idx = req_count + i;
62            let (name, ty, default) = &self.optional_args[i];
63
64            if idx < args.len() {
65                if !ty.matches(&args[idx]) {
66                    return Err(anyhow!(
67                        "Invalid type for optional argument '{}'. Expected {:?}, got {:?}",
68                        name,
69                        ty,
70                        args[idx]
71                    ));
72                }
73            } else {
74                args.push(default.clone());
75            }
76        }
77
78        Ok(args)
79    }
80}
81
82/// Value types for procedure signatures.
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
84pub enum ValueType {
85    Int,
86    Float,
87    String,
88    Bool,
89    List,
90    Map,
91    Node,
92    Relationship,
93    Path,
94    Any,
95}
96
97impl ValueType {
98    pub fn matches(&self, val: &Value) -> bool {
99        match self {
100            ValueType::Int => val.is_i64() || val.is_u64(),
101            ValueType::Float => val.is_f64() || val.is_i64() || val.is_u64(),
102            ValueType::String => val.is_string(),
103            ValueType::Bool => val.is_boolean(),
104            ValueType::List => val.is_array(),
105            ValueType::Map => val.is_object(),
106            ValueType::Node => val.is_string() || val.is_u64(), // VID string or u64
107            ValueType::Relationship => val.is_u64() || val.is_object(),
108            ValueType::Path => val.is_object(), // Path struct
109            ValueType::Any => true,
110        }
111    }
112}
113
114/// Result row from algorithm execution.
115#[derive(Debug, Clone)]
116pub struct AlgoResultRow {
117    /// Column values in order matching `yields`.
118    pub values: Vec<Value>,
119}
120
121/// Trait for algorithm procedures.
122///
123/// Implement this to expose an algorithm via `CALL algo.name(...)`.
124pub trait AlgoProcedure: Send + Sync {
125    /// Procedure name (e.g., "algo.pageRank").
126    fn name(&self) -> &str;
127
128    /// Procedure signature for validation and documentation.
129    fn signature(&self) -> ProcedureSignature;
130
131    /// Execute the procedure with given arguments.
132    ///
133    /// Returns a stream of result rows.
134    fn execute(
135        &self,
136        ctx: AlgoContext,
137        args: Vec<Value>,
138    ) -> Pin<Box<dyn Stream<Item = Result<AlgoResultRow>> + Send + 'static>>;
139}
140
141use std::sync::Arc;
142use uni_store::runtime::L0Manager;
143use uni_store::storage::manager::StorageManager;
144
145/// Execution context for algorithm procedures.
146pub struct AlgoContext {
147    pub storage: Arc<StorageManager>,
148    /// L0 manager for scanning in-memory vertices not yet flushed.
149    pub l0_manager: Option<Arc<L0Manager>>,
150}
151
152impl AlgoContext {
153    /// Create a new algorithm context.
154    pub fn new(storage: Arc<StorageManager>, l0_manager: Option<Arc<L0Manager>>) -> Self {
155        Self {
156            storage,
157            l0_manager,
158        }
159    }
160}
161
162// Placeholder procedure implementations will be added in Phase 3.3
uni_algo/algo/procedures.rs

uni_algo/algo/
procedures.rs
