use std::fmt::Write;
use crate::{Predicate, QueryAst, QueryStep, TraverseDirection};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DrivingTable {
Nodes,
FtsNodes,
VecNodes,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ExecutionHints {
pub recursion_limit: usize,
pub hard_limit: usize,
}
pub fn choose_driving_table(ast: &QueryAst) -> DrivingTable {
let has_deterministic_id_filter = ast.steps.iter().any(|step| {
matches!(
step,
QueryStep::Filter(Predicate::LogicalIdEq(_) | Predicate::SourceRefEq(_))
)
});
if has_deterministic_id_filter {
DrivingTable::Nodes
} else if ast
.steps
.iter()
.any(|step| matches!(step, QueryStep::VectorSearch { .. }))
{
DrivingTable::VecNodes
} else if ast
.steps
.iter()
.any(|step| matches!(step, QueryStep::TextSearch { .. }))
{
DrivingTable::FtsNodes
} else {
DrivingTable::Nodes
}
}
pub fn execution_hints(ast: &QueryAst) -> ExecutionHints {
let step_limit = ast
.steps
.iter()
.find_map(|step| {
if let QueryStep::Traverse { max_depth, .. } = step {
Some(*max_depth)
} else {
None
}
})
.unwrap_or(0);
let expansion_limit = ast
.expansions
.iter()
.map(|expansion| expansion.max_depth)
.max()
.unwrap_or(0);
let recursion_limit = step_limit.max(expansion_limit);
ExecutionHints {
recursion_limit,
hard_limit: ast.final_limit.unwrap_or(1000),
}
}
#[allow(clippy::too_many_lines)]
pub fn shape_signature(ast: &QueryAst) -> String {
let mut signature = String::new();
let _ = write!(&mut signature, "Root({})", ast.root_kind);
for step in &ast.steps {
match step {
QueryStep::Search { limit, .. } => {
let _ = write!(&mut signature, "-Search(limit={limit})");
}
QueryStep::VectorSearch { limit, .. } => {
let _ = write!(&mut signature, "-Vector(limit={limit})");
}
QueryStep::TextSearch { limit, .. } => {
let _ = write!(&mut signature, "-Text(limit={limit})");
}
QueryStep::Traverse {
direction,
label,
max_depth,
filter: _,
} => {
let dir = match direction {
TraverseDirection::In => "in",
TraverseDirection::Out => "out",
};
let _ = write!(
&mut signature,
"-Traverse(direction={dir},label={label},depth={max_depth})"
);
}
QueryStep::Filter(predicate) => match predicate {
Predicate::LogicalIdEq(_) => signature.push_str("-Filter(logical_id_eq)"),
Predicate::KindEq(_) => signature.push_str("-Filter(kind_eq)"),
Predicate::JsonPathEq { path, .. } => {
let _ = write!(&mut signature, "-Filter(json_eq:{path})");
}
Predicate::JsonPathCompare { path, op, .. } => {
let op = match op {
crate::ComparisonOp::Gt => "gt",
crate::ComparisonOp::Gte => "gte",
crate::ComparisonOp::Lt => "lt",
crate::ComparisonOp::Lte => "lte",
};
let _ = write!(&mut signature, "-Filter(json_cmp:{path}:{op})");
}
Predicate::SourceRefEq(_) => signature.push_str("-Filter(source_ref_eq)"),
Predicate::ContentRefNotNull => {
signature.push_str("-Filter(content_ref_not_null)");
}
Predicate::ContentRefEq(_) => signature.push_str("-Filter(content_ref_eq)"),
Predicate::JsonPathFusedEq { path, .. } => {
let _ = write!(&mut signature, "-Filter(json_fused_eq:{path})");
}
Predicate::JsonPathFusedTimestampCmp { path, op, .. } => {
let op = match op {
crate::ComparisonOp::Gt => "gt",
crate::ComparisonOp::Gte => "gte",
crate::ComparisonOp::Lt => "lt",
crate::ComparisonOp::Lte => "lte",
};
let _ = write!(&mut signature, "-Filter(json_fused_ts_cmp:{path}:{op})");
}
Predicate::JsonPathFusedBoolEq { path, .. } => {
let _ = write!(&mut signature, "-Filter(json_fused_bool_eq:{path})");
}
Predicate::EdgePropertyEq { path, .. } => {
let _ = write!(&mut signature, "-Filter(edge_eq:{path})");
}
Predicate::EdgePropertyCompare { path, op, .. } => {
let op = match op {
crate::ComparisonOp::Gt => "gt",
crate::ComparisonOp::Gte => "gte",
crate::ComparisonOp::Lt => "lt",
crate::ComparisonOp::Lte => "lte",
};
let _ = write!(&mut signature, "-Filter(edge_cmp:{path}:{op})");
}
Predicate::JsonPathFusedIn { path, values } => {
let _ = write!(
&mut signature,
"-Filter(json_fused_in:{path}:n={})",
values.len()
);
}
Predicate::JsonPathIn { path, values } => {
let _ = write!(&mut signature, "-Filter(json_in:{path}:n={})", values.len());
}
},
}
}
for expansion in &ast.expansions {
let dir = match expansion.direction {
TraverseDirection::In => "in",
TraverseDirection::Out => "out",
};
let edge_filter_str = match &expansion.edge_filter {
None => String::new(),
Some(Predicate::EdgePropertyEq { path, .. }) => {
format!(",edge_eq:{path}")
}
Some(Predicate::EdgePropertyCompare { path, op, .. }) => {
let op_str = match op {
crate::ComparisonOp::Gt => "gt",
crate::ComparisonOp::Gte => "gte",
crate::ComparisonOp::Lt => "lt",
crate::ComparisonOp::Lte => "lte",
};
format!(",edge_cmp:{path}:{op_str}")
}
Some(p) => unreachable!("edge_filter predicate {p:?} not handled in shape_signature"),
};
let _ = write!(
&mut signature,
"-Expand(slot={},direction={dir},label={},depth={}{})",
expansion.slot, expansion.label, expansion.max_depth, edge_filter_str
);
}
for edge_expansion in &ast.edge_expansions {
let dir = match edge_expansion.direction {
TraverseDirection::In => "in",
TraverseDirection::Out => "out",
};
let _ = write!(
&mut signature,
"-EdgeExpand(slot={},direction={dir},label={},depth={},edge_filter={},endpoint_filter={})",
edge_expansion.slot,
edge_expansion.label,
edge_expansion.max_depth,
edge_expansion.edge_filter.is_some(),
edge_expansion.endpoint_filter.is_some(),
);
}
if let Some(limit) = ast.final_limit {
let _ = write!(&mut signature, "-Limit({limit})");
}
signature
}
#[cfg(test)]
mod tests {
use crate::{DrivingTable, QueryBuilder, TraverseDirection};
use super::{choose_driving_table, execution_hints};
#[test]
fn deterministic_filter_overrides_vector_driver() {
let ast = QueryBuilder::nodes("Meeting")
.vector_search("budget", 5)
.filter_logical_id_eq("meeting-123")
.into_ast();
assert_eq!(choose_driving_table(&ast), DrivingTable::Nodes);
}
#[test]
fn hard_limit_honors_user_specified_limit_below_default() {
let ast = QueryBuilder::nodes("Meeting")
.traverse(TraverseDirection::Out, "HAS_TASK", 3)
.limit(10)
.into_ast();
let hints = execution_hints(&ast);
assert_eq!(
hints.hard_limit, 10,
"hard_limit must honor user's final_limit"
);
}
#[test]
fn hard_limit_defaults_to_1000_when_no_limit_set() {
let ast = QueryBuilder::nodes("Meeting")
.traverse(TraverseDirection::Out, "HAS_TASK", 3)
.into_ast();
let hints = execution_hints(&ast);
assert_eq!(hints.hard_limit, 1000, "hard_limit defaults to 1000");
}
#[test]
fn shape_signature_differs_for_different_edge_filters() {
use crate::{ComparisonOp, ExpansionSlot, Predicate, QueryAst, ScalarValue};
let base_expansion = ExpansionSlot {
slot: "tasks".to_owned(),
direction: TraverseDirection::Out,
label: "HAS_TASK".to_owned(),
max_depth: 1,
filter: None,
edge_filter: None,
};
let ast_no_filter = QueryAst {
root_kind: "Meeting".to_owned(),
steps: vec![],
expansions: vec![base_expansion.clone()],
edge_expansions: vec![],
final_limit: None,
};
let ast_with_eq_filter = QueryAst {
root_kind: "Meeting".to_owned(),
steps: vec![],
expansions: vec![ExpansionSlot {
edge_filter: Some(Predicate::EdgePropertyEq {
path: "$.rel".to_owned(),
value: ScalarValue::Text("cites".to_owned()),
}),
..base_expansion.clone()
}],
edge_expansions: vec![],
final_limit: None,
};
let ast_with_cmp_filter = QueryAst {
root_kind: "Meeting".to_owned(),
steps: vec![],
expansions: vec![ExpansionSlot {
edge_filter: Some(Predicate::EdgePropertyCompare {
path: "$.weight".to_owned(),
op: ComparisonOp::Gt,
value: ScalarValue::Integer(5),
}),
..base_expansion
}],
edge_expansions: vec![],
final_limit: None,
};
let sig_no_filter = super::shape_signature(&ast_no_filter);
let sig_eq_filter = super::shape_signature(&ast_with_eq_filter);
let sig_cmp_filter = super::shape_signature(&ast_with_cmp_filter);
assert_ne!(
sig_no_filter, sig_eq_filter,
"no edge_filter and EdgePropertyEq must produce different signatures"
);
assert_ne!(
sig_no_filter, sig_cmp_filter,
"no edge_filter and EdgePropertyCompare must produce different signatures"
);
assert_ne!(
sig_eq_filter, sig_cmp_filter,
"EdgePropertyEq and EdgePropertyCompare must produce different signatures"
);
}
#[test]
fn shape_signature_differs_for_different_edge_expansions() {
use crate::{EdgeExpansionSlot, Predicate, QueryAst, ScalarValue};
let base = EdgeExpansionSlot {
slot: "cites".to_owned(),
direction: TraverseDirection::Out,
label: "CITES".to_owned(),
max_depth: 1,
endpoint_filter: None,
edge_filter: None,
};
let ast_no_edge_expansions = QueryAst {
root_kind: "Paper".to_owned(),
steps: vec![],
expansions: vec![],
edge_expansions: vec![],
final_limit: None,
};
let ast_with_edge_expansion = QueryAst {
root_kind: "Paper".to_owned(),
steps: vec![],
expansions: vec![],
edge_expansions: vec![base.clone()],
final_limit: None,
};
let ast_with_different_slot_name = QueryAst {
root_kind: "Paper".to_owned(),
steps: vec![],
expansions: vec![],
edge_expansions: vec![EdgeExpansionSlot {
slot: "references".to_owned(),
..base.clone()
}],
final_limit: None,
};
let ast_with_edge_filter = QueryAst {
root_kind: "Paper".to_owned(),
steps: vec![],
expansions: vec![],
edge_expansions: vec![EdgeExpansionSlot {
edge_filter: Some(Predicate::EdgePropertyEq {
path: "$.rel".to_owned(),
value: ScalarValue::Text("primary".to_owned()),
}),
..base.clone()
}],
final_limit: None,
};
let ast_with_endpoint_filter = QueryAst {
root_kind: "Paper".to_owned(),
steps: vec![],
expansions: vec![],
edge_expansions: vec![EdgeExpansionSlot {
endpoint_filter: Some(Predicate::KindEq("Paper".to_owned())),
..base
}],
final_limit: None,
};
let sig_none = super::shape_signature(&ast_no_edge_expansions);
let sig_basic = super::shape_signature(&ast_with_edge_expansion);
let sig_diff_slot = super::shape_signature(&ast_with_different_slot_name);
let sig_edge = super::shape_signature(&ast_with_edge_filter);
let sig_endpoint = super::shape_signature(&ast_with_endpoint_filter);
assert_ne!(
sig_none, sig_basic,
"presence of edge_expansions must change signature"
);
assert_ne!(
sig_basic, sig_diff_slot,
"different edge_expansion slot names must change signature"
);
assert_ne!(
sig_basic, sig_edge,
"edge_filter presence must change signature"
);
assert_ne!(
sig_basic, sig_endpoint,
"endpoint_filter presence must change signature"
);
}
}