use std::collections::{HashMap, HashSet};
use std::ops::ControlFlow;
use std::sync::Arc;
use datafusion::arrow::compute::can_cast_types;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::catalog::default_table_source::provider_as_source;
use datafusion::catalog::CatalogProviderList;
use datafusion::common::{Column, DFSchema, TableReference};
use datafusion::error::{DataFusionError, Result};
use datafusion::logical_expr::expr_fn::cast;
use datafusion::logical_expr::planner::{
PlannedRelation, RelationPlanner, RelationPlannerContext, RelationPlanning,
};
use datafusion::logical_expr::simplify::SimplifyContext;
use datafusion::logical_expr::TableSource;
use datafusion::logical_expr::{
EmptyRelation, Expr, ExprSchemable, Extension, LogicalPlan, LogicalPlanBuilder,
};
use datafusion::optimizer::simplify_expressions::ExprSimplifier;
use datafusion::sql::sqlparser::ast::{
self, visit_expressions_mut, FunctionArg, FunctionArgExpr, FunctionArgOperator,
FunctionArguments, TableFactor, TableFunctionArgs, UnaryOperator, Value, VisitMut,
};
use datafusion::sql::sqlparser::dialect::GenericDialect;
use datafusion::sql::sqlparser::parser::Parser;
use futures::FutureExt;
use ursa_core::algo::AggKind;
use ursa_core::catalog::{Call, FunctionSpec};
use super::functions::{
self, did_you_mean, is_per_node, nested_signature, plan, signature, Literal, ALGORITHMS,
VALUES_QUERY,
};
use super::kernels;
use super::node::{
endpoint_id_type, weight_column, GraphAlgorithmNode, GraphTraversalNode, Output,
WEIGHTED_DESTINATION, WEIGHTED_SOURCE,
};
use super::path_search::{PathColumns, PathSearchNode, Search};
use super::traversal::{self, TraversalParam, TraversalParamType, TraversalSpec};
use super::{GraphDef, GraphRegistry, PropertyGraph, VertexTable};
#[derive(Debug)]
pub struct GraphFunctions {
registry: Arc<GraphRegistry>,
catalogs: Arc<dyn CatalogProviderList>,
}
impl GraphFunctions {
pub fn new(registry: Arc<GraphRegistry>, catalogs: Arc<dyn CatalogProviderList>) -> Self {
Self { registry, catalogs }
}
}
#[derive(Clone, Copy)]
enum Target {
Function(&'static FunctionSpec),
Algorithms,
Traversal(&'static TraversalSpec),
}
impl Target {
fn name(&self) -> &'static str {
match self {
Target::Function(spec) => spec.name,
Target::Algorithms => ALGORITHMS,
Target::Traversal(spec) => spec.name,
}
}
fn signature(&self) -> String {
match self {
Target::Function(spec) => signature(spec),
Target::Algorithms => ALGORITHMS_SIGNATURE.to_string(),
Target::Traversal(spec) => spec.signature(),
}
}
}
const ALGORITHMS_SIGNATURE: &str =
"algorithms(<graph>, degree(direction => 'both'), community => louvain(seed => 1), ...)";
impl RelationPlanner for GraphFunctions {
fn plan_relation(
&self,
relation: TableFactor,
context: &mut dyn RelationPlannerContext,
) -> Result<RelationPlanning> {
let TableFactor::Table {
name,
args: Some(args),
..
} = &relation
else {
return Ok(RelationPlanning::Original(Box::new(relation)));
};
let [part] = name.0.as_slice() else {
return Ok(RelationPlanning::Original(Box::new(relation)));
};
let Some(ident) = part.as_ident() else {
return Ok(RelationPlanning::Original(Box::new(relation)));
};
let function_name = context.normalize_ident(ident.clone());
if function_name == super::pattern::FUNCTION {
return self.plan_graph_table(relation, context);
}
let target = if function_name == ALGORITHMS {
Target::Algorithms
} else if let Some(spec) = functions::lookup(&function_name) {
Target::Function(spec)
} else if let Some(spec) = traversal::lookup(&function_name) {
Target::Traversal(spec)
} else {
if self.first_argument_is_a_graph(args, context) {
return Err(no_such_function(
&function_name,
"",
functions::names().collect(),
));
}
return Ok(RelationPlanning::Original(Box::new(relation)));
};
let TableFactor::Table {
args: Some(args),
alias,
with_hints,
version,
with_ordinality,
partitions,
json_path,
sample,
index_hints,
..
} = relation
else {
unreachable!("matched above")
};
if !with_hints.is_empty()
|| version.is_some()
|| with_ordinality
|| !partitions.is_empty()
|| json_path.is_some()
|| sample.is_some()
|| !index_hints.is_empty()
|| args.settings.is_some()
{
return Err(plan(format!(
"{}() takes no table modifiers; write {}",
target.name(),
target.signature()
)));
}
let (graph, rest) = self.graph_argument(target, args, context)?;
let edges = self.edge_table(&graph, context)?;
let vertices = self.vertex_keys(&graph, context)?;
let mut weights = WeightSet::new(edges.schema().as_ref().clone());
let (outputs, per_node) = match target {
Target::Function(spec) if spec.reads_attribute() => {
let nodes = |c: &dyn RelationPlannerContext| {
self.table_source(&TableReference::partial(graph.name.as_str(), "nodes"), c)
};
let node = values_node(spec, &graph, edges, vertices, rest, context, &nodes)?;
let plan = LogicalPlan::Extension(Extension {
node: Arc::new(node),
});
return planned(plan, target, alias);
}
Target::Function(spec) => {
let (call, weight) = resolve_call(spec, rest, context, &mut weights, "")?;
let outputs = functions::outputs(spec)
.iter()
.enumerate()
.map(|(field, o)| Output {
name: o.name.to_string(),
call,
field,
weight,
})
.collect();
(outputs, is_per_node(spec))
}
Target::Algorithms => (algorithms(rest, context, &mut weights)?, true),
Target::Traversal(spec) => {
let node = traversal_node(spec, &graph, edges, vertices, rest, context)?;
let plan = LogicalPlan::Extension(Extension {
node: Arc::new(node),
});
return planned(plan, target, alias);
}
};
let node = node(&graph, edges, vertices, weights, outputs, per_node)?;
let plan = LogicalPlan::Extension(Extension {
node: Arc::new(node),
});
planned(plan, target, alias)
}
}
fn planned(
plan: LogicalPlan,
target: Target,
alias: Option<ast::TableAlias>,
) -> Result<RelationPlanning> {
let plan = match alias {
Some(_) => plan,
None => LogicalPlanBuilder::from(plan)
.alias(target.name())?
.build()?,
};
Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
plan, alias,
))))
}
type NodeTable<'a> = dyn Fn(&dyn RelationPlannerContext) -> Result<Arc<dyn TableSource>> + 'a;
#[allow(clippy::too_many_arguments)]
fn values_node(
spec: &'static FunctionSpec,
graph: &GraphDef,
edges: LogicalPlan,
vertices: Option<LogicalPlan>,
args: Vec<FunctionArg>,
context: &mut dyn RelationPlannerContext,
nodes: &NodeTable,
) -> Result<GraphAlgorithmNode> {
let f = spec.name;
let mut values: Option<ast::Expr> = None;
let mut attribute: Option<ast::Expr> = None;
let mut literals = Vec::new();
for arg in args {
if let FunctionArg::Named {
name,
arg: FunctionArgExpr::Expr(expr),
operator: FunctionArgOperator::RightArrow,
} = &arg
{
let name = context.normalize_ident(name.clone());
let slot = match name.as_str() {
"values" => &mut values,
"attribute" => &mut attribute,
_ => {
literals.push(arg);
continue;
}
};
if slot.replace(expr.clone()).is_some() {
return Err(plan(format!("{f}() was given '{name}' more than once")));
}
continue;
}
literals.push(arg);
}
let mut weights = WeightSet::new(edges.schema().as_ref().clone());
let (call, _) = resolve_call(spec, literals, context, &mut weights, "")?;
let by_order = matches!(
call,
Call::NeighborAgg {
agg: AggKind::Count | AggKind::NUnique | AggKind::Min | AggKind::Max,
..
}
);
let edges = edge_inputs(graph, edges, vertices, weights)?;
let id_type = endpoint_id_type(&graph.name, &edges.input, &edges.source, &edges.destination)?;
let in_param = |param: &'static str| {
move |e: DataFusionError| {
plan(format!(
"{f}({param} =>): {}",
crate::user_message(&e).replace('\n', " ")
))
}
};
let (param, values) = match (values, attribute) {
(Some(_), Some(_)) => {
return Err(plan(format!(
"{f}() takes 'values' or 'attribute', not both: values => {VALUES_QUERY} is a \
query of node ids and values, attribute => <expr> an expression over the \
graph's node table"
)))
}
(None, None) => {
return Err(plan(format!(
"{f}() needs the value of each node: values => {VALUES_QUERY}, a query of node \
ids and values, or attribute => <expr>, an expression over the graph's node \
table. Write {}",
signature(spec)
)))
}
(Some(ast::Expr::Subquery(query)), None) => {
let values = context
.plan(TableFactor::Derived {
lateral: false,
subquery: query,
alias: None,
sample: None,
})
.map_err(in_param("values"))?;
("values", values)
}
(Some(other), None) => {
return Err(plan(format!(
"{f}(values =>) must be a query of two columns, node ids then values, as \
values => {VALUES_QUERY}; got {other}"
)))
}
(None, Some(expr)) => {
let table = TableReference::partial(graph.name.as_str(), "nodes");
let scan = LogicalPlanBuilder::scan(table.clone(), nodes(&*context)?, None)?.build()?;
let value = context
.sql_to_expr(expr, scan.schema())
.map_err(in_param("attribute"))?;
let values = LogicalPlanBuilder::from(scan)
.project(vec![
Expr::Column(Column::new(Some(table), "id")),
value.alias("value"),
])
.and_then(LogicalPlanBuilder::build)
.map_err(in_param("attribute"))?;
("attribute", values)
}
};
let schema = values.schema();
let columns = schema.fields().len();
if columns != 2 {
return Err(plan(format!(
"{f}({param} =>) must be a query of two columns, node ids then values; it returns \
{columns} columns"
)));
}
let (ids, value) = (schema.field(0).data_type(), schema.field(1).data_type());
if !traversal::is_node_id_type(ids, &id_type) {
let kind = match id_type {
DataType::Int64 => "integers",
_ => "strings",
};
return Err(plan(format!(
"{f}({param} =>): the first column must be node ids of graph '{}', which are \
{kind}; got a column of type {ids}",
graph.name
)));
}
if kernels::is_ordered(value) {
if !by_order {
let agg = call
.values()
.into_iter()
.find(|(p, _)| p.name == "agg")
.map(|(_, v)| functions::sql_value(&v))
.unwrap_or_default();
let what = if kernels::is_string(value) {
"strings".to_string()
} else {
format!("of type {value}")
};
return Err(plan(format!(
"{f}(agg => {agg}) needs numbers, and the values are {what}; only 'count', \
'n_unique', 'min' and 'max' take strings, dates and times"
)));
}
} else if !(value == &DataType::Null || can_cast_types(value, &DataType::Float64)) {
return Err(plan(format!(
"{f}({param} =>): the values must be numbers, or strings, dates or times for \
'count', 'n_unique', 'min' and 'max'; got a column of type {value}"
)));
}
let outputs = functions::outputs(spec)
.iter()
.enumerate()
.map(|(field, o)| Output {
name: o.name.to_string(),
call,
field,
weight: None,
})
.collect();
GraphAlgorithmNode::try_new(
graph.name.clone(),
Arc::new(edges.input),
None,
edges.vertices,
Some(Arc::new(values)),
edges.source,
edges.destination,
edges.weights,
outputs,
true,
)
}
fn traversal_node(
spec: &'static TraversalSpec,
graph: &GraphDef,
edges: LogicalPlan,
vertices: Option<LogicalPlan>,
args: Vec<FunctionArg>,
context: &mut dyn RelationPlannerContext,
) -> Result<GraphTraversalNode> {
let f = spec.name;
let mut weights = WeightSet::new(edges.schema().as_ref().clone());
let mut literals = Vec::new();
let node_params: Vec<&TraversalParam> = spec.node_params().collect();
let mut nodes: Vec<Option<ast::Expr>> = vec![None; node_params.len()];
let mut seen = HashSet::new();
for arg in args {
let FunctionArg::Named {
name,
arg,
operator: FunctionArgOperator::RightArrow,
} = arg
else {
return Err(plan(format!(
"{f}() takes its parameters by name, as name => value; got {arg}. Write {}",
spec.signature()
)));
};
let name = context.normalize_ident(name);
let Some(param) = spec.param(&name) else {
return Err(traversal::unknown_param(spec, &name));
};
if !seen.insert(name.clone()) {
return Err(plan(format!("{f}() was given '{name}' more than once")));
}
let FunctionArgExpr::Expr(expr) = arg else {
return Err(plan(format!(
"{f}({name} =>) must be a value or an expression; got {arg}"
)));
};
if param.takes_nodes() {
let i = node_params
.iter()
.position(|p| p.name == param.name)
.expect("a node parameter");
nodes[i] = Some(expr);
} else if param.is_weight() {
weights.add(f, expr, context)?;
} else {
let arg = FunctionArgExpr::Expr(expr);
let value = literal(&arg).ok_or_else(|| {
plan(format!(
"{f}({name} =>) must be a literal (a number, a string or NULL); got {arg}"
))
})?;
literals.push((name, value));
}
}
let nodes = node_params
.iter()
.zip(nodes)
.map(|(p, n)| n.ok_or_else(|| traversal::missing_param(spec, p)))
.collect::<Result<Vec<_>>>()?;
let edges = edge_inputs(graph, edges, vertices, weights)?;
let id_type = endpoint_id_type(&graph.name, &edges.input, &edges.source, &edges.destination)?;
let ids = NodeIds {
spec,
graph: &graph.name,
id_type: &id_type,
};
let mut inputs = Vec::with_capacity(nodes.len());
let mut written = Vec::with_capacity(nodes.len());
for (p, expr) in node_params.iter().zip(nodes) {
written.push(expr.to_string());
inputs.push(Arc::new(ids.input(p, expr, context)?));
}
let traversal = traversal::resolve(spec, &literals)?;
GraphTraversalNode::try_new(
graph.name.clone(),
traversal,
Arc::new(edges.input),
edges.weighted,
edges.vertices,
edges.weights.into_iter().next(),
inputs,
written,
edges.source,
edges.destination,
)
}
struct NodeIds<'a> {
spec: &'static TraversalSpec,
graph: &'a str,
id_type: &'a DataType,
}
impl NodeIds<'_> {
fn input(
&self,
param: &TraversalParam,
expr: ast::Expr,
context: &mut dyn RelationPlannerContext,
) -> Result<LogicalPlan> {
let (f, p) = (self.spec.name, param.name);
match expr {
ast::Expr::Subquery(query) => {
let subquery = context
.plan(TableFactor::Derived {
lateral: false,
subquery: query,
alias: None,
sample: None,
})
.map_err(|e| {
plan(format!(
"{f}({p} =>): {}",
crate::user_message(&e).replace('\n', " ")
))
})?;
let columns = subquery.schema().fields().len();
if columns != 1 {
return Err(plan(format!(
"{f}({p} =>) must be a query of one column of node ids; it returns \
{columns} columns"
)));
}
let ty = subquery.schema().field(0).data_type();
if !traversal::is_node_id_type(ty, self.id_type) {
return Err(self.not_ids(param, &format!("a query of type {ty}")));
}
Ok(subquery)
}
ast::Expr::Array(array) => {
if param.ty == TraversalParamType::Node {
return Err(plan(format!(
"{f}({p} =>) takes one node id, not a list; got {array}"
)));
}
self.values(param, array.elem, context)
}
one => self.values(param, vec![one], context),
}
}
fn values(
&self,
param: &TraversalParam,
elems: Vec<ast::Expr>,
context: &mut dyn RelationPlannerContext,
) -> Result<LogicalPlan> {
let (f, p) = (self.spec.name, param.name);
if elems.is_empty() {
let schema = Schema::new(vec![Field::new("id", DataType::Null, true)]);
return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
produce_one_row: false,
schema: Arc::new(DFSchema::try_from(schema)?),
}));
}
let empty = Arc::new(DFSchema::empty());
let simplifier = ExprSimplifier::new(
SimplifyContext::builder()
.with_schema(Arc::clone(&empty))
.build(),
);
let rows = elems
.into_iter()
.map(|e| {
let written = e.to_string();
let hint = match &e {
ast::Expr::Identifier(i) => {
format!(
" (a string id is quoted: '{}')",
i.value.replace('\'', "''")
)
}
_ => String::new(),
};
let not_an_id = |err: DataFusionError| {
plan(format!(
"{f}({p} =>): {written} is not a node id{hint}: {}",
crate::user_message(&err).replace('\n', " ")
))
};
let expr = context
.sql_to_expr(e, &empty)
.and_then(|expr| simplifier.simplify(expr))
.map_err(not_an_id)?;
let ty = expr.get_type(&empty).map_err(not_an_id)?;
if !castable(&ty, self.id_type) {
return Err(self.not_ids(param, &format!("{written}, of type {ty}")));
}
if !traversal::is_node_id_type(&ty, self.id_type) {
return Err(self.not_ids(param, &written));
}
let expr = match expr {
Expr::Literal(value, metadata) => {
if let (DataType::Int64, Some(Some(text))) =
(self.id_type, value.try_as_str())
{
if !traversal::is_integer_text(text) {
return Err(self.not_ids(param, &written));
}
}
let value = value
.cast_to(self.id_type)
.map_err(|_| self.not_ids(param, &written))?;
Expr::Literal(value, metadata)
}
other => other.cast_to(self.id_type, &empty)?,
};
Ok(vec![expr])
})
.collect::<Result<Vec<_>>>()?;
LogicalPlanBuilder::values(rows).and_then(LogicalPlanBuilder::build)
}
fn not_ids(&self, param: &TraversalParam, got: &str) -> DataFusionError {
plan(traversal::not_node_ids(
self.spec.name,
param,
self.graph,
self.id_type,
got,
))
}
}
fn castable(from: &DataType, id_type: &DataType) -> bool {
from == &DataType::Null || can_cast_types(from, id_type)
}
fn no_such_function(name: &str, prefix: &str, candidates: Vec<&str>) -> DataFusionError {
let hint = did_you_mean(name, candidates.iter().copied())
.map(|s| format!("; did you mean '{s}'?"))
.unwrap_or_else(|| format!(". The graph functions are: {}", candidates.join(", ")));
plan(format!("{prefix}there is no graph function '{name}'{hint}"))
}
struct WeightSet {
schema: DFSchema,
exprs: Vec<Expr>,
written: Vec<String>,
}
impl WeightSet {
fn new(schema: DFSchema) -> Self {
Self {
schema,
exprs: Vec::new(),
written: Vec::new(),
}
}
fn add(
&mut self,
function: &str,
arg: ast::Expr,
context: &mut dyn RelationPlannerContext,
) -> Result<Option<usize>> {
if matches!(&arg, ast::Expr::Value(v) if v.value == Value::Null) {
return Ok(None);
}
let written = arg.to_string();
let context_error = |e: DataFusionError| {
plan(format!(
"{function}(weight => {written}): {}",
crate::user_message(&e).replace('\n', " ")
))
};
let expr = context
.sql_to_expr(arg, &self.schema)
.map_err(context_error)?;
let data_type = expr.get_type(&self.schema).map_err(context_error)?;
if !is_numeric(&data_type) {
return Err(plan(format!(
"{function}(weight => {written}): weight expression must be numeric; it \
produced a column of type {data_type} (a string or other non-numeric column \
cannot be a weight)"
)));
}
if let Some(i) = self.exprs.iter().position(|e| *e == expr) {
return Ok(Some(i));
}
self.exprs.push(expr);
self.written.push(written);
Ok(Some(self.exprs.len() - 1))
}
}
fn is_numeric(t: &DataType) -> bool {
use DataType::*;
matches!(
t,
Int8 | Int16
| Int32
| Int64
| UInt8
| UInt16
| UInt32
| UInt64
| Float16
| Float32
| Float64
| Decimal32(_, _)
| Decimal64(_, _)
| Decimal128(_, _)
| Decimal256(_, _)
)
}
fn resolve_call(
spec: &'static FunctionSpec,
args: Vec<FunctionArg>,
context: &mut dyn RelationPlannerContext,
weights: &mut WeightSet,
prefix: &str,
) -> Result<(Call, Option<usize>)> {
let in_context = |e: DataFusionError| match e {
DataFusionError::Plan(m) => plan(format!("{prefix}{m}")),
other => other,
};
let mut literals = Vec::new();
let mut weight: Option<Option<usize>> = None;
for arg in args {
match arg {
FunctionArg::Named {
name,
arg,
operator: FunctionArgOperator::RightArrow,
} => {
let name = context.normalize_ident(name);
if name == "weight" && spec.weighted() {
if weight.is_some() {
return Err(plan(format!(
"{prefix}{}() was given 'weight' more than once",
spec.name
)));
}
let FunctionArgExpr::Expr(expr) = arg else {
return Err(plan(format!(
"{prefix}{}(weight =>) must be an expression over the edge \
table's columns; got {arg}",
spec.name
)));
};
weight = Some(weights.add(spec.name, expr, context).map_err(in_context)?);
continue;
}
let known = name == "dtype" || functions::sql_params(spec).any(|p| p.name == name);
if !known {
return Err(in_context(functions::unknown_param(spec, &name)));
}
let value = literal(&arg).ok_or_else(|| {
plan(format!(
"{prefix}{}({name} =>) must be a literal (a number, a string, a \
boolean or NULL); got {arg}",
spec.name
))
})?;
literals.push((name, value));
}
other => {
let write = if prefix.is_empty() {
signature(spec)
} else {
nested_signature(spec)
};
return Err(plan(format!(
"{prefix}{}() takes its parameters by name, as name => value; got {other}. \
Write {write}",
spec.name
)));
}
}
}
let call = functions::resolve(spec, &literals).map_err(in_context)?;
Ok((call, weight.flatten()))
}
fn algorithms(
args: Vec<FunctionArg>,
context: &mut dyn RelationPlannerContext,
weights: &mut WeightSet,
) -> Result<Vec<Output>> {
let prefix = format!("{ALGORITHMS}(): ");
let per_node = || {
functions::specs()
.filter(|s| is_per_node(s) && !s.reads_attribute())
.map(|s| s.name)
};
let mut outputs = Vec::new();
for arg in &args {
let (key, inner) = match arg {
FunctionArg::Named {
name,
arg,
operator: FunctionArgOperator::RightArrow,
} => (Some(context.normalize_ident(name.clone())), arg),
FunctionArg::Unnamed(arg) => (None, arg),
other => {
return Err(plan(format!(
"{prefix}each argument after the graph is a call to a per-node graph \
function, as degree() or name => louvain(seed => 1); got {other}"
)))
}
};
let function = match inner {
FunctionArgExpr::Expr(ast::Expr::Function(f)) => f,
other => {
let hint = match other {
FunctionArgExpr::Expr(ast::Expr::Identifier(i))
if functions::lookup(&context.normalize_ident(i.clone())).is_some() =>
{
format!(" (write {}())", context.normalize_ident(i.clone()))
}
_ => String::new(),
};
return Err(plan(format!(
"{prefix}each argument after the graph is a call to a per-node graph \
function, as degree() or name => louvain(seed => 1); got {arg}{hint}"
)));
}
};
let name = match function.name.0.as_slice() {
[part] => part.as_ident().map(|i| context.normalize_ident(i.clone())),
_ => None,
}
.unwrap_or_else(|| function.name.to_string());
if let Some(t) = traversal::lookup(&name) {
return Err(plan(format!(
"{prefix}{name}() is a traversal, which returns rows of its own, not one per \
node; call it on its own: SELECT * FROM {}",
t.signature()
)));
}
let Some(spec) = functions::lookup(&name) else {
return Err(no_such_function(&name, &prefix, per_node().collect()));
};
if !is_per_node(spec) {
return Err(plan(format!(
"{prefix}{name}() is a statistic, with one row for the whole graph, not one \
per node; call it on its own: SELECT * FROM {name}(<graph>)"
)));
}
if spec.reads_attribute() {
return Err(plan(format!(
"{prefix}{name}() reads node values, so it runs on its own: SELECT * FROM {}; \
join its result on id",
signature(spec)
)));
}
let modified = function.uses_odbc_syntax
|| !matches!(function.parameters, FunctionArguments::None)
|| function.filter.is_some()
|| function.null_treatment.is_some()
|| function.over.is_some()
|| !function.within_group.is_empty();
let nested_args = match &function.args {
FunctionArguments::List(list)
if !modified && list.duplicate_treatment.is_none() && list.clauses.is_empty() =>
{
list.args.clone()
}
FunctionArguments::None if !modified => Vec::new(),
_ => {
return Err(plan(format!(
"{prefix}{name}() is called with named parameters only; got {function}. \
Write {}",
nested_signature(spec)
)))
}
};
if let Some(FunctionArg::Unnamed(FunctionArgExpr::Expr(ast::Expr::Identifier(_)))) =
nested_args.first()
{
return Err(plan(format!(
"{prefix}{name}() takes its parameters by name, as name => value; got {}. \
The graph is algorithms()'s first argument only: write {}",
nested_args[0],
nested_signature(spec)
)));
}
let (call, weight) = resolve_call(spec, nested_args, context, weights, &prefix)?;
let several = spec.outputs.len() > 1;
for (field, o) in spec.outputs.iter().enumerate() {
let column = match (&key, several) {
(Some(key), true) => format!("{key}_{}", o.name),
(Some(key), false) => key.clone(),
(None, _) => o.name.to_string(),
};
outputs.push(Output {
name: column,
call,
field,
weight,
});
}
}
if outputs.is_empty() {
return Err(plan(format!(
"{ALGORITHMS}() needs at least one call after the graph: {ALGORITHMS_SIGNATURE}"
)));
}
let mut seen = HashSet::new();
for o in &outputs {
if o.name == "id" {
return Err(plan(format!(
"{prefix}no column can be named 'id', which holds the node ids; give \
{}() another name with name => {}(...)",
o.call.name(),
o.call.name()
)));
}
if !seen.insert(o.name.as_str()) {
return Err(plan(format!(
"{prefix}two columns are named '{}'; give each call its own name, as in \
in_degree => degree(direction => 'in'), out_degree => degree()",
o.name
)));
}
}
Ok(outputs)
}
fn node(
graph: &GraphDef,
edges: LogicalPlan,
vertices: Option<LogicalPlan>,
weights: WeightSet,
outputs: Vec<Output>,
per_node: bool,
) -> Result<GraphAlgorithmNode> {
let edges = edge_inputs(graph, edges, vertices, weights)?;
GraphAlgorithmNode::try_new(
graph.name.clone(),
Arc::new(edges.input),
edges.weighted,
edges.vertices,
None,
edges.source,
edges.destination,
edges.weights,
outputs,
per_node,
)
}
struct EdgeInputs {
input: LogicalPlan,
weighted: Option<Arc<LogicalPlan>>,
vertices: Option<Arc<LogicalPlan>>,
source: Expr,
destination: Expr,
weights: Vec<String>,
}
fn edge_inputs(
graph: &GraphDef,
edges: LogicalPlan,
vertices: Option<LogicalPlan>,
weights: WeightSet,
) -> Result<EdgeInputs> {
let table = graph.edges.table.clone();
let column = |name: &str| Expr::Column(Column::new(Some(table.clone()), name));
let (source, destination) = (
column(&graph.edges.source),
column(&graph.edges.destination),
);
let no_endpoints = |e: DataFusionError| {
DataFusionError::Plan(format!(
"graph '{}': the edge table '{table}' has no column for its edges' source '{}' and \
destination '{}': {}",
graph.name,
graph.edges.source,
graph.edges.destination,
crate::user_message(&e)
))
};
let input = LogicalPlanBuilder::from(edges.clone())
.project(vec![source.clone(), destination.clone()])
.map_err(no_endpoints)?
.build()?;
if let (Some(keys), Some(v)) = (&vertices, &graph.vertices) {
let id_type = endpoint_id_type(&graph.name, &input, &source, &destination)?;
check_key_type(&graph.name, &v.table, keys, &id_type)?;
}
let weighted = if weights.exprs.is_empty() {
None
} else {
let mut exprs = vec![
source.clone().alias(WEIGHTED_SOURCE),
destination.clone().alias(WEIGHTED_DESTINATION),
];
exprs.extend(
weights
.exprs
.into_iter()
.enumerate()
.map(|(i, e)| e.alias(weight_column(i))),
);
Some(Arc::new(
LogicalPlanBuilder::from(edges).project(exprs)?.build()?,
))
};
Ok(EdgeInputs {
input,
weighted,
vertices: vertices.map(Arc::new),
source,
destination,
weights: weights.written,
})
}
fn keys_plan(
tables: &[&VertexTable],
id_type: &DataType,
source_of: &dyn Fn(&TableReference) -> Result<Arc<dyn TableSource>>,
) -> Result<Option<LogicalPlan>> {
let mut plan: Option<LogicalPlanBuilder> = None;
for v in tables {
let key = Expr::Column(Column::new(Some(v.table.clone()), &v.key));
let scan = LogicalPlanBuilder::scan(v.table.clone(), source_of(&v.table)?, None)?;
let key_type = key.get_type(scan.schema())?;
let key = match key_type == *id_type {
true => key,
false => cast(key, id_type.clone()),
};
let keys = scan.project(vec![key.alias(VERTEX_KEY)])?.build()?;
plan = Some(match plan {
None => LogicalPlanBuilder::from(keys),
Some(p) => p.union(keys)?,
});
}
plan.map(LogicalPlanBuilder::build).transpose()
}
const VERTEX_KEY: &str = "__ursa_key";
fn check_key_type(
graph: &str,
table: &TableReference,
keys: &LogicalPlan,
id_type: &DataType,
) -> Result<()> {
let key_type = keys.schema().field(0).data_type();
if super::node::id_type(key_type).as_ref() == Some(id_type) {
return Ok(());
}
Err(plan(format!(
"graph '{graph}': the key of vertex table '{table}' is {key_type}, but the edge \
endpoints are {id_type}; both must be integers or both strings"
)))
}
impl GraphFunctions {
fn vertex_keys(
&self,
graph: &GraphDef,
context: &dyn RelationPlannerContext,
) -> Result<Option<LogicalPlan>> {
let Some(v) = &graph.vertices else {
return Ok(None);
};
let source = match self.registry.file_table(&graph.name, &v.table) {
Some(provider) => provider_as_source(provider),
None => self.table_source(&v.table, context).map_err(|e| {
plan(format!(
"graph '{}': its vertex table cannot be read: {}",
graph.name,
crate::user_message(&e)
))
})?,
};
let key = Expr::Column(Column::new(Some(v.table.clone()), &v.key));
let plan = LogicalPlanBuilder::scan(v.table.clone(), source, None)?
.project(vec![key.alias(VERTEX_KEY)])
.map_err(|e| {
plan(format!(
"graph '{}': the vertex table '{}' has no key column '{}': {}",
graph.name,
v.table,
v.key,
crate::user_message(&e)
))
})?
.build()?;
Ok(Some(plan))
}
fn plan_graph_table(
&self,
relation: TableFactor,
context: &mut dyn RelationPlannerContext,
) -> Result<RelationPlanning> {
let TableFactor::Table {
args: Some(args),
alias,
..
} = relation
else {
unreachable!("matched by the caller")
};
let body = match args.args.as_slice() {
[FunctionArg::Unnamed(FunctionArgExpr::Expr(ast::Expr::Value(v)))] => match &v.value {
Value::SingleQuotedString(s) => s.clone(),
_ => return Err(graph_table_usage()),
},
_ => return Err(graph_table_usage()),
};
let pattern = super::pattern::parse(&body)?;
let graph = self.graph(&pattern.graph)?;
let source_of = |table: &TableReference, context: &dyn RelationPlannerContext| match self
.registry
.file_table(&graph.name, table)
{
Some(provider) => Ok(provider_as_source(provider)),
None => self.table_source(table, context),
};
let lowered = {
let context: &dyn RelationPlannerContext = context;
let schema_of = |table: &TableReference| -> Result<Arc<Schema>> {
Ok(source_of(table, context)?.schema())
};
super::pattern::lower(&pattern, &graph, &schema_of)?
};
let mut conditions = lowered
.conditions
.iter()
.map(|c| parse_sql_expr(c))
.collect::<Result<Vec<_>>>()?;
let mut columns = parse_select_items(&lowered.columns)?;
let paths: Vec<String> = lowered
.segments
.iter()
.filter_map(|s| s.path.clone())
.collect();
let mut used: HashMap<String, PathColumns> = HashMap::new();
for condition in &mut conditions {
path_functions(condition, &paths, &mut used, context)?;
}
for item in &mut columns {
if let ast::SelectItem::UnnamedExpr(expr @ ast::Expr::Function(_)) = item {
let written = expr.to_string();
let function = written.split('(').next().unwrap_or("").to_lowercase();
if super::path_search::FUNCTIONS
.iter()
.any(|(f, _)| *f == function)
{
*item = ast::SelectItem::ExprWithAlias {
expr: expr.clone(),
alias: ast::Ident::with_quote('"', written),
};
}
}
path_functions(item, &paths, &mut used, context)?;
}
let relation_of = |source: &super::pattern::Source,
context: &dyn RelationPlannerContext|
-> Result<LogicalPlan> {
let table = source_of(&source.table, context)?;
let mut scan = LogicalPlanBuilder::scan(source.table.clone(), table, None)?;
if let Some(columns) = &source.columns {
scan = scan.project(columns.iter().map(|(column, name)| {
Expr::Column(Column::new_unqualified(column)).alias(name)
}))?;
}
scan.alias(source.var.as_str())?.build()
};
let column = |(var, name): &super::pattern::VarColumn| {
let source = lowered.sources.iter().find(|s| &s.var == var);
let name = source
.and_then(|s| s.columns.as_ref())
.and_then(|columns| columns.iter().find(|(c, _)| c == name).map(|(_, n)| n))
.unwrap_or(name);
Expr::Column(Column::new(Some(var.as_str()), name))
};
let mut from: Option<LogicalPlanBuilder> = None;
let mut relations = Vec::new();
for source in &lowered.sources {
relations.push(relation_of(source, context)?);
}
for segment in &lowered.segments {
let start = lowered
.sources
.iter()
.find(|s| s.var == segment.start)
.expect("a segment starts at a node");
let start = relation_of(start, context)?;
let mut keep = Vec::new();
for condition in &conditions {
for conjunct in conjuncts(condition) {
if let Ok(expr) = context.sql_to_expr(conjunct.clone(), start.schema()) {
keep.push(expr);
}
}
}
let mut starts = LogicalPlanBuilder::from(start);
if let Some(predicate) = keep.into_iter().reduce(Expr::and) {
starts = starts.filter(predicate)?;
}
let starts = starts
.project(vec![
column(&segment.start_key).alias(super::path_search::START)
])?
.distinct()?
.build()?;
let table = source_of(&segment.edges, context)?;
let edge_column =
|name: &str| Expr::Column(Column::new(Some(segment.edges.clone()), name));
let (source, destination) = (
edge_column(&segment.source),
edge_column(&segment.destination),
);
let edges = LogicalPlanBuilder::scan(segment.edges.clone(), table, None)?
.project(vec![source.clone(), destination.clone()])?
.build()?;
let element = graph.edge_tables.iter().find(|e| {
e.table == segment.edges
&& e.source == segment.source
&& e.destination == segment.destination
});
let ends = element.map(|e| {
(
graph.endpoint_vertex(e, true),
graph.endpoint_vertex(e, false),
)
});
let vertices = match ends {
Some((Some(s), Some(d))) => {
let id_type = endpoint_id_type(&graph.name, &edges, &source, &destination)?;
let tables = if s.name == d.name {
vec![s]
} else {
vec![s, d]
};
let source_of = |t: &TableReference| source_of(t, &*context);
keys_plan(&tables, &id_type, &source_of)?.map(Arc::new)
}
_ => None,
};
let node = PathSearchNode::try_new(
graph.name.clone(),
Search {
direction: segment.direction,
quantifier: segment.quantifier,
selector: segment.selector,
mode: segment.mode,
},
segment
.path
.as_ref()
.and_then(|p| used.get(p).copied())
.unwrap_or_default(),
Arc::new(edges),
vertices,
Arc::new(starts),
source,
destination,
)?;
relations.push(
LogicalPlanBuilder::from(LogicalPlan::Extension(Extension {
node: Arc::new(node),
}))
.alias(segment.var.as_str())?
.build()?,
);
}
for relation in relations {
from = Some(match from {
None => LogicalPlanBuilder::from(relation),
Some(b) => b.cross_join(relation)?,
});
}
let base = from.expect("a pattern has a node").build()?;
let mut branches = Vec::new();
for joins in &lowered.branches {
let mut filter: Vec<Expr> = joins
.iter()
.map(|join| match join {
super::pattern::Join::Equal(a, b) => column(a).eq(column(b)),
super::pattern::Join::NotEqual(a, b) => column(a).not_eq(column(b)),
})
.collect();
for condition in &conditions {
filter.push(context.sql_to_expr(condition.clone(), base.schema())?);
}
let mut branch = LogicalPlanBuilder::from(base.clone());
if let Some(predicate) = filter.into_iter().reduce(Expr::and) {
branch = branch.filter(predicate)?;
}
let schema = branch.schema().clone();
let mut projection = Vec::new();
for item in &columns {
match item {
ast::SelectItem::UnnamedExpr(expr) => {
projection.push(context.sql_to_expr(expr.clone(), &schema)?)
}
ast::SelectItem::ExprWithAlias { expr, alias } => projection.push(
context
.sql_to_expr(expr.clone(), &schema)?
.alias(context.normalize_ident(alias.clone())),
),
ast::SelectItem::QualifiedWildcard(
ast::SelectItemQualifiedWildcardKind::ObjectName(name),
_,
) => {
let var = context.normalize_ident(
name.0
.last()
.and_then(|p| p.as_ident().cloned())
.unwrap_or_else(|| ast::Ident::new(name.to_string())),
);
let columns: Vec<Expr> = schema
.iter()
.filter(|(q, _)| q.is_some_and(|q| q.table() == var))
.map(|(q, f)| Expr::Column(Column::new(q.cloned(), f.name())))
.collect();
if columns.is_empty() {
return Err(plan(format!(
"in GRAPH_TABLE: COLUMNS ({var}.*) names no variable '{var}'"
)));
}
projection.extend(columns);
}
other => {
return Err(plan(format!(
"in GRAPH_TABLE: COLUMNS takes expressions, as a.name AS name; got {other}"
)))
}
}
}
branches.push(branch.project(projection)?.build()?);
}
let mut branches = branches.into_iter();
let mut plan = LogicalPlanBuilder::from(branches.next().expect("one branch at least"));
for branch in branches {
plan = plan.union(branch)?;
}
let plan = match alias {
Some(_) => plan.build()?,
None => plan.alias(super::pattern::FUNCTION)?.build()?,
};
Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
plan, alias,
))))
}
fn first_argument_is_a_graph(
&self,
args: &TableFunctionArgs,
context: &dyn RelationPlannerContext,
) -> bool {
match args.args.first() {
Some(FunctionArg::Unnamed(FunctionArgExpr::Expr(ast::Expr::Identifier(ident)))) => self
.registry
.get(&context.normalize_ident(ident.clone()))
.is_some(),
_ => false,
}
}
fn graph_argument(
&self,
target: Target,
args: TableFunctionArgs,
context: &dyn RelationPlannerContext,
) -> Result<(Arc<GraphDef>, Vec<FunctionArg>)> {
let name = target.name();
let mut args = args.args.into_iter();
let missing = || {
plan(format!(
"{name}() needs a graph as its first argument: {}",
target.signature()
))
};
let graph = match args.next() {
None => return Err(missing()),
Some(FunctionArg::Unnamed(FunctionArgExpr::Expr(ast::Expr::Identifier(ident)))) => {
self.graph(&context.normalize_ident(ident))?
}
Some(FunctionArg::Unnamed(FunctionArgExpr::Expr(ast::Expr::Function(_))))
if matches!(target, Target::Algorithms) =>
{
return Err(missing())
}
Some(FunctionArg::Unnamed(other)) => {
return Err(plan(format!(
"the first argument of {name}() must be the name of a graph; got {other}{}",
quoted_name_hint(&other),
)))
}
Some(_) => return Err(missing()),
};
let mut args = args.peekable();
let mut labels = Vec::new();
while let Some(FunctionArg::Unnamed(FunctionArgExpr::Expr(ast::Expr::Identifier(ident)))) =
args.peek()
{
if matches!(target, Target::Algorithms) {
let label = context.normalize_ident(ident.clone());
if graph.vertex_table(&label).is_none() && graph.edge_table(&label).is_none() {
break;
}
}
let Some(FunctionArg::Unnamed(FunctionArgExpr::Expr(ast::Expr::Identifier(ident)))) =
args.next()
else {
unreachable!("peeked")
};
labels.push(context.normalize_ident(ident));
}
let pair = match labels.as_slice() {
[] => graph.pair(None, None)?,
[vertex, edge] => graph.pair(Some(vertex), Some(edge))?,
_ => {
return Err(plan(format!(
"{name}() takes a graph, optionally followed by a vertex table and an edge \
table ({name}({}, <vertex table>, <edge table>)), then named parameters; \
got {} name{} after the graph",
graph.name,
labels.len(),
if labels.len() == 1 { "" } else { "s" }
)))
}
};
Ok((Arc::new(pair), args.collect()))
}
fn graph(&self, name: &str) -> Result<Arc<PropertyGraph>> {
if let Some(graph) = self.registry.get(name) {
return Ok(graph);
}
let names = self.registry.names();
let hint = match did_you_mean(name, names.iter().map(String::as_str)) {
Some(s) => format!("; did you mean '{s}'?"),
None if names.is_empty() => {
"; no graphs are defined (define one with CREATE PROPERTY GRAPH)".to_string()
}
None => format!(". The graphs are: {}", names.join(", ")),
};
Err(plan(format!("graph '{name}' does not exist{hint}")))
}
}
fn conjuncts(expr: &ast::Expr) -> Vec<&ast::Expr> {
match expr {
ast::Expr::BinaryOp {
left,
op: ast::BinaryOperator::And,
right,
} => {
let mut parts = conjuncts(left);
parts.extend(conjuncts(right));
parts
}
ast::Expr::Nested(inner) => conjuncts(inner),
other => vec![other],
}
}
fn path_functions<V: VisitMut>(
node: &mut V,
paths: &[String],
used: &mut HashMap<String, PathColumns>,
context: &dyn RelationPlannerContext,
) -> Result<()> {
let flow = visit_expressions_mut(node, |expr: &mut ast::Expr| {
let ast::Expr::Function(f) = expr else {
return ControlFlow::Continue(());
};
let [name] = f.name.0.as_slice() else {
return ControlFlow::Continue(());
};
let Some(name) = name.as_ident() else {
return ControlFlow::Continue(());
};
let function = name.value.to_lowercase();
let Some((_, column)) = super::path_search::FUNCTIONS
.iter()
.find(|(f, _)| *f == function)
else {
return ControlFlow::Continue(());
};
let arg = match &f.args {
FunctionArguments::List(list) => match list.args.as_slice() {
[FunctionArg::Unnamed(FunctionArgExpr::Expr(ast::Expr::Identifier(i)))] => {
Some(context.normalize_ident(i.clone()))
}
_ => None,
},
_ => None,
};
let Some(path) = arg.filter(|a| paths.contains(a)) else {
return ControlFlow::Break(plan(format!(
"in GRAPH_TABLE: {function}() takes a path variable, as {function}(p) for MATCH p = (a)-[e]->*(b){}",
if paths.is_empty() {
String::new()
} else {
format!("; the paths are {}", paths.join(", "))
}
)));
};
let needs = used.entry(path.clone()).or_default();
match *column {
super::path_search::VERTICES => needs.vertices = true,
super::path_search::EDGES => needs.edges = true,
super::path_search::ELEMENTS => needs.elements = true,
_ => {}
}
*expr = ast::Expr::CompoundIdentifier(vec![
ast::Ident::with_quote('"', path),
ast::Ident::with_quote('"', *column),
]);
ControlFlow::Continue(())
});
match flow {
ControlFlow::Break(e) => Err(e),
ControlFlow::Continue(()) => Ok(()),
}
}
fn parse_sql_expr(text: &str) -> Result<ast::Expr> {
Parser::new(&GenericDialect {})
.try_with_sql(text)
.and_then(|mut p| p.parse_expr())
.map_err(|e| plan(format!("in GRAPH_TABLE: {e} in '{text}'")))
}
fn parse_select_items(text: &str) -> Result<Vec<ast::SelectItem>> {
Parser::new(&GenericDialect {})
.try_with_sql(text)
.and_then(|mut p| p.parse_comma_separated(Parser::parse_select_item))
.map_err(|e| plan(format!("in GRAPH_TABLE: {e} in COLUMNS ({text})")))
}
fn graph_table_usage() -> DataFusionError {
DataFusionError::Plan(
"graph_table() takes one string, a GRAPH_TABLE body; write \
GRAPH_TABLE (graph MATCH (a)-[e]->(b) COLUMNS (a.id, b.id))"
.to_string(),
)
}
fn quoted_name_hint(arg: &FunctionArgExpr) -> String {
match arg {
FunctionArgExpr::Expr(ast::Expr::Value(v)) => match &v.value {
Value::SingleQuotedString(s) => format!(" (write the name without quotes: {s})"),
_ => String::new(),
},
_ => String::new(),
}
}
fn literal(arg: &FunctionArgExpr) -> Option<Literal> {
let FunctionArgExpr::Expr(expr) = arg else {
return None;
};
fn number(expr: &ast::Expr) -> Option<String> {
match expr {
ast::Expr::Value(v) => match &v.value {
Value::Number(n, false) => Some(n.to_string()),
_ => None,
},
ast::Expr::UnaryOp {
op: UnaryOperator::Minus,
expr,
} => number(expr).map(|n| match n.strip_prefix('-') {
Some(positive) => positive.to_string(),
None => format!("-{n}"),
}),
ast::Expr::UnaryOp {
op: UnaryOperator::Plus,
expr,
} => number(expr),
ast::Expr::Nested(inner) => number(inner),
_ => None,
}
}
if let Some(n) = number(expr) {
return Some(Literal::Number(n));
}
match expr {
ast::Expr::Value(v) => match &v.value {
Value::SingleQuotedString(s) => Some(Literal::String(s.clone())),
Value::Boolean(b) => Some(Literal::Boolean(*b)),
Value::Null => Some(Literal::Null),
_ => None,
},
_ => None,
}
}
impl GraphFunctions {
fn table_source(
&self,
table: &TableReference,
context: &dyn RelationPlannerContext,
) -> Result<Arc<dyn TableSource>> {
let provider = context.context_provider();
if let Ok(source) = provider.get_table_source(table.clone()) {
return Ok(source);
}
let options = provider.options();
let resolved = table.clone().resolve(
&options.catalog.default_catalog,
&options.catalog.default_schema,
);
let not_found = || plan(format!("table '{table}' not found"));
let schema = self
.catalogs
.catalog(&resolved.catalog)
.and_then(|c| c.schema(&resolved.schema))
.ok_or_else(not_found)?;
match schema.table(&resolved.table).now_or_never() {
Some(Ok(Some(table))) => Ok(provider_as_source(table)),
Some(Ok(None)) => Err(not_found()),
Some(Err(e)) => Err(e),
None => Err(plan(format!(
"'{table}' must be a table or view in the session (create a view over it)"
))),
}
}
fn edge_table(
&self,
graph: &GraphDef,
context: &dyn RelationPlannerContext,
) -> Result<LogicalPlan> {
let table = graph.edges.table.clone();
if let Some(file) = self.registry.file_table(&graph.name, &table) {
return super::edges_plan_from(graph, provider_as_source(file));
}
let source = self.table_source(&table, context).map_err(|e| {
plan(format!(
"graph '{}': its edge table cannot be read: {}",
graph.name,
crate::user_message(&e)
))
})?;
super::edges_plan_from(graph, source)
}
}