use super::super::ast::{Clause, ConstraintCommand, Expression, SchemaCommand};
use crate::datatypes::values::Value;
use crate::graph::core::pattern_matching::PatternElement;
use std::collections::HashMap;
use std::sync::{Arc, OnceLock};
pub(super) struct VectorScoreFilterSpec {
pub(super) variable: String,
pub(super) prop_name: String,
pub(super) query_vec: Vec<f32>,
pub(super) scorer: crate::graph::algorithms::vector::Scorer,
pub(super) threshold: f64,
pub(super) greater_than: bool,
pub(super) inclusive: bool,
}
pub(super) struct DistanceFilterSpec {
pub(super) variable: String,
pub(super) lat_prop: String,
pub(super) lon_prop: String,
pub(super) center_lat: f64,
pub(super) center_lon: f64,
pub(super) threshold: f64,
pub(super) less_than: bool,
pub(super) inclusive: bool,
}
pub(super) struct ContainsFilterSpec {
pub(super) container_variable: String,
pub(super) contained: ContainsTarget,
pub(super) negated: bool,
}
pub(super) enum ContainsTarget {
ConstantPoint(f64, f64),
Variable { name: String },
}
pub(super) enum ResolvedSpatial {
Point(f64, f64),
Geometry(Arc<geo::Geometry<f64>>, Option<geo::Rect<f64>>),
}
pub(super) type GeomWithBBox = (Arc<geo::Geometry<f64>>, Option<geo::Rect<f64>>);
pub(super) struct NodeSpatialData {
pub(super) geometry: Option<GeomWithBBox>,
pub(super) location: Option<(f64, f64)>,
pub(super) shapes: HashMap<String, GeomWithBBox>,
pub(super) points: HashMap<String, (f64, f64)>,
}
pub(super) struct ScoredRowRef {
pub(super) score: f64,
pub(super) index: usize,
}
impl PartialEq for ScoredRowRef {
fn eq(&self, other: &Self) -> bool {
self.index == other.index
}
}
impl Eq for ScoredRowRef {}
impl PartialOrd for ScoredRowRef {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for ScoredRowRef {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
other
.score
.partial_cmp(&self.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| self.index.cmp(&other.index))
}
}
pub(super) struct VectorScoreCache {
pub(super) keys: Option<Vec<ArgKey>>,
pub(super) prop_name: String,
pub(super) query_vec: Vec<f32>,
pub(super) scorer: crate::graph::algorithms::vector::Scorer,
}
impl VectorScoreCache {
pub(super) fn key_for(args: &[Expression]) -> Option<Vec<ArgKey>> {
args[1..].iter().map(ArgKey::of).collect()
}
pub(super) fn matches(&self, args: &[Expression]) -> bool {
self.keys.as_ref().is_some_and(|keys| {
keys.len() + 1 == args.len()
&& keys
.iter()
.zip(&args[1..])
.all(|(key, arg)| key.matches(arg))
})
}
}
const VECTOR_SCORE_CACHE_SLOTS: usize = 4;
#[derive(Default)]
pub(super) struct VectorScoreCaches {
slots: [OnceLock<VectorScoreCache>; VECTOR_SCORE_CACHE_SLOTS],
}
impl VectorScoreCaches {
pub(super) fn get(&self, args: &[Expression]) -> Option<&VectorScoreCache> {
self.slots
.iter()
.filter_map(OnceLock::get)
.find(|entry| entry.matches(args))
}
pub(super) fn park(
&self,
entry: VectorScoreCache,
) -> Result<&VectorScoreCache, VectorScoreCache> {
if entry.keys.is_none() {
return Err(entry);
}
let mut entry = entry;
for slot in &self.slots {
match slot.set(entry) {
Ok(()) => {
return Ok(slot.get().expect("this thread just filled the slot"));
}
Err(returned) => entry = returned,
}
}
Err(entry)
}
}
#[cfg(test)]
thread_local! {
pub(super) static VECTOR_SCORE_PREPARES: std::cell::Cell<usize> =
const { std::cell::Cell::new(0) };
}
#[derive(PartialEq, Debug)]
pub(super) enum ArgKey {
Literal(Value),
Param(String),
LiteralList(Vec<Value>),
}
impl ArgKey {
pub(super) fn of(expr: &Expression) -> Option<Self> {
match expr {
Expression::Literal(value) => Some(ArgKey::Literal(value.clone())),
Expression::Parameter(name) => Some(ArgKey::Param(name.clone())),
Expression::ListLiteral(items) => items
.iter()
.map(|item| match item {
Expression::Literal(value) => Some(value.clone()),
_ => None,
})
.collect::<Option<Vec<Value>>>()
.map(ArgKey::LiteralList),
_ => None,
}
}
pub(super) fn matches(&self, expr: &Expression) -> bool {
match (self, expr) {
(ArgKey::Literal(value), Expression::Literal(other)) => value == other,
(ArgKey::Param(name), Expression::Parameter(other)) => name == other,
(ArgKey::LiteralList(values), Expression::ListLiteral(items)) => values.len()
== items.len()
&& values.iter().zip(items).all(
|(value, item)| matches!(item, Expression::Literal(other) if value == other),
),
_ => false,
}
}
}
pub(super) struct TextBm25Cache {
pub(super) node_type: String,
pub(super) keys: Option<(ArgKey, ArgKey)>,
pub(super) query_text: Option<String>,
pub(super) prepared: crate::graph::algorithms::text_index::bm25::PreparedQuery,
pub(super) prop_name: String,
pub(super) generation: u64,
}
pub fn clause_display_name(clause: &Clause) -> String {
match clause {
Clause::Match(m) => {
let types: Vec<&str> = m
.patterns
.iter()
.flat_map(|p| p.elements.iter())
.filter_map(|e| {
if let PatternElement::Node(n) = e {
n.node_type.as_deref()
} else {
None
}
})
.collect();
if types.is_empty() {
"Match".into()
} else {
format!("Match :{}", types.join(", :"))
}
}
Clause::OptionalMatch(m) => {
let types: Vec<&str> = m
.patterns
.iter()
.flat_map(|p| p.elements.iter())
.filter_map(|e| {
if let PatternElement::Node(n) = e {
n.node_type.as_deref()
} else {
None
}
})
.collect();
if types.is_empty() {
"OptionalMatch".into()
} else {
format!("OptionalMatch :{}", types.join(", :"))
}
}
Clause::Where(_) => "Where".into(),
Clause::Return(_) => "Return".into(),
Clause::With(_) => "With".into(),
Clause::OrderBy(_) => "OrderBy".into(),
Clause::Skip(_) => "Skip".into(),
Clause::Limit(_) => "Limit".into(),
Clause::Unwind(_) => "Unwind".into(),
Clause::LoadCsv(l) => {
if l.with_headers {
"LoadCsv (with headers)".into()
} else {
"LoadCsv".into()
}
}
Clause::Union(_) => "Union".into(),
Clause::Create(_) => "Create".into(),
Clause::Set(_) => "Set".into(),
Clause::Delete(_) => "Delete".into(),
Clause::Remove(_) => "Remove".into(),
Clause::Merge(_) => "Merge".into(),
Clause::Foreach { .. } => "Foreach".into(),
Clause::Call(_) => "Call".into(),
Clause::Schema(command) => match command {
SchemaCommand::CreateIndex(_) => "CreateIndex".into(),
SchemaCommand::UnsupportedIndexType { index_type, .. } => {
format!("CreateIndex ({})", index_type.keyword())
}
SchemaCommand::DropIndex(_) => "DropIndex".into(),
SchemaCommand::ShowIndexes => "ShowIndexes".into(),
SchemaCommand::ShowProcedures { .. } => "ShowProcedures".into(),
SchemaCommand::ShowFunctions { .. } => "ShowFunctions".into(),
SchemaCommand::ShowOntology => "ShowOntology".into(),
SchemaCommand::Constraint(ConstraintCommand::Create(_)) => "CreateConstraint".into(),
SchemaCommand::Constraint(ConstraintCommand::Drop { .. }) => "DropConstraint".into(),
SchemaCommand::Constraint(ConstraintCommand::Show) => "ShowConstraints".into(),
},
Clause::CallSubquery { .. } => "CallSubquery".into(),
Clause::FusedOptionalMatchAggregate { .. } => "FusedOptionalMatchAggregate".into(),
Clause::FusedVectorScoreTopK { .. } => "FusedVectorScoreTopK".into(),
Clause::FusedTextBm25TopK { .. } => "FusedTextBm25TopK".into(),
Clause::FusedMatchReturnAggregate { .. } => "FusedMatchReturnAggregate".into(),
Clause::FusedMatchWithAggregate { .. } => "FusedMatchWithAggregate".into(),
Clause::FusedOrderByTopK { .. } => "FusedOrderByTopK".into(),
Clause::FusedCountAll { .. } => "FusedCountAll".into(),
Clause::FusedCountAllEdges { .. } => "FusedCountAllEdges".into(),
Clause::FusedCountByType { .. } => "FusedCountByType".into(),
Clause::FusedCountEdgesByType { .. } => "FusedCountEdgesByType".into(),
Clause::FusedCountTypedNode { node_type, .. } => {
format!("FusedCountTypedNode :{node_type}")
}
Clause::FusedCountTypedEdge { edge_type, .. } => {
format!("FusedCountTypedEdge :{edge_type}")
}
Clause::FusedCountAnchoredEdges {
anchor_idx,
anchor_direction,
edge_types,
..
} => {
let arrow = match anchor_direction {
petgraph::Direction::Outgoing => "→",
petgraph::Direction::Incoming => "←",
};
let t = edge_types
.as_ref()
.map_or_else(|| "*".to_string(), |types| types.join("|"));
format!("FusedCountAnchoredEdges (anchor#{anchor_idx} {arrow} :{t})")
}
Clause::FusedNodeScanAggregate {
where_predicate, ..
} => format!("FusedNodeScanAggregate{}", filter_suffix(where_predicate)),
Clause::FusedNodeScanTopK {
limit,
where_predicate,
..
} => format!(
"FusedNodeScanTopK (k={limit}){}",
filter_suffix(where_predicate)
),
Clause::SpatialJoin {
container_type,
probe_type,
..
} => format!("SpatialJoin :{container_type} ⊇ :{probe_type}"),
}
}
fn filter_suffix(where_predicate: &Option<super::super::ast::Predicate>) -> &'static str {
if where_predicate.is_some() {
" +filter"
} else {
""
}
}