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), ("graph-arena", 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 node_properties: Vec<String>,
pub edge_properties: Vec<String>,
pub project_all: bool,
}
impl GraphProjectionSpec {
#[must_use]
pub fn from_config_object(cfg: &serde_json::Map<String, serde_json::Value>) -> Self {
fn string_array(v: &serde_json::Value) -> Vec<String> {
v.as_array()
.map(|arr| {
arr.iter()
.filter_map(|s| s.as_str().map(str::to_owned))
.collect()
})
.unwrap_or_default()
}
let node_labels = cfg.get("nodeLabels").map(string_array).unwrap_or_default();
let edge_types = cfg
.get("edgeTypes")
.or_else(|| cfg.get("relationshipTypes"))
.map(string_array)
.unwrap_or_default();
let weight_property = cfg
.get("weightProperty")
.and_then(serde_json::Value::as_str)
.map(str::to_owned);
let include_reverse = cfg
.get("includeReverse")
.and_then(serde_json::Value::as_bool)
.unwrap_or(true);
let node_properties = cfg
.get("nodeProperties")
.map(string_array)
.unwrap_or_default();
let edge_properties = cfg
.get("edgeProperties")
.map(string_array)
.unwrap_or_default();
let project_all = cfg
.get("projectAll")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
Self {
node_labels,
edge_types,
weight_property,
include_reverse,
node_properties,
edge_properties,
project_all,
}
}
pub const CONFIG_KEYS: &'static [&'static str] = &[
"nodeLabels",
"edgeTypes",
"relationshipTypes",
"weightProperty",
"includeReverse",
"nodeProperties",
"edgeProperties",
"projectAll",
"nodeQuery",
"edgeQuery",
"weightColumn",
"name",
"scopes",
];
pub const QUERY_CONFIG_KEYS: &'static [&'static str] = &["nodeQuery", "edgeQuery", "name"];
#[must_use]
pub fn is_query_graph_ref(cfg: &serde_json::Map<String, serde_json::Value>) -> bool {
Self::QUERY_CONFIG_KEYS.iter().any(|k| cfg.contains_key(*k))
}
pub fn scopes_from_config_object(
cfg: &serde_json::Map<String, serde_json::Value>,
) -> Result<Vec<GraphScopeSpec>, String> {
let Some(raw) = cfg.get("scopes") else {
return Ok(Vec::new());
};
let map = raw
.as_object()
.ok_or_else(|| "`scopes` must be an object of {name: projection-config}".to_string())?;
let mut out = Vec::with_capacity(map.len());
for (name, value) in map {
if name.is_empty() {
return Err("a scope name must not be empty".to_string());
}
if name == "graph" {
return Err(
"`graph` is not a valid scope name: it is the primary projection, \
reached with `gc.graph()` rather than `gc.graph_named(..)`"
.to_string(),
);
}
let obj = value.as_object().ok_or_else(|| {
format!("scope `{name}` must be a projection-config object, got {value}")
})?;
let graph_ref = Self::is_query_graph_ref(obj).then(|| value.clone());
out.push(GraphScopeSpec {
name: name.clone(),
spec: Self::from_config_object(obj),
graph_ref,
});
}
Ok(out)
}
pub fn reject_scopes(
cfg: &serde_json::Map<String, serde_json::Value>,
algorithm: &str,
) -> Result<(), FnError> {
let has_scopes = cfg
.get("scopes")
.and_then(serde_json::Value::as_object)
.is_some_and(|m| !m.is_empty());
if !has_scopes {
return Ok(());
}
Err(FnError::new(
0x86E,
format!(
"{algorithm} does not take named `scopes`: it runs a fixed algorithm over \
one projection. Named scopes are a guest-authored-algorithm feature -- \
the guest is what decides which scope to read."
),
))
}
#[must_use]
pub fn take_config_from_args(
args: &mut Vec<serde_json::Value>,
) -> Option<serde_json::Map<String, serde_json::Value>> {
let is_config = args
.last()
.and_then(serde_json::Value::as_object)
.is_some_and(|o| Self::CONFIG_KEYS.iter().any(|k| o.contains_key(*k)));
if !is_config {
return None;
}
match args.pop() {
Some(serde_json::Value::Object(cfg)) => Some(cfg),
_ => None,
}
}
#[must_use]
pub fn take_from_args(args: &mut Vec<serde_json::Value>) -> Option<Self> {
Self::take_config_from_args(args).map(|cfg| Self::from_config_object(&cfg))
}
}
#[derive(Clone, Debug)]
pub struct GraphScopeSpec {
pub name: String,
pub spec: GraphProjectionSpec,
pub graph_ref: Option<serde_json::Value>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ProjectionKnob {
NodeLabels,
EdgeTypes,
WeightProperty,
IncludeReverse,
NodeProperties,
EdgeProperties,
CypherNodeQuery,
CypherEdgeQuery,
CypherWeightColumn,
NamedGraph,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum KnobReach {
GuestNative,
GuestQuerySeam,
HostOnly,
}
impl ProjectionKnob {
pub const ALL: &'static [ProjectionKnob] = &[
ProjectionKnob::NodeLabels,
ProjectionKnob::EdgeTypes,
ProjectionKnob::WeightProperty,
ProjectionKnob::IncludeReverse,
ProjectionKnob::NodeProperties,
ProjectionKnob::EdgeProperties,
ProjectionKnob::CypherNodeQuery,
ProjectionKnob::CypherEdgeQuery,
ProjectionKnob::CypherWeightColumn,
ProjectionKnob::NamedGraph,
];
#[must_use]
pub fn config_key(self) -> &'static str {
match self {
ProjectionKnob::NodeLabels => "nodeLabels",
ProjectionKnob::EdgeTypes => "edgeTypes",
ProjectionKnob::WeightProperty => "weightProperty",
ProjectionKnob::IncludeReverse => "includeReverse",
ProjectionKnob::NodeProperties => "nodeProperties",
ProjectionKnob::EdgeProperties => "edgeProperties",
ProjectionKnob::CypherNodeQuery => "nodeQuery",
ProjectionKnob::CypherEdgeQuery => "edgeQuery",
ProjectionKnob::CypherWeightColumn => "weightColumn",
ProjectionKnob::NamedGraph => "name",
}
}
#[must_use]
pub fn reach(self) -> KnobReach {
match self {
ProjectionKnob::NodeLabels
| ProjectionKnob::EdgeTypes
| ProjectionKnob::WeightProperty
| ProjectionKnob::IncludeReverse
| ProjectionKnob::NodeProperties
| ProjectionKnob::EdgeProperties => KnobReach::GuestNative,
ProjectionKnob::CypherNodeQuery
| ProjectionKnob::CypherEdgeQuery
| ProjectionKnob::CypherWeightColumn
| ProjectionKnob::NamedGraph => KnobReach::GuestQuerySeam,
}
}
}
#[cfg(test)]
mod projection_knob_contract {
use super::{GraphProjectionSpec, KnobReach, ProjectionKnob};
#[test]
fn every_projection_knob_is_classified_and_keyed() {
let mut keys = std::collections::HashSet::new();
for knob in ProjectionKnob::ALL {
let key = knob.config_key();
assert!(keys.insert(key), "duplicate projection config key {key}");
assert!(
GraphProjectionSpec::CONFIG_KEYS.contains(&key),
"{key} missing from GraphProjectionSpec::CONFIG_KEYS"
);
let _ = knob.reach(); }
}
#[test]
fn guest_native_knobs_round_trip_through_the_shared_parser() {
for &knob in ProjectionKnob::ALL {
if knob.reach() != KnobReach::GuestNative {
continue;
}
let key = knob.config_key();
let sample = match knob {
ProjectionKnob::IncludeReverse => serde_json::json!(false),
ProjectionKnob::WeightProperty => serde_json::json!("w"),
_ => serde_json::json!(["X"]),
};
let mut obj = serde_json::Map::new();
obj.insert(key.to_string(), sample);
let spec = GraphProjectionSpec::from_config_object(&obj);
let honored = match knob {
ProjectionKnob::NodeLabels => spec.node_labels == ["X"],
ProjectionKnob::EdgeTypes => spec.edge_types == ["X"],
ProjectionKnob::WeightProperty => spec.weight_property.as_deref() == Some("w"),
ProjectionKnob::IncludeReverse => !spec.include_reverse,
ProjectionKnob::NodeProperties => spec.node_properties == ["X"],
ProjectionKnob::EdgeProperties => spec.edge_properties == ["X"],
_ => unreachable!("only GuestNative knobs reach here"),
};
assert!(honored, "from_config_object did not honor `{key}`");
}
}
}
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(),
}
}
fn cfg(json: &str) -> serde_json::Map<String, serde_json::Value> {
match serde_json::from_str(json).expect("valid json") {
serde_json::Value::Object(o) => o,
other => panic!("expected an object, got {other}"),
}
}
#[test]
fn a_scopes_only_object_is_recognised_as_the_projection_config() {
let mut args: Vec<serde_json::Value> =
serde_json::from_str(r#"[1, {"scopes": {"agg": {"nodeLabels": ["N"]}}}]"#)
.expect("valid json");
let spec = super::GraphProjectionSpec::take_from_args(&mut args);
assert!(
spec.is_some(),
"the trailing object must be taken as config"
);
assert_eq!(args.len(), 1, "only the guest's own arg may remain");
}
#[test]
fn scopes_parse_per_scope_and_keep_their_mode() {
let scopes = super::GraphProjectionSpec::scopes_from_config_object(&cfg(r#"{"scopes": {
"agg": {"nodeLabels": ["Cell"], "edgeTypes": ["AGG"], "weightProperty": "w"},
"flow": {"nodeQuery": "MATCH (c) RETURN id(c) AS id"}
}}"#))
.expect("well-formed scopes");
assert_eq!(scopes.len(), 2);
let agg = scopes.iter().find(|s| s.name == "agg").expect("agg");
assert_eq!(agg.spec.node_labels, vec!["Cell".to_string()]);
assert_eq!(agg.spec.weight_property.as_deref(), Some("w"));
assert!(
agg.graph_ref.is_none(),
"a Native scope must not be routed to the resolver"
);
let flow = scopes.iter().find(|s| s.name == "flow").expect("flow");
assert!(
flow.graph_ref.is_some(),
"a Cypher scope must reach the resolver verbatim"
);
}
#[test]
fn an_absent_scopes_key_yields_no_scopes() {
let scopes =
super::GraphProjectionSpec::scopes_from_config_object(&cfg(r#"{"nodeLabels": ["N"]}"#))
.expect("no scopes is fine");
assert!(scopes.is_empty());
}
#[test]
fn a_malformed_scopes_map_is_rejected_with_the_offending_name() {
for (json, needle) in [
(r#"{"scopes": ["agg"]}"#, "must be an object"),
(r#"{"scopes": {"agg": 7}}"#, "agg"),
(r#"{"scopes": {"": {}}}"#, "must not be empty"),
(r#"{"scopes": {"graph": {}}}"#, "primary projection"),
] {
let err = super::GraphProjectionSpec::scopes_from_config_object(&cfg(json))
.expect_err("must be rejected");
assert!(
err.contains(needle),
"error for {json} must mention `{needle}`, got: {err}"
);
}
}
#[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");
}
}