use std::sync::Arc;
use datafusion::execution::SendableRecordBatchStream;
use futures::future::BoxFuture;
use uni_common::core::id::Vid;
use crate::errors::FnError;
#[derive(Clone, Debug, Default)]
pub struct AlgorithmSignature {
pub output_fields: Vec<arrow_schema::Field>,
pub docs: String,
pub args: Vec<crate::traits::procedure::NamedArgType>,
pub slices: Vec<SliceReq>,
pub df_composable: bool,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SliceReq {
pub slice: smol_str::SmolStr,
pub version: u16,
}
pub const HOST_CAPABILITY_SLICES: &[(&str, u16)] = &[("graph-compute", 1)];
impl AlgorithmSignature {
pub fn check_slices(&self, host_slices: &[(&str, u16)]) -> Result<(), FnError> {
for req in &self.slices {
let satisfied = host_slices
.iter()
.any(|(name, ver)| *name == req.slice.as_str() && *ver >= req.version);
if !satisfied {
return Err(FnError::new(
0x86A,
format!(
"algorithm requires capability slice `{}@{}` the host does not provide",
req.slice, req.version
),
));
}
}
Ok(())
}
pub fn coerce_config_json(&self, config_json: &str) -> Result<String, FnError> {
use crate::traits::scalar::ArgType;
if self.args.is_empty() {
return Ok(config_json.to_owned());
}
let mut provided: Vec<serde_json::Value> = if config_json.trim().is_empty() {
Vec::new()
} else {
serde_json::from_str(config_json)
.map_err(|e| FnError::new(0x86E, format!("bad positional config json: {e}")))?
};
if provided.len() > self.args.len() {
return Err(FnError::new(
0x86E,
format!(
"too many arguments: got {}, expected at most {}",
provided.len(),
self.args.len()
),
));
}
let mut out = Vec::with_capacity(self.args.len());
for (i, arg) in self.args.iter().enumerate() {
match provided.get_mut(i) {
Some(value) => {
let value = std::mem::replace(value, serde_json::Value::Null);
if !matches!(arg.ty, ArgType::CypherValue)
&& !json_matches_argtype(&value, &arg.ty)
{
return Err(FnError::new(
0x86E,
format!("argument `{}` (position {i}) has the wrong type", arg.name),
));
}
out.push(value);
}
None => match &arg.default {
Some(default) => out.push(scalar_default_to_json(default)),
None => {
return Err(FnError::new(
0x86E,
format!("missing required argument `{}` (position {i})", arg.name),
));
}
},
}
}
serde_json::to_string(&out)
.map_err(|e| FnError::new(0x86E, format!("re-encoding coerced config: {e}")))
}
}
fn json_matches_argtype(value: &serde_json::Value, ty: &crate::traits::scalar::ArgType) -> bool {
use arrow_schema::DataType;
use crate::traits::scalar::ArgType;
match ty {
ArgType::CypherValue => true,
ArgType::Vector { .. } => value.is_array(),
ArgType::Variadic(inner) => json_matches_argtype(value, inner),
ArgType::Primitive(dt) => match dt {
DataType::Boolean => value.is_boolean(),
DataType::Utf8 | DataType::LargeUtf8 => value.is_string(),
DataType::Float16 | DataType::Float32 | DataType::Float64 => value.is_number(),
d if d.is_integer() => value.is_i64() || value.is_u64(),
_ => true,
},
}
}
fn scalar_default_to_json(default: &datafusion::scalar::ScalarValue) -> serde_json::Value {
use datafusion::scalar::ScalarValue;
match default {
ScalarValue::Null => serde_json::Value::Null,
ScalarValue::Boolean(Some(b)) => serde_json::Value::Bool(*b),
ScalarValue::Float32(Some(x)) => serde_json::json!(*x),
ScalarValue::Float64(Some(x)) => serde_json::json!(*x),
ScalarValue::Int8(Some(x)) => serde_json::json!(*x),
ScalarValue::Int16(Some(x)) => serde_json::json!(*x),
ScalarValue::Int32(Some(x)) => serde_json::json!(*x),
ScalarValue::Int64(Some(x)) => serde_json::json!(*x),
ScalarValue::UInt8(Some(x)) => serde_json::json!(*x),
ScalarValue::UInt16(Some(x)) => serde_json::json!(*x),
ScalarValue::UInt32(Some(x)) => serde_json::json!(*x),
ScalarValue::UInt64(Some(x)) => serde_json::json!(*x),
ScalarValue::Utf8(Some(s)) | ScalarValue::LargeUtf8(Some(s)) => {
serde_json::Value::String(s.clone())
}
_ => serde_json::Value::Null,
}
}
#[non_exhaustive]
pub struct AlgorithmContext<'a> {
pub config_json: &'a str,
pub host: Option<&'a dyn AlgorithmHost>,
}
impl std::fmt::Debug for AlgorithmContext<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AlgorithmContext")
.field("config_json", &self.config_json)
.field("host_bound", &self.host.is_some())
.finish()
}
}
impl<'a> AlgorithmContext<'a> {
#[must_use]
pub fn new(config_json: &'a str) -> Self {
Self {
config_json,
host: None,
}
}
#[must_use]
pub fn with_host(mut self, host: &'a dyn AlgorithmHost) -> Self {
self.host = Some(host);
self
}
}
pub trait AlgorithmHost: Send + Sync {
fn as_any(&self) -> &dyn std::any::Any;
fn project(
&self,
spec: &GraphProjectionSpec,
) -> BoxFuture<'static, Result<Arc<dyn GraphView>, FnError>> {
let _ = spec;
Box::pin(async {
Err(FnError::new(
0x805,
"AlgorithmHost: project() is not supported by this host",
))
})
}
}
#[derive(Clone, Debug, Default)]
pub struct GraphProjectionSpec {
pub node_labels: Vec<String>,
pub edge_types: Vec<String>,
pub weight_property: Option<String>,
pub include_reverse: bool,
}
pub trait GraphView: Send + Sync {
fn vertex_count(&self) -> usize;
fn edge_count(&self) -> usize;
fn out_neighbors(&self, slot: u32) -> &[u32];
fn out_degree(&self, slot: u32) -> u32;
fn in_neighbors(&self, slot: u32) -> &[u32];
fn in_degree(&self, slot: u32) -> u32;
fn has_reverse(&self) -> bool;
fn out_weight(&self, slot: u32, edge_idx: usize) -> f64;
fn has_weights(&self) -> bool;
fn to_vid(&self, slot: u32) -> Vid;
fn to_slot(&self, vid: Vid) -> Option<u32>;
fn vertices(&self) -> Box<dyn Iterator<Item = (u32, Vid)> + '_>;
}
pub trait AlgorithmProvider: Send + Sync {
fn signature(&self) -> &AlgorithmSignature;
fn run(&self, ctx: AlgorithmContext<'_>) -> Result<SendableRecordBatchStream, FnError>;
}
#[cfg(test)]
mod tests {
use arrow_schema::DataType;
use datafusion::scalar::ScalarValue;
use super::{AlgorithmSignature, HOST_CAPABILITY_SLICES, SliceReq};
use crate::traits::procedure::NamedArgType;
use crate::traits::scalar::ArgType;
fn sig_with(args: Vec<NamedArgType>, slices: Vec<SliceReq>) -> AlgorithmSignature {
AlgorithmSignature {
args,
slices,
..Default::default()
}
}
fn arg(name: &str, ty: ArgType, default: Option<ScalarValue>) -> NamedArgType {
NamedArgType {
name: name.into(),
ty,
default,
doc: String::new(),
}
}
#[test]
fn check_slices_accepts_available_and_rejects_missing() {
let ok = sig_with(
vec![],
vec![SliceReq {
slice: "graph-compute".into(),
version: 1,
}],
);
assert!(ok.check_slices(HOST_CAPABILITY_SLICES).is_ok());
let too_new = sig_with(
vec![],
vec![SliceReq {
slice: "graph-compute".into(),
version: 2,
}],
);
let err = too_new
.check_slices(HOST_CAPABILITY_SLICES)
.expect_err("graph-compute@2 must be refused");
assert_eq!(err.code, 0x86A, "slice mismatch is 0x86A");
let unknown = sig_with(
vec![],
vec![SliceReq {
slice: "tensor-compute".into(),
version: 1,
}],
);
assert_eq!(
unknown
.check_slices(HOST_CAPABILITY_SLICES)
.unwrap_err()
.code,
0x86A
);
assert!(
sig_with(vec![], vec![])
.check_slices(HOST_CAPABILITY_SLICES)
.is_ok()
);
}
#[test]
fn coerce_config_passes_through_when_untyped() {
let s = sig_with(vec![], vec![]);
assert_eq!(s.coerce_config_json("[1, 2, 3]").unwrap(), "[1, 2, 3]");
}
#[test]
fn coerce_config_fills_defaults_and_validates() {
let s = sig_with(
vec![
arg("src", ArgType::CypherValue, None),
arg(
"alpha",
ArgType::Primitive(DataType::Float64),
Some(ScalarValue::Float64(Some(0.85))),
),
],
vec![],
);
let out = s.coerce_config_json("[5]").unwrap();
let arr: Vec<serde_json::Value> = serde_json::from_str(&out).unwrap();
assert_eq!(arr.len(), 2, "the omitted default is appended");
assert_eq!(arr[0], serde_json::json!(5));
assert!((arr[1].as_f64().unwrap() - 0.85).abs() < 1e-12);
let err = s.coerce_config_json("[]").expect_err("src is required");
assert_eq!(err.code, 0x86E);
let err = s
.coerce_config_json(r#"[5, "not-a-number"]"#)
.expect_err("alpha must be numeric");
assert_eq!(err.code, 0x86E);
assert_eq!(s.coerce_config_json("[5, 0.9, 1]").unwrap_err().code, 0x86E);
let arr_src = s.coerce_config_json("[[1, 2, 3], 0.9]").unwrap();
let parsed: Vec<serde_json::Value> = serde_json::from_str(&arr_src).unwrap();
assert!(parsed[0].is_array(), "CypherValue accepts an array");
}
}