#![forbid(unsafe_code)]
#![allow(clippy::unnecessary_literal_bound)]
use std::cmp::Ordering;
use std::collections::HashMap;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::path::PathBuf;
use std::sync::Arc;
use datafusion::arrow::datatypes::{DataType, Field, Fields, Schema};
use datafusion::common::{DFSchema, DFSchemaRef, Result as DfResult, TableReference};
use datafusion::logical_expr::{Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore};
use graphforge_core::OntologyMode;
use graphforge_ir::{Direction, IrLiteral};
macro_rules! impl_partial_ord {
($t:ty) => {
impl PartialOrd for $t {
fn partial_cmp(&self, _other: &Self) -> Option<Ordering> {
None
}
}
};
}
fn passthrough_schema(inputs: &[&LogicalPlan]) -> DFSchemaRef {
inputs
.first()
.map_or_else(|| Arc::new(DFSchema::empty()), |p| p.schema().clone())
}
pub const VAR_LEN_EDGE_LIST_FIELD: &str = "rels";
#[must_use]
pub fn var_len_edge_list_field(prop_fields: &[Field]) -> Arc<Field> {
let mut fields = vec![
Field::new("edge_uuid", DataType::FixedSizeBinary(16), false),
Field::new("src_uuid", DataType::FixedSizeBinary(16), false),
Field::new("dst_uuid", DataType::FixedSizeBinary(16), false),
Field::new("rel_type", DataType::Utf8, true),
];
fields.extend(prop_fields.iter().cloned());
let struct_fields = Fields::from(fields);
let item = Arc::new(Field::new("item", DataType::Struct(struct_fields), true));
Arc::new(Field::new(
VAR_LEN_EDGE_LIST_FIELD,
DataType::List(item),
true,
))
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct VarLenExpandNode {
pub input: Arc<LogicalPlan>,
pub rel_type_name: String,
pub min_hops: u16,
pub max_hops: Option<u16>,
pub src_var: u32,
pub dst_var: u32,
pub edge_var: u32,
pub direction: Direction,
pub rel_ty: Option<u32>,
pub dir: PathBuf,
pub mode: OntologyMode,
schema: DFSchemaRef,
}
impl VarLenExpandNode {
#[allow(clippy::too_many_arguments)]
#[must_use]
pub fn new(
input: Arc<LogicalPlan>,
rel_type_name: impl Into<String>,
min_hops: u16,
max_hops: Option<u16>,
src_var: u32,
dst_var: u32,
edge_var: u32,
direction: Direction,
rel_ty: Option<u32>,
dir: PathBuf,
mode: OntologyMode,
dst_fields: Vec<Arc<Field>>,
edge_field: Arc<Field>,
) -> Self {
let schema = Self::build_schema(&input, dst_var, edge_var, dst_fields, edge_field);
Self {
input,
rel_type_name: rel_type_name.into(),
min_hops,
max_hops,
src_var,
dst_var,
edge_var,
direction,
rel_ty,
dir,
mode,
schema,
}
}
fn build_schema(
input: &Arc<LogicalPlan>,
dst_var: u32,
edge_var: u32,
dst_fields: Vec<Arc<Field>>,
edge_field: Arc<Field>,
) -> DFSchemaRef {
let dst_qualifier = TableReference::bare(format!("var_{dst_var}"));
let edge_qualifier = TableReference::bare(format!("var_{edge_var}"));
let mut qualified: Vec<(Option<TableReference>, Arc<Field>)> = input
.schema()
.iter()
.map(|(q, f)| (q.cloned(), Arc::clone(f)))
.collect();
qualified.extend(
dst_fields
.into_iter()
.map(|f| (Some(dst_qualifier.clone()), f)),
);
qualified.push((Some(edge_qualifier), edge_field));
Arc::new(
DFSchema::new_with_metadata(qualified, std::collections::HashMap::new())
.expect("VarLenExpandNode schema must include qualified destination + edge fields"),
)
}
}
impl_partial_ord!(VarLenExpandNode);
impl UserDefinedLogicalNodeCore for VarLenExpandNode {
fn name(&self) -> &str {
"VarLenExpand"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
vec![]
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
let max = self.max_hops.map_or("*".to_owned(), |h| h.to_string());
let arrow = match self.direction {
Direction::Out => "->",
Direction::In => "<-",
Direction::Undirected => "--",
};
write!(
f,
"VarLenExpand: rel={}, hops={}..{}, dir={arrow}, edge=var_{}",
self.rel_type_name, self.min_hops, max, self.edge_var
)
}
fn with_exprs_and_inputs(&self, _exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DfResult<Self> {
let input = Arc::new(
inputs
.into_iter()
.next()
.unwrap_or_else(|| (*self.input).clone()),
);
let dst_qual = format!("var_{}", self.dst_var);
let dst_fields: Vec<Arc<Field>> = self
.schema
.iter()
.filter(|(q, _)| q.map(TableReference::table) == Some(dst_qual.as_str()))
.map(|(_, f)| Arc::clone(f))
.collect();
let edge_qual = format!("var_{}", self.edge_var);
let edge_field = self
.schema
.iter()
.find(|(q, _)| q.map(TableReference::table) == Some(edge_qual.as_str()))
.map_or_else(|| var_len_edge_list_field(&[]), |(_, f)| Arc::clone(f));
Ok(Self::new(
input,
self.rel_type_name.clone(),
self.min_hops,
self.max_hops,
self.src_var,
self.dst_var,
self.edge_var,
self.direction,
self.rel_ty,
self.dir.clone(),
self.mode,
dst_fields,
edge_field,
))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ExpandNode {
pub input: Arc<LogicalPlan>,
pub rel_type_name: String,
pub src_var: u32,
pub dst_var: u32,
pub edge_var: u32,
pub direction: Direction,
pub rel_ty: Option<u32>,
pub dir: PathBuf,
pub mode: OntologyMode,
pub edge_prop_count: usize,
schema: DFSchemaRef,
}
impl ExpandNode {
#[allow(clippy::too_many_arguments)]
#[must_use]
pub fn new(
input: Arc<LogicalPlan>,
rel_type_name: impl Into<String>,
src_var: u32,
dst_var: u32,
edge_var: u32,
direction: Direction,
rel_ty: Option<u32>,
dir: PathBuf,
mode: OntologyMode,
edge_fields: Vec<Arc<Field>>,
edge_prop_fields: Vec<Arc<Field>>,
dst_fields: Vec<Arc<Field>>,
) -> Self {
let edge_prop_count = edge_prop_fields.len();
let schema = Self::build_schema(
&input,
dst_var,
edge_var,
edge_fields,
edge_prop_fields,
dst_fields,
);
Self {
input,
rel_type_name: rel_type_name.into(),
src_var,
dst_var,
edge_var,
direction,
rel_ty,
dir,
mode,
edge_prop_count,
schema,
}
}
fn build_schema(
input: &Arc<LogicalPlan>,
dst_var: u32,
edge_var: u32,
edge_fields: Vec<Arc<Field>>,
edge_prop_fields: Vec<Arc<Field>>,
dst_fields: Vec<Arc<Field>>,
) -> DFSchemaRef {
let edge_qualifier = TableReference::bare(format!("var_{edge_var}"));
let dst_qualifier = TableReference::bare(format!("var_{dst_var}"));
let mut qualified: Vec<(Option<TableReference>, Arc<Field>)> = input
.schema()
.iter()
.map(|(q, f)| (q.cloned(), Arc::clone(f)))
.collect();
qualified.extend(
edge_fields
.into_iter()
.map(|f| (Some(edge_qualifier.clone()), f)),
);
qualified.extend(edge_prop_fields.into_iter().map(|f| {
let f = Arc::new(f.as_ref().clone().with_nullable(true));
(Some(edge_qualifier.clone()), f)
}));
qualified.extend(
dst_fields
.into_iter()
.map(|f| (Some(dst_qualifier.clone()), f)),
);
Arc::new(
DFSchema::new_with_metadata(qualified, std::collections::HashMap::new())
.expect("ExpandNode schema must include qualified edge + destination fields"),
)
}
}
impl_partial_ord!(ExpandNode);
impl UserDefinedLogicalNodeCore for ExpandNode {
fn name(&self) -> &str {
"Expand"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
vec![]
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
let arrow = match self.direction {
Direction::Out => "->",
Direction::In => "<-",
Direction::Undirected => "--",
};
write!(
f,
"Expand: rel={}, dir={arrow}, edge=var_{}",
self.rel_type_name, self.edge_var
)
}
fn with_exprs_and_inputs(&self, _exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DfResult<Self> {
let input = Arc::new(
inputs
.into_iter()
.next()
.unwrap_or_else(|| (*self.input).clone()),
);
let edge_qual = format!("var_{}", self.edge_var);
let all_edge: Vec<Arc<Field>> = self
.schema
.iter()
.filter(|(q, _)| q.map(TableReference::table) == Some(edge_qual.as_str()))
.map(|(_, f)| Arc::clone(f))
.collect();
let topo_count = all_edge.len().saturating_sub(self.edge_prop_count);
let edge_fields = all_edge[..topo_count].to_vec();
let edge_prop_fields = all_edge[topo_count..].to_vec();
let dst_qual = format!("var_{}", self.dst_var);
let dst_fields: Vec<Arc<Field>> = self
.schema
.iter()
.filter(|(q, _)| q.map(TableReference::table) == Some(dst_qual.as_str()))
.map(|(_, f)| Arc::clone(f))
.collect();
Ok(Self::new(
input,
self.rel_type_name.clone(),
self.src_var,
self.dst_var,
self.edge_var,
self.direction,
self.rel_ty,
self.dir.clone(),
self.mode,
edge_fields,
edge_prop_fields,
dst_fields,
))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct OptionalMatchNode {
pub outer: Arc<LogicalPlan>,
pub optional: Arc<LogicalPlan>,
pub join_keys: Vec<(usize, usize)>,
pub inner_keep_idx: Vec<usize>,
schema: DFSchemaRef,
}
impl OptionalMatchNode {
#[must_use]
pub fn new(
outer: Arc<LogicalPlan>,
optional: Arc<LogicalPlan>,
join_keys: Vec<(usize, usize)>,
inner_keep_idx: Vec<usize>,
) -> Self {
let schema = Self::build_schema(&outer, &optional, &inner_keep_idx);
Self {
outer,
optional,
join_keys,
inner_keep_idx,
schema,
}
}
fn build_schema(
outer: &Arc<LogicalPlan>,
optional: &Arc<LogicalPlan>,
inner_keep_idx: &[usize],
) -> DFSchemaRef {
let mut qualified: Vec<(Option<TableReference>, Arc<Field>)> = outer
.schema()
.iter()
.map(|(q, f)| (q.cloned(), Arc::clone(f)))
.collect();
let inner: Vec<(Option<TableReference>, Arc<Field>)> = optional
.schema()
.iter()
.map(|(q, f)| (q.cloned(), Arc::clone(f)))
.collect();
qualified.extend(inner_keep_idx.iter().map(|&i| {
let (q, f) = &inner[i];
let nullable = Arc::new(f.as_ref().clone().with_nullable(true));
(q.clone(), nullable)
}));
Arc::new(
DFSchema::new_with_metadata(qualified, std::collections::HashMap::new())
.expect("OptionalMatchNode schema must be constructible from outer + inner fields"),
)
}
}
impl_partial_ord!(OptionalMatchNode);
impl UserDefinedLogicalNodeCore for OptionalMatchNode {
fn name(&self) -> &str {
"OptionalMatch"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.outer, &self.optional]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
vec![]
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "OptionalMatch: keys={}", self.join_keys.len())
}
fn with_exprs_and_inputs(
&self,
_exprs: Vec<Expr>,
mut inputs: Vec<LogicalPlan>,
) -> DfResult<Self> {
let optional = Arc::new(inputs.pop().unwrap_or_else(|| (*self.optional).clone()));
let outer = Arc::new(inputs.pop().unwrap_or_else(|| (*self.outer).clone()));
Ok(Self::new(
outer,
optional,
self.join_keys.clone(),
self.inner_keep_idx.clone(),
))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct PathUniqueNode {
pub input: Arc<LogicalPlan>,
schema: DFSchemaRef,
}
impl PathUniqueNode {
#[must_use]
pub fn new(input: Arc<LogicalPlan>) -> Self {
let schema = passthrough_schema(&[&input]);
Self { input, schema }
}
}
impl_partial_ord!(PathUniqueNode);
impl UserDefinedLogicalNodeCore for PathUniqueNode {
fn name(&self) -> &str {
"PathUnique"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
vec![]
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "PathUnique")
}
fn with_exprs_and_inputs(&self, _exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DfResult<Self> {
let input = Arc::new(
inputs
.into_iter()
.next()
.unwrap_or_else(|| (*self.input).clone()),
);
Ok(Self::new(input))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct OntologyInferNode {
pub input: Arc<LogicalPlan>,
pub rel_type_name: String,
pub rule_id: String,
pub confidence_model: String,
schema: DFSchemaRef,
}
impl OntologyInferNode {
#[must_use]
pub fn new(
input: Arc<LogicalPlan>,
rel_type_name: impl Into<String>,
rule_id: impl Into<String>,
confidence_model: impl Into<String>,
) -> Self {
let schema = passthrough_schema(&[&input]);
Self {
input,
rel_type_name: rel_type_name.into(),
rule_id: rule_id.into(),
confidence_model: confidence_model.into(),
schema,
}
}
}
impl_partial_ord!(OntologyInferNode);
impl UserDefinedLogicalNodeCore for OntologyInferNode {
fn name(&self) -> &str {
"OntologyInfer"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
vec![]
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"OntologyInfer: rel={} rule_id={}",
self.rel_type_name, self.rule_id
)
}
fn with_exprs_and_inputs(&self, _exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DfResult<Self> {
let input = Arc::new(
inputs
.into_iter()
.next()
.unwrap_or_else(|| (*self.input).clone()),
);
Ok(Self::new(
input,
self.rel_type_name.clone(),
self.rule_id.clone(),
self.confidence_model.clone(),
))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct GraphMergeNode {
schema: DFSchemaRef,
}
impl GraphMergeNode {
#[must_use]
pub fn new() -> Self {
Self {
schema: Arc::new(DFSchema::empty()),
}
}
}
impl_partial_ord!(GraphMergeNode);
impl Default for GraphMergeNode {
fn default() -> Self {
Self::new()
}
}
impl UserDefinedLogicalNodeCore for GraphMergeNode {
fn name(&self) -> &str {
"GraphMerge"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
vec![]
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "GraphMerge")
}
fn with_exprs_and_inputs(
&self,
_exprs: Vec<Expr>,
_inputs: Vec<LogicalPlan>,
) -> DfResult<Self> {
Ok(Self::new())
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct ResolvedNodeSpec {
pub var: u32,
pub label_ids: Vec<u32>,
pub label_names: Vec<String>,
pub properties: Vec<(String, IrLiteral)>,
pub computed_properties: Vec<(String, Expr)>,
pub is_reference: bool,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ResolvedEdgeSpec {
pub var: u32,
pub src: u32,
pub dst: u32,
pub rel_type_id: Option<u32>,
pub rel_type_name: Option<String>,
pub direction: Direction,
pub properties: Vec<(String, IrLiteral)>,
pub computed_properties: Vec<(String, Expr)>,
}
#[derive(Debug, Clone)]
pub struct GraphCreateNode {
pub input: Arc<LogicalPlan>,
pub nodes: Vec<ResolvedNodeSpec>,
pub edges: Vec<ResolvedEdgeSpec>,
pub dir: PathBuf,
pub mode: OntologyMode,
schema: DFSchemaRef,
emit_rows: bool,
}
impl GraphCreateNode {
#[must_use]
pub fn summary_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("nodes_created", DataType::UInt64, false),
Field::new("edges_created", DataType::UInt64, false),
Field::new("properties_set", DataType::UInt64, false),
Field::new("labels_added", DataType::UInt64, false),
]))
}
#[must_use]
pub fn new(
input: Arc<LogicalPlan>,
nodes: Vec<ResolvedNodeSpec>,
edges: Vec<ResolvedEdgeSpec>,
dir: PathBuf,
mode: OntologyMode,
) -> Self {
let schema = Arc::new(
DFSchema::try_from(Self::summary_schema())
.expect("write-summary schema is always valid"),
);
Self {
input,
nodes,
edges,
dir,
mode,
schema,
emit_rows: false,
}
}
#[must_use]
pub fn new_emitting(
input: Arc<LogicalPlan>,
nodes: Vec<ResolvedNodeSpec>,
edges: Vec<ResolvedEdgeSpec>,
dir: PathBuf,
mode: OntologyMode,
output_schema: DFSchemaRef,
) -> Self {
Self {
input,
nodes,
edges,
dir,
mode,
schema: output_schema,
emit_rows: true,
}
}
#[must_use]
pub fn emits_rows(&self) -> bool {
self.emit_rows
}
}
impl_partial_ord!(GraphCreateNode);
impl PartialEq for GraphCreateNode {
fn eq(&self, other: &Self) -> bool {
self.input == other.input
&& self.nodes == other.nodes
&& self.edges == other.edges
&& self.dir == other.dir
&& self.mode == other.mode
&& self.emit_rows == other.emit_rows
}
}
impl Eq for GraphCreateNode {}
fn hash_literal<H: Hasher>(lit: &IrLiteral, state: &mut H) {
match lit {
IrLiteral::Null => 0u8.hash(state),
IrLiteral::Bool(b) => {
1u8.hash(state);
b.hash(state);
}
IrLiteral::Int(i) => {
2u8.hash(state);
i.hash(state);
}
IrLiteral::Float(f) => {
3u8.hash(state);
f.to_bits().hash(state);
}
IrLiteral::Str(s) => {
4u8.hash(state);
s.hash(state);
}
IrLiteral::Uuid(uuid) => {
14u8.hash(state);
uuid.hash(state);
}
IrLiteral::Duration {
months,
days,
seconds,
nanos,
} => {
5u8.hash(state);
months.hash(state);
days.hash(state);
seconds.hash(state);
nanos.hash(state);
}
IrLiteral::DateTime(t) => {
6u8.hash(state);
t.hash(state);
}
IrLiteral::Date(d) => {
7u8.hash(state);
d.hash(state);
}
IrLiteral::LocalDateTime { days, nanos } => {
8u8.hash(state);
days.hash(state);
nanos.hash(state);
}
IrLiteral::Time(n) => {
9u8.hash(state);
n.hash(state);
}
IrLiteral::ZonedTime { nanos, offset } => {
10u8.hash(state);
nanos.hash(state);
offset.hash(state);
}
IrLiteral::ZonedDateTime {
days,
nanos,
offset,
zone,
} => {
11u8.hash(state);
days.hash(state);
nanos.hash(state);
offset.hash(state);
zone.hash(state);
}
IrLiteral::List(items) => {
12u8.hash(state);
items.len().hash(state);
for it in items {
hash_literal(it, state);
}
}
IrLiteral::Map(entries) => {
13u8.hash(state);
entries.len().hash(state);
for (key, value) in entries {
key.hash(state);
hash_literal(value, state);
}
}
}
}
fn hash_props<H: Hasher>(props: &[(String, IrLiteral)], state: &mut H) {
props.len().hash(state);
for (k, v) in props {
k.hash(state);
hash_literal(v, state);
}
}
fn hash_computed<H: Hasher>(props: &[(String, Expr)], state: &mut H) {
props.len().hash(state);
for (k, _) in props {
k.hash(state);
}
}
impl Hash for GraphCreateNode {
fn hash<H: Hasher>(&self, state: &mut H) {
self.input.hash(state);
self.nodes.len().hash(state);
for n in &self.nodes {
n.var.hash(state);
n.label_ids.hash(state);
n.label_names.hash(state);
n.is_reference.hash(state);
hash_props(&n.properties, state);
hash_computed(&n.computed_properties, state);
}
self.edges.len().hash(state);
for e in &self.edges {
e.var.hash(state);
e.src.hash(state);
e.dst.hash(state);
e.rel_type_id.hash(state);
e.rel_type_name.hash(state);
e.direction.hash(state);
hash_props(&e.properties, state);
hash_computed(&e.computed_properties, state);
}
self.dir.hash(state);
self.mode.hash(state);
self.emit_rows.hash(state);
}
}
impl UserDefinedLogicalNodeCore for GraphCreateNode {
fn name(&self) -> &str {
"GraphCreate"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
self.nodes
.iter()
.flat_map(|n| n.computed_properties.iter())
.chain(self.edges.iter().flat_map(|e| e.computed_properties.iter()))
.map(|(_, e)| e.clone())
.collect()
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"GraphCreate: nodes={}, edges={}",
self.nodes.len(),
self.edges.len()
)
}
fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DfResult<Self> {
let input = Arc::new(
inputs
.into_iter()
.next()
.unwrap_or_else(|| (*self.input).clone()),
);
let mut nodes = self.nodes.clone();
let mut edges = self.edges.clone();
let total: usize = nodes
.iter()
.map(|n| n.computed_properties.len())
.sum::<usize>()
+ edges
.iter()
.map(|e| e.computed_properties.len())
.sum::<usize>();
if exprs.len() == total {
let mut it = exprs.into_iter();
for n in &mut nodes {
for (_, e) in &mut n.computed_properties {
*e = it.next().expect("count checked above");
}
}
for e in &mut edges {
for (_, x) in &mut e.computed_properties {
*x = it.next().expect("count checked above");
}
}
}
if self.emit_rows {
Ok(Self::new_emitting(
input,
nodes,
edges,
self.dir.clone(),
self.mode,
self.schema.clone(),
))
} else {
Ok(Self::new(input, nodes, edges, self.dir.clone(), self.mode))
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct DeleteTarget {
pub var: u32,
pub is_edge: bool,
}
#[derive(Debug, Clone)]
pub struct GraphDeleteNode {
pub input: Arc<LogicalPlan>,
pub targets: Vec<DeleteTarget>,
pub detach: bool,
pub dir: PathBuf,
pub mode: OntologyMode,
schema: DFSchemaRef,
}
impl GraphDeleteNode {
#[must_use]
pub fn summary_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("nodes_deleted", DataType::UInt64, false),
Field::new("edges_deleted", DataType::UInt64, false),
]))
}
#[must_use]
pub fn new(
input: Arc<LogicalPlan>,
targets: Vec<DeleteTarget>,
detach: bool,
dir: PathBuf,
mode: OntologyMode,
) -> Self {
let schema = Arc::new(
DFSchema::try_from(Self::summary_schema())
.expect("write-summary schema is always valid"),
);
Self {
input,
targets,
detach,
dir,
mode,
schema,
}
}
}
impl_partial_ord!(GraphDeleteNode);
impl PartialEq for GraphDeleteNode {
fn eq(&self, other: &Self) -> bool {
self.input == other.input
&& self.targets == other.targets
&& self.detach == other.detach
&& self.dir == other.dir
&& self.mode == other.mode
}
}
impl Eq for GraphDeleteNode {}
impl Hash for GraphDeleteNode {
fn hash<H: Hasher>(&self, state: &mut H) {
self.input.hash(state);
self.targets.hash(state);
self.detach.hash(state);
self.dir.hash(state);
self.mode.hash(state);
}
}
impl UserDefinedLogicalNodeCore for GraphDeleteNode {
fn name(&self) -> &str {
"GraphDelete"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
vec![]
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"GraphDelete: targets={}, detach={}",
self.targets.len(),
self.detach
)
}
fn with_exprs_and_inputs(&self, _exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DfResult<Self> {
let input = Arc::new(
inputs
.into_iter()
.next()
.unwrap_or_else(|| (*self.input).clone()),
);
Ok(Self::new(
input,
self.targets.clone(),
self.detach,
self.dir.clone(),
self.mode,
))
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct SetTarget {
pub var: u32,
pub is_edge: bool,
pub prop_name: String,
pub value: Expr,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RemoveTarget {
pub var: u32,
pub is_edge: bool,
pub prop_name: String,
}
#[derive(Debug, Clone)]
pub struct GraphSetNode {
pub input: Arc<LogicalPlan>,
pub targets: Vec<SetTarget>,
pub type_id_to_entity_name: HashMap<u32, String>,
pub dir: PathBuf,
pub mode: OntologyMode,
schema: DFSchemaRef,
}
impl GraphSetNode {
#[must_use]
pub fn summary_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![Field::new(
"properties_set",
DataType::UInt64,
false,
)]))
}
#[must_use]
pub fn new(
input: Arc<LogicalPlan>,
targets: Vec<SetTarget>,
type_id_to_entity_name: HashMap<u32, String>,
dir: PathBuf,
mode: OntologyMode,
) -> Self {
let schema = Arc::new(
DFSchema::try_from(Self::summary_schema())
.expect("write-summary schema is always valid"),
);
Self {
input,
targets,
type_id_to_entity_name,
dir,
mode,
schema,
}
}
}
impl_partial_ord!(GraphSetNode);
impl PartialEq for GraphSetNode {
fn eq(&self, other: &Self) -> bool {
self.input == other.input
&& self.targets == other.targets
&& self.dir == other.dir
&& self.mode == other.mode
&& self.type_id_to_entity_name == other.type_id_to_entity_name
}
}
impl Eq for GraphSetNode {}
impl Hash for GraphSetNode {
fn hash<H: Hasher>(&self, state: &mut H) {
self.input.hash(state);
for t in &self.targets {
t.var.hash(state);
t.is_edge.hash(state);
t.prop_name.hash(state);
}
self.dir.hash(state);
self.mode.hash(state);
}
}
impl UserDefinedLogicalNodeCore for GraphSetNode {
fn name(&self) -> &str {
"GraphSet"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
self.targets.iter().map(|t| t.value.clone()).collect()
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "GraphSet: targets={}", self.targets.len())
}
fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DfResult<Self> {
let input = Arc::new(
inputs
.into_iter()
.next()
.unwrap_or_else(|| (*self.input).clone()),
);
let mut targets = self.targets.clone();
if exprs.len() == targets.len() {
for (t, e) in targets.iter_mut().zip(exprs) {
t.value = e;
}
}
Ok(Self::new(
input,
targets,
self.type_id_to_entity_name.clone(),
self.dir.clone(),
self.mode,
))
}
}
#[derive(Debug, Clone)]
pub struct GraphRemoveNode {
pub input: Arc<LogicalPlan>,
pub targets: Vec<RemoveTarget>,
pub type_id_to_entity_name: HashMap<u32, String>,
pub dir: PathBuf,
pub mode: OntologyMode,
schema: DFSchemaRef,
}
impl GraphRemoveNode {
#[must_use]
pub fn summary_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![Field::new(
"properties_removed",
DataType::UInt64,
false,
)]))
}
#[must_use]
pub fn new(
input: Arc<LogicalPlan>,
targets: Vec<RemoveTarget>,
type_id_to_entity_name: HashMap<u32, String>,
dir: PathBuf,
mode: OntologyMode,
) -> Self {
let schema = Arc::new(
DFSchema::try_from(Self::summary_schema())
.expect("write-summary schema is always valid"),
);
Self {
input,
targets,
type_id_to_entity_name,
dir,
mode,
schema,
}
}
}
impl_partial_ord!(GraphRemoveNode);
impl PartialEq for GraphRemoveNode {
fn eq(&self, other: &Self) -> bool {
self.input == other.input
&& self.targets == other.targets
&& self.dir == other.dir
&& self.mode == other.mode
&& self.type_id_to_entity_name == other.type_id_to_entity_name
}
}
impl Eq for GraphRemoveNode {}
impl Hash for GraphRemoveNode {
fn hash<H: Hasher>(&self, state: &mut H) {
self.input.hash(state);
self.targets.hash(state);
self.dir.hash(state);
self.mode.hash(state);
}
}
impl UserDefinedLogicalNodeCore for GraphRemoveNode {
fn name(&self) -> &str {
"GraphRemove"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
vec![]
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "GraphRemove: targets={}", self.targets.len())
}
fn with_exprs_and_inputs(&self, _exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DfResult<Self> {
let input = Arc::new(
inputs
.into_iter()
.next()
.unwrap_or_else(|| (*self.input).clone()),
);
Ok(Self::new(
input,
self.targets.clone(),
self.type_id_to_entity_name.clone(),
self.dir.clone(),
self.mode,
))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct UnwindNode {
pub input: Arc<LogicalPlan>,
pub list_expr: Expr,
pub alias: String,
element_field: Arc<Field>,
schema: DFSchemaRef,
}
impl UnwindNode {
#[must_use]
pub fn new(
input: Arc<LogicalPlan>,
list_expr: Expr,
alias: impl Into<String>,
element_field: &Field,
) -> Self {
let alias = alias.into();
let element_field = Arc::new(element_field.clone().with_name(&alias).with_nullable(true));
let schema = Self::build_schema(&input, &alias, Arc::clone(&element_field));
Self {
input,
list_expr,
alias,
element_field,
schema,
}
}
fn build_schema(
input: &Arc<LogicalPlan>,
alias: &str,
element_field: Arc<Field>,
) -> DFSchemaRef {
let mut qualified: Vec<(Option<TableReference>, Arc<Field>)> = input
.schema()
.iter()
.map(|(q, f)| (q.cloned(), Arc::clone(f)))
.collect();
if let DataType::Struct(fields) = element_field.data_type()
&& fields.iter().any(|field| {
matches!(
field.name().as_str(),
"node_uuid" | "edge_uuid" | "src_uuid" | "dst_uuid" | "nodes" | "relationships"
)
})
{
let alias_ref = TableReference::bare(alias.to_owned());
qualified.extend(fields.iter().map(|field| {
(
Some(alias_ref.clone()),
Arc::new(field.as_ref().clone().with_nullable(true)),
)
}));
} else {
qualified.push((None, element_field));
}
Arc::new(
DFSchema::new_with_metadata(qualified, std::collections::HashMap::new())
.expect("UnwindNode schema must include the element column"),
)
}
fn element_field(&self) -> Arc<Field> {
Arc::clone(&self.element_field)
}
fn bound_element_field(&self, list_expr: &Expr, input: &LogicalPlan) -> Arc<Field> {
match list_expr.get_type(input.schema().as_ref()) {
Ok(
DataType::List(field)
| DataType::LargeList(field)
| DataType::FixedSizeList(field, _),
) => Arc::new(
field
.as_ref()
.clone()
.with_name(&self.alias)
.with_nullable(true),
),
_ => self.element_field(),
}
}
}
impl_partial_ord!(UnwindNode);
impl UserDefinedLogicalNodeCore for UnwindNode {
fn name(&self) -> &str {
"Unwind"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn expressions(&self) -> Vec<Expr> {
vec![self.list_expr.clone()]
}
fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "Unwind: alias={}", self.alias)
}
fn with_exprs_and_inputs(
&self,
mut exprs: Vec<Expr>,
inputs: Vec<LogicalPlan>,
) -> DfResult<Self> {
let list_expr = exprs.pop().unwrap_or_else(|| self.list_expr.clone());
let input = Arc::new(
inputs
.into_iter()
.next()
.unwrap_or_else(|| (*self.input).clone()),
);
let element_field = self.bound_element_field(&list_expr, &input);
Ok(Self::new(
input,
list_expr,
self.alias.clone(),
&element_field,
))
}
}
#[cfg(test)]
mod tests {
use super::*;
use datafusion::logical_expr::{Extension, LogicalPlan};
fn empty_plan() -> Arc<LogicalPlan> {
use datafusion::logical_expr::LogicalPlanBuilder;
Arc::new(LogicalPlanBuilder::empty(false).build().unwrap())
}
fn wrap(node: impl UserDefinedLogicalNodeCore) -> LogicalPlan {
LogicalPlan::Extension(Extension {
node: Arc::new(node),
})
}
fn dst_node_fields() -> Vec<Arc<Field>> {
vec![
Arc::new(Field::new("node_id", DataType::UInt64, false)),
Arc::new(Field::new("type_id", DataType::UInt32, false)),
]
}
fn var_len(rel: &str, min_hops: u16, max_hops: Option<u16>) -> VarLenExpandNode {
VarLenExpandNode::new(
empty_plan(),
rel,
min_hops,
max_hops,
0,
1,
2,
Direction::Out,
Some(7),
PathBuf::from("/tmp/gf"),
OntologyMode::Strict,
dst_node_fields(),
var_len_edge_list_field(&[]),
)
}
fn scan_plan(var: u32, fields: Vec<Field>) -> Arc<LogicalPlan> {
use datafusion::logical_expr::LogicalPlanBuilder;
use datafusion::logical_expr::logical_plan::LogicalTableSource;
let schema = Arc::new(Schema::new(fields));
Arc::new(
LogicalPlanBuilder::scan(
format!("var_{var}"),
Arc::new(LogicalTableSource::new(schema)),
None,
)
.unwrap()
.build()
.unwrap(),
)
}
fn opt_match(
optional: Arc<LogicalPlan>,
join_keys: Vec<(usize, usize)>,
inner_keep_idx: Vec<usize>,
) -> OptionalMatchNode {
OptionalMatchNode::new(empty_plan(), optional, join_keys, inner_keep_idx)
}
fn unwind(list_expr: datafusion::logical_expr::Expr, alias: &str) -> UnwindNode {
let element_field = Field::new("elem", DataType::Int64, true);
UnwindNode::new(empty_plan(), list_expr, alias, &element_field)
}
#[test]
fn var_len_expand_name() {
let n = var_len("KNOWS", 1, Some(3));
assert_eq!(UserDefinedLogicalNodeCore::name(&n), "VarLenExpand");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&n).len(), 1);
}
#[test]
fn var_len_expand_schema_extends_input_with_dst_and_edge_fields() {
use datafusion::common::TableReference;
let n = var_len("KNOWS", 1, Some(2));
let schema = UserDefinedLogicalNodeCore::schema(&n);
assert_eq!(schema.fields().len(), 3);
let dst = TableReference::bare("var_1");
assert!(schema.field_with_qualified_name(&dst, "node_id").is_ok());
assert!(schema.field_with_qualified_name(&dst, "type_id").is_ok());
let edge = TableReference::bare("var_2");
let edge_field = schema
.field_with_qualified_name(&edge, VAR_LEN_EDGE_LIST_FIELD)
.expect("edge-list column present");
assert!(matches!(edge_field.data_type(), DataType::List(_)));
}
#[test]
fn optional_match_name() {
let n = opt_match(empty_plan(), vec![], vec![]);
assert_eq!(UserDefinedLogicalNodeCore::name(&n), "OptionalMatch");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&n).len(), 2);
}
#[test]
fn optional_match_schema_appends_nullable_inner_fields() {
let optional = scan_plan(
1,
vec![
Field::new("node_id", DataType::UInt64, false),
Field::new("type_id", DataType::UInt32, false),
],
);
let n = opt_match(optional, vec![], vec![0, 1]);
let schema = UserDefinedLogicalNodeCore::schema(&n);
assert_eq!(schema.fields().len(), 2);
assert!(schema.field(0).is_nullable());
assert!(schema.field(1).is_nullable());
let var1 = TableReference::bare("var_1");
assert!(schema.field_with_qualified_name(&var1, "node_id").is_ok());
}
#[test]
fn optional_match_schema_excludes_shared_var_columns() {
let optional = scan_plan(
7,
vec![
Field::new("node_id", DataType::UInt64, false),
Field::new("type_id", DataType::UInt32, false),
Field::new("payload", DataType::UInt64, false),
],
);
let n = opt_match(optional, vec![(0, 0)], vec![2]);
let schema = UserDefinedLogicalNodeCore::schema(&n);
assert_eq!(schema.fields().len(), 1, "only the non-shared col is kept");
assert_eq!(schema.field(0).name(), "payload");
assert!(schema.field(0).is_nullable());
}
#[test]
fn path_unique_name() {
let n = PathUniqueNode::new(empty_plan());
assert_eq!(UserDefinedLogicalNodeCore::name(&n), "PathUnique");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&n).len(), 1);
}
#[test]
fn ontology_infer_name() {
let n = OntologyInferNode::new(
empty_plan(),
"MANAGES",
"transitive:MANAGES",
"conservative_min",
);
assert_eq!(UserDefinedLogicalNodeCore::name(&n), "OntologyInfer");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&n).len(), 1);
}
#[test]
fn graph_merge_name() {
let n = GraphMergeNode::new();
assert_eq!(UserDefinedLogicalNodeCore::name(&n), "GraphMerge");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&n).len(), 0);
}
#[test]
fn unwind_name() {
use datafusion::logical_expr::lit;
let n = unwind(lit(1i64), "x");
assert_eq!(UserDefinedLogicalNodeCore::name(&n), "Unwind");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&n).len(), 1);
assert_eq!(UserDefinedLogicalNodeCore::expressions(&n).len(), 1);
}
#[test]
fn unwind_schema_appends_element_column() {
use datafusion::logical_expr::lit;
let n = unwind(lit(1i64), "x");
let schema = UserDefinedLogicalNodeCore::schema(&n);
assert_eq!(schema.fields().len(), 1);
assert!(
schema.field(0).is_nullable(),
"unwound element column is nullable"
);
assert_eq!(schema.field(0).name(), "x");
assert!(schema.field_with_unqualified_name("x").is_ok());
}
#[test]
fn graph_create_name_and_wrap() {
let node = GraphCreateNode::new(
empty_plan(),
vec![ResolvedNodeSpec {
var: 0,
label_ids: vec![3],
label_names: vec!["Person".to_owned()],
properties: vec![("name".to_owned(), IrLiteral::Str("Alice".to_owned()))],
computed_properties: vec![],
is_reference: false,
}],
vec![],
PathBuf::from("/tmp/gf"),
OntologyMode::Strict,
);
assert_eq!(UserDefinedLogicalNodeCore::name(&node), "GraphCreate");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&node).len(), 1);
assert_eq!(UserDefinedLogicalNodeCore::schema(&node).fields().len(), 4);
assert!(matches!(wrap(node), LogicalPlan::Extension(_)));
}
#[test]
fn graph_create_eq_and_hash_are_consistent() {
use std::collections::hash_map::DefaultHasher;
let mk = || {
GraphCreateNode::new(
empty_plan(),
vec![ResolvedNodeSpec {
var: 0,
label_ids: vec![],
label_names: vec![],
properties: vec![("score".to_owned(), IrLiteral::Float(1.5))],
computed_properties: vec![],
is_reference: false,
}],
vec![],
PathBuf::from("/tmp/gf"),
OntologyMode::Exploratory,
)
};
let (a, b) = (mk(), mk());
assert_eq!(a, b);
let mut ha = DefaultHasher::new();
let mut hb = DefaultHasher::new();
a.hash(&mut ha);
b.hash(&mut hb);
assert_eq!(ha.finish(), hb.finish());
}
#[test]
fn graph_create_hashes_every_literal_shape_deterministically() {
use std::collections::hash_map::DefaultHasher;
let literals = vec![
IrLiteral::Null,
IrLiteral::Bool(true),
IrLiteral::Int(7),
IrLiteral::Float(1.5),
IrLiteral::Str("value".into()),
IrLiteral::Uuid([1; 16]),
IrLiteral::Duration {
months: 1,
days: 2,
seconds: 3,
nanos: 4,
},
IrLiteral::DateTime(5),
IrLiteral::Date(6),
IrLiteral::LocalDateTime { days: 7, nanos: 8 },
IrLiteral::Time(9),
IrLiteral::ZonedTime {
nanos: 10,
offset: 3_600,
},
IrLiteral::ZonedDateTime {
days: 11,
nanos: 12,
offset: -3_600,
zone: Some("UTC".into()),
},
IrLiteral::List(vec![IrLiteral::Int(13)]),
IrLiteral::Map(vec![("key".into(), IrLiteral::Int(14))]),
];
let make = || {
GraphCreateNode::new(
empty_plan(),
vec![ResolvedNodeSpec {
var: 0,
label_ids: vec![1],
label_names: vec!["Person".into()],
properties: literals
.iter()
.cloned()
.enumerate()
.map(|(index, literal)| (index.to_string(), literal))
.collect(),
computed_properties: vec![],
is_reference: false,
}],
vec![],
PathBuf::from("/tmp/gf"),
OntologyMode::Strict,
)
};
let (first, second) = (make(), make());
let mut first_hash = DefaultHasher::new();
let mut second_hash = DefaultHasher::new();
first.hash(&mut first_hash);
second.hash(&mut second_hash);
assert_eq!(first, second);
assert_eq!(first_hash.finish(), second_hash.finish());
}
fn set_node(value: Expr) -> GraphSetNode {
GraphSetNode::new(
empty_plan(),
vec![SetTarget {
var: 0,
is_edge: false,
prop_name: "age".to_owned(),
value,
}],
HashMap::new(),
PathBuf::from("/tmp/gf"),
OntologyMode::Exploratory,
)
}
fn remove_node() -> GraphRemoveNode {
GraphRemoveNode::new(
empty_plan(),
vec![RemoveTarget {
var: 0,
is_edge: false,
prop_name: "age".to_owned(),
}],
HashMap::new(),
PathBuf::from("/tmp/gf"),
OntologyMode::Exploratory,
)
}
#[test]
fn graph_set_name_schema_and_expr_surface() {
use datafusion::logical_expr::lit;
let n = set_node(lit(42i64));
assert_eq!(UserDefinedLogicalNodeCore::name(&n), "GraphSet");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&n).len(), 1);
assert_eq!(UserDefinedLogicalNodeCore::schema(&n).fields().len(), 1);
assert_eq!(UserDefinedLogicalNodeCore::expressions(&n).len(), 1);
assert!(matches!(wrap(n), LogicalPlan::Extension(_)));
}
#[test]
fn graph_set_round_trips_value_expr() {
use datafusion::logical_expr::{col, lit};
let value = col("var_0.age") + lit(1i64);
let n = set_node(value.clone());
let exprs = UserDefinedLogicalNodeCore::expressions(&n);
assert_eq!(exprs, vec![value.clone()]);
let rebuilt = UserDefinedLogicalNodeCore::with_exprs_and_inputs(&n, exprs, vec![]).unwrap();
assert_eq!(rebuilt.targets[0].value, value);
assert_eq!(n, rebuilt);
}
#[test]
fn graph_set_eq_distinguishes_value_expr() {
use datafusion::logical_expr::lit;
assert_ne!(set_node(lit(1i64)), set_node(lit(2i64)));
}
#[test]
fn graph_remove_name_schema_and_no_exprs() {
let n = remove_node();
assert_eq!(UserDefinedLogicalNodeCore::name(&n), "GraphRemove");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&n).len(), 1);
assert_eq!(UserDefinedLogicalNodeCore::schema(&n).fields().len(), 1);
assert!(UserDefinedLogicalNodeCore::expressions(&n).is_empty());
assert!(matches!(wrap(n), LogicalPlan::Extension(_)));
}
#[test]
fn graph_remove_eq_and_hash_are_consistent() {
use std::collections::hash_map::DefaultHasher;
let (a, b) = (remove_node(), remove_node());
assert_eq!(a, b);
let mut ha = DefaultHasher::new();
let mut hb = DefaultHasher::new();
a.hash(&mut ha);
b.hash(&mut hb);
assert_eq!(ha.finish(), hb.finish());
}
#[test]
fn graph_create_rebuilds_computed_properties_and_preserves_emit_mode() {
use datafusion::logical_expr::lit;
let output_schema = Arc::new(DFSchema::empty());
let node = GraphCreateNode::new_emitting(
empty_plan(),
vec![ResolvedNodeSpec {
var: 0,
label_ids: vec![1],
label_names: vec!["Person".into()],
properties: vec![],
computed_properties: vec![("score".into(), lit(1_i64))],
is_reference: false,
}],
vec![ResolvedEdgeSpec {
var: 1,
src: 0,
dst: 2,
rel_type_id: Some(3),
rel_type_name: Some("KNOWS".into()),
direction: Direction::Out,
properties: vec![],
computed_properties: vec![("weight".into(), lit(2_i64))],
}],
PathBuf::from("/tmp/gf"),
OntologyMode::Strict,
output_schema.clone(),
);
assert!(node.emits_rows());
assert_eq!(UserDefinedLogicalNodeCore::expressions(&node).len(), 2);
let rebuilt = UserDefinedLogicalNodeCore::with_exprs_and_inputs(
&node,
vec![lit(10_i64), lit(20_i64)],
vec![(*empty_plan()).clone()],
)
.unwrap();
assert!(rebuilt.emits_rows());
assert_eq!(rebuilt.schema(), &output_schema);
assert_eq!(rebuilt.nodes[0].computed_properties[0].1, lit(10_i64));
assert_eq!(rebuilt.edges[0].computed_properties[0].1, lit(20_i64));
let short =
UserDefinedLogicalNodeCore::with_exprs_and_inputs(&node, vec![lit(99_i64)], vec![])
.unwrap();
assert_eq!(short.nodes[0].computed_properties[0].1, lit(1_i64));
assert_eq!(short.edges[0].computed_properties[0].1, lit(2_i64));
}
#[test]
fn graph_delete_contract_round_trips_and_distinguishes_detach() {
use std::collections::hash_map::DefaultHasher;
let node = GraphDeleteNode::new(
empty_plan(),
vec![DeleteTarget {
var: 7,
is_edge: false,
}],
true,
PathBuf::from("/tmp/gf"),
OntologyMode::Exploratory,
);
assert_eq!(UserDefinedLogicalNodeCore::name(&node), "GraphDelete");
assert_eq!(UserDefinedLogicalNodeCore::inputs(&node).len(), 1);
assert_eq!(UserDefinedLogicalNodeCore::schema(&node).fields().len(), 2);
assert!(UserDefinedLogicalNodeCore::expressions(&node).is_empty());
let rebuilt = UserDefinedLogicalNodeCore::with_exprs_and_inputs(
&node,
vec![],
vec![(*empty_plan()).clone()],
)
.unwrap();
assert_eq!(node, rebuilt);
assert_ne!(
node,
GraphDeleteNode::new(
empty_plan(),
node.targets.clone(),
false,
node.dir.clone(),
node.mode,
)
);
let mut first = DefaultHasher::new();
let mut second = DefaultHasher::new();
node.hash(&mut first);
rebuilt.hash(&mut second);
assert_eq!(first.finish(), second.finish());
}
#[test]
fn optimizer_nodes_are_intentionally_not_orderable() {
use datafusion::logical_expr::lit;
let set = set_node(lit(1_i64));
let remove = remove_node();
let create = GraphCreateNode::new(
empty_plan(),
vec![],
vec![],
PathBuf::from("/tmp/gf"),
OntologyMode::Strict,
);
let delete = GraphDeleteNode::new(
empty_plan(),
vec![],
false,
PathBuf::from("/tmp/gf"),
OntologyMode::Strict,
);
assert_eq!(create.partial_cmp(&create), None);
assert_eq!(delete.partial_cmp(&delete), None);
assert_eq!(set.partial_cmp(&set), None);
assert_eq!(remove.partial_cmp(&remove), None);
}
#[test]
fn all_stubs_wrap_as_extension() {
use datafusion::logical_expr::lit;
let plans = vec![
wrap(var_len("*", 1, None)),
wrap(opt_match(empty_plan(), vec![], vec![])),
wrap(PathUniqueNode::new(empty_plan())),
wrap(OntologyInferNode::new(
empty_plan(),
"REL",
"transitive:REL",
"conservative_min",
)),
wrap(GraphMergeNode::new()),
wrap(unwind(lit(1i64), "x")),
wrap(set_node(lit(1i64))),
wrap(remove_node()),
];
for p in &plans {
assert!(
matches!(p, LogicalPlan::Extension(_)),
"expected Extension, got {p:?}"
);
}
}
#[test]
fn fmt_for_explain_is_non_empty() {
use datafusion::logical_expr::lit;
struct ExplainWrapper<'a, T: UserDefinedLogicalNodeCore>(&'a T);
impl<T: UserDefinedLogicalNodeCore> fmt::Display for ExplainWrapper<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt_for_explain(f)
}
}
fn explain<T: UserDefinedLogicalNodeCore>(n: &T) -> String {
ExplainWrapper(n).to_string()
}
assert!(!explain(&var_len("KNOWS", 1, Some(3))).is_empty());
assert!(!explain(&opt_match(empty_plan(), vec![], vec![])).is_empty());
assert!(!explain(&PathUniqueNode::new(empty_plan())).is_empty());
assert!(
!explain(&OntologyInferNode::new(
empty_plan(),
"MANAGES",
"transitive:MANAGES",
"conservative_min"
))
.is_empty()
);
assert!(!explain(&GraphMergeNode::new()).is_empty());
assert!(!explain(&unwind(lit(1i64), "x")).is_empty());
assert!(explain(&var_len("KNOWS", 1, Some(3))).contains("KNOWS"));
assert!(
explain(&OntologyInferNode::new(
empty_plan(),
"MANAGES",
"transitive:MANAGES",
"conservative_min"
))
.contains("MANAGES")
);
assert!(explain(&unwind(lit(1i64), "x")).contains("x"));
}
#[test]
fn extension_stub_trait_contracts_cover_rebuilds_directions_and_hashing() {
use datafusion::logical_expr::lit;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut incoming = var_len("KNOWS", 1, None);
incoming.direction = Direction::In;
let mut undirected = var_len("KNOWS", 1, None);
undirected.direction = Direction::Undirected;
struct Explain<'a, T: UserDefinedLogicalNodeCore>(&'a T);
impl<T: UserDefinedLogicalNodeCore> fmt::Display for Explain<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt_for_explain(f)
}
}
assert!(Explain(&incoming).to_string().contains("dir=<-"));
assert!(Explain(&undirected).to_string().contains("dir=--"));
let path_unique = PathUniqueNode::new(empty_plan());
assert!(
UserDefinedLogicalNodeCore::schema(&path_unique)
.fields()
.is_empty()
);
assert!(UserDefinedLogicalNodeCore::expressions(&path_unique).is_empty());
let rebuilt = path_unique
.with_exprs_and_inputs(vec![], vec![empty_plan().as_ref().clone()])
.unwrap();
assert_eq!(UserDefinedLogicalNodeCore::name(&rebuilt), "PathUnique");
let merge = GraphMergeNode::default();
assert!(
UserDefinedLogicalNodeCore::schema(&merge)
.fields()
.is_empty()
);
assert!(UserDefinedLogicalNodeCore::expressions(&merge).is_empty());
assert_eq!(
UserDefinedLogicalNodeCore::name(&merge.with_exprs_and_inputs(vec![], vec![]).unwrap()),
"GraphMerge"
);
let create = GraphCreateNode::new(
empty_plan(),
vec![],
vec![],
PathBuf::from("/tmp/gf"),
OntologyMode::Strict,
);
let delete = GraphDeleteNode::new(
empty_plan(),
vec![],
true,
PathBuf::from("/tmp/gf"),
OntologyMode::Strict,
);
let set = set_node(lit(1_i64));
let remove = remove_node();
for rendered in [
Explain(&create).to_string(),
Explain(&delete).to_string(),
Explain(&set).to_string(),
Explain(&remove).to_string(),
] {
assert!(!rendered.is_empty());
}
for hashable in [&create as &dyn HashProbe, &set as &dyn HashProbe] {
let mut hasher = DefaultHasher::new();
hashable.hash_into(&mut hasher);
let _ = hasher.finish();
}
trait HashProbe {
fn hash_into(&self, state: &mut DefaultHasher);
}
impl<T: Hash> HashProbe for T {
fn hash_into(&self, state: &mut DefaultHasher) {
self.hash(state);
}
}
}
fn expand_node(edge_prop_fields: Vec<Arc<Field>>) -> ExpandNode {
let edge_fields = vec![
Arc::new(Field::new(
"edge_uuid",
DataType::FixedSizeBinary(16),
false,
)),
Arc::new(Field::new("edge_id", DataType::UInt64, false)),
Arc::new(Field::new("src_id", DataType::UInt64, false)),
Arc::new(Field::new("dst_id", DataType::UInt64, false)),
];
ExpandNode::new(
empty_plan(),
"KNOWS",
0,
2,
1,
Direction::Out,
Some(7),
PathBuf::from("/tmp/p"),
OntologyMode::Strict,
edge_fields,
edge_prop_fields,
dst_node_fields(),
)
}
#[test]
fn expand_node_schema_order_qualifiers_and_prop_nullability() {
use datafusion::common::TableReference;
let node = expand_node(vec![Arc::new(Field::new("since", DataType::Int64, false))]);
let schema = UserDefinedLogicalNodeCore::schema(&node);
let names: Vec<String> = schema
.iter()
.map(|(q, f)| {
format!(
"{}.{}",
q.map(TableReference::table).unwrap_or("?"),
f.name()
)
})
.collect();
assert_eq!(
names,
[
"var_1.edge_uuid",
"var_1.edge_id",
"var_1.src_id",
"var_1.dst_id",
"var_1.since",
"var_2.node_id",
"var_2.type_id",
]
);
let edge = TableReference::bare("var_1");
let since = schema.field_with_qualified_name(&edge, "since").unwrap();
assert!(since.is_nullable(), "prop columns are LEFT-join nullable");
let uuid = schema
.field_with_qualified_name(&edge, "edge_uuid")
.unwrap();
assert!(!uuid.is_nullable(), "topology columns keep non-null");
assert_eq!(node.edge_prop_count, 1);
}
#[test]
fn expand_node_with_exprs_and_inputs_round_trips() {
let node = expand_node(vec![Arc::new(Field::new("since", DataType::Int64, true))]);
let rebuilt = UserDefinedLogicalNodeCore::with_exprs_and_inputs(
&node,
vec![],
vec![(*empty_plan()).clone()],
)
.unwrap();
assert_eq!(
UserDefinedLogicalNodeCore::schema(&rebuilt).as_arrow(),
UserDefinedLogicalNodeCore::schema(&node).as_arrow(),
"schema survives optimizer-style input replacement"
);
assert_eq!(rebuilt.edge_prop_count, 1);
assert_eq!(rebuilt.rel_type_name, "KNOWS");
}
}