use std::collections::BTreeSet;
use ahash::AHashMap;
use dr_strange_core::compute::expr::{ArithOp, CmpOp, LogicOp};
use dr_strange_core::types::Dir;
use dr_strange_core::{
Agg, Algo, Binding, Expr, GraphChannel, HybridSpec, HybridWeights, KeywordChannel, LogicalPlan,
LouvainOptions, NodeId, NodeRef, PageRankOptions, ProjExpr, ProjItem, Projection, PropValue,
SortKey, Source, Step, TupleSortKey, VectorChannel,
};
use crate::Embedder;
use crate::ast::*;
fn resolve(arg: &VecArg, embedder: Option<&dyn Embedder>) -> Result<Vec<f32>, String> {
match arg {
VecArg::Vector(v) => Ok(v.clone()),
VecArg::Text(t) => {
let e = embedder.ok_or_else(|| {
"SEARCH by text needs an embedding provider (set the API key), \
or pass a literal vector like NEAR [..]"
.to_string()
})?;
e.embed(t)
.map_err(|e| format!("embedding the query text failed: {e}"))
}
}
}
enum Hop<'a> {
Rel(&'a RelPat),
Beam(&'a BeamClause),
}
pub fn compile(
mut q: Query,
embedder: Option<&dyn Embedder>,
params: &crate::Params,
) -> Result<LogicalPlan, String> {
let first_node = &q.source.first;
let source = compile_source(&q.source, embedder, params)?;
let mut hops: Vec<(Hop, &NodePat)> = Vec::new();
hops.extend(q.source.rest.iter().map(|(r, n)| (Hop::Rel(r), n)));
hops.extend(q.beams.iter().map(|b| (Hop::Beam(b), &b.node)));
let mut nodes: Vec<&NodePat> = vec![first_node];
nodes.extend(hops.iter().map(|(_, n)| *n));
let last = nodes.len() - 1;
let var_slot = bind_variables(&nodes)?;
let where_clause = q.where_clause.take();
let mut slot_filters = push_down_where(where_clause, &var_slot, nodes.len())?;
let source = key_seek_rewrite(source, &mut slot_filters[0], first_node, params)?;
let mut steps: Vec<Step> = Vec::new();
for e in slot_filters[0].drain(..) {
let scope = Scope {
var_slot: &var_slot,
current: 0,
};
steps.push(Step::Filter(compile_expr(&e, embedder, params, &scope)?));
}
compile_hops(
&mut steps,
&hops,
&nodes,
&mut slot_filters,
&var_slot,
embedder,
params,
)?;
let project = compile_return(&q, &mut steps, &var_slot, last, embedder, params)?;
Ok(LogicalPlan {
source,
steps,
project,
})
}
fn bind_variables<'a>(nodes: &[&'a NodePat]) -> Result<AHashMap<&'a str, usize>, String> {
let mut var_slot: AHashMap<&'a str, usize> = AHashMap::new();
for (idx, n) in nodes.iter().enumerate() {
if let Some(v) = &n.var
&& var_slot.insert(v.as_str(), idx).is_some()
{
return Err(format!(
"variable `{v}` is bound twice; reusing a pattern variable isn't supported yet"
));
}
}
Ok(var_slot)
}
fn push_down_where(
where_clause: Option<PExpr>,
var_slot: &AHashMap<&str, usize>,
slots: usize,
) -> Result<Vec<Vec<PExpr>>, String> {
let mut slot_filters: Vec<Vec<PExpr>> = vec![Vec::new(); slots];
let Some(w) = where_clause else {
return Ok(slot_filters);
};
for conj in split_and(w) {
let vars = referenced_vars(&conj);
let slot = match vars.len() {
0 => slots - 1,
1 => {
let v = vars.iter().next().unwrap();
let slot = *var_slot
.get(v.as_str())
.ok_or_else(|| format!("WHERE refers to unknown variable `{v}`"))?;
match reads_row_channel(&conj) {
true => slots - 1,
false => slot,
}
}
_ => {
return Err(
"a WHERE condition may reference only one pattern variable; \
cross-variable predicates aren't supported yet"
.to_string(),
);
}
};
slot_filters[slot].push(conj);
}
Ok(slot_filters)
}
fn key_seek_rewrite(
source: Source,
slot0: &mut Vec<PExpr>,
first_node: &NodePat,
params: &crate::Params,
) -> Result<Source, String> {
if !matches!(source, Source::ScanAll | Source::ScanLabel(_)) {
return Ok(source);
}
let var = first_node.var.as_deref();
let Some(pos) = slot0
.iter()
.position(|e| key_seek_keys(e, var, params).is_some())
else {
return Ok(source);
};
let keys = key_seek_keys(&slot0[pos], var, params).expect("just matched")?;
slot0.remove(pos);
if let Source::ScanLabel(label) = &source {
slot0.insert(
0,
PExpr::HasLabel {
var: var.unwrap_or_default().to_string(),
label: label.clone(),
},
);
}
Ok(Source::SeekKeys(keys))
}
#[allow(clippy::too_many_arguments)]
fn compile_hops(
steps: &mut Vec<Step>,
hops: &[(Hop<'_>, &NodePat)],
nodes: &[&NodePat],
slot_filters: &mut [Vec<PExpr>],
var_slot: &AHashMap<&str, usize>,
embedder: Option<&dyn Embedder>,
params: &crate::Params,
) -> Result<(), String> {
for idx in 1..nodes.len() {
match &hops[idx - 1].0 {
Hop::Rel(rel) => match rel.var_len {
None => steps.push(Step::Expand {
dir: rel.dir,
edge_type: rel.ty.clone(),
}),
Some(VarLen { min, max }) => {
let max = max.ok_or_else(|| {
"unbounded variable-length relationships aren't supported; \
give an upper bound, e.g. *1..3"
.to_string()
})?;
if min > max {
return Err(format!(
"variable-length range is empty: min ({min}) exceeds max ({max})"
));
}
steps.push(Step::ExpandVar {
dir: rel.dir,
edge_type: rel.ty.clone(),
min,
max,
});
}
},
Hop::Beam(b) => steps.push(Step::ExpandBeam {
dir: b.dir,
edge_type: b.edge_type.clone(),
property: vector_property(b.property.as_deref()),
query: resolve(&b.query, embedder)?,
metric: b.metric,
width: b.width,
depth: b.depth,
}),
}
if let Some(l) = &nodes[idx].label {
steps.push(Step::Filter(Expr::HasLabel(l.clone())));
}
let scope = Scope {
var_slot,
current: idx,
};
for e in slot_filters[idx].drain(..) {
steps.push(Step::Filter(compile_expr(&e, embedder, params, &scope)?));
}
}
Ok(())
}
fn compile_return(
q: &Query,
steps: &mut Vec<Step>,
var_slot: &AHashMap<&str, usize>,
last: usize,
embedder: Option<&dyn Embedder>,
params: &crate::Params,
) -> Result<Option<Projection>, String> {
let scope = Scope {
var_slot,
current: last,
};
let node_rows = q
.ret
.items
.iter()
.any(|i| matches!(i, ReturnItem::Star | ReturnItem::Var(_)));
if node_rows && q.ret.items.len() > 1 {
return Err(
"RETURN mixes a node with columns of values; a node isn't a value, \
so return its properties instead (`RETURN n.name, count(*)`)"
.to_string(),
);
}
if !node_rows {
return projection(q, &scope, embedder, params).map(Some);
}
if let Some(ReturnItem::Var(v)) = q.ret.items.first() {
let slot = *var_slot
.get(v.as_str())
.ok_or_else(|| format!("RETURN refers to unknown variable `{v}`"))?;
if slot != last {
return Err(format!(
"RETURN must name the pattern's last variable; returning an earlier \
variable (`{v}`) isn't supported yet — project its values instead \
(`RETURN {v}.name`)"
));
}
}
if q.ret.distinct {
steps.push(Step::Distinct);
}
if !q.order_by.is_empty() {
steps.push(Step::Sort(compile_sort(
&q.order_by,
var_slot,
last,
embedder,
params,
)?));
}
if let Some(s) = q.skip {
steps.push(Step::Skip(s));
}
if let Some(l) = q.limit {
steps.push(Step::Limit(l));
}
Ok(None)
}
fn projection(
q: &Query,
scope: &Scope<'_, '_>,
embedder: Option<&dyn Embedder>,
params: &crate::Params,
) -> Result<Projection, String> {
let mut items = Vec::with_capacity(q.ret.items.len());
for item in &q.ret.items {
items.push(match item {
ReturnItem::Value { expr, name } => {
ProjItem::value(name.clone(), compile_expr(expr, embedder, params, scope)?)
}
ReturnItem::Agg {
func,
arg,
distinct,
name,
} => {
let arg = arg
.as_ref()
.map(|e| compile_expr(e, embedder, params, scope))
.transpose()?;
ProjItem::agg(
name.clone(),
Agg {
func: *func,
arg,
distinct: *distinct,
},
)
}
ReturnItem::Star | ReturnItem::Var(_) => unreachable!("node RETURN"),
});
}
let order_by = order_columns(&q.order_by, &items, embedder, params, scope)?;
Ok(Projection {
items,
distinct: q.ret.distinct,
order_by,
skip: q.skip,
limit: q.limit,
})
}
fn compile_sort(
keys: &[OrderKey],
var_slot: &AHashMap<&str, usize>,
last: usize,
embedder: Option<&dyn Embedder>,
params: &crate::Params,
) -> Result<Vec<SortKey>, String> {
let mut out = Vec::with_capacity(keys.len());
for key in keys {
let e = match &key.target {
SortTarget::Expr(e) => e,
SortTarget::Name(_) => {
return Err(format!(
"ORDER BY `{}`: a bare name orders by a RETURN alias, and this query \
returns nodes rather than columns",
key.text
));
}
SortTarget::Agg { .. } => {
return Err(format!(
"ORDER BY `{}`: an aggregate orders a projection, and this query \
returns nodes rather than columns (`RETURN n.name, {}`)",
key.text, key.text
));
}
};
for v in referenced_vars(e) {
let slot = *var_slot
.get(v.as_str())
.ok_or_else(|| format!("ORDER BY refers to unknown variable `{v}`"))?;
if slot != last {
return Err("ORDER BY may reference only the returned (last) variable".to_string());
}
}
out.push(SortKey {
expr: compile_expr(
e,
embedder,
params,
&Scope {
var_slot,
current: last,
},
)?,
descending: key.descending,
});
}
Ok(out)
}
fn order_columns(
keys: &[OrderKey],
items: &[ProjItem],
embedder: Option<&dyn Embedder>,
params: &crate::Params,
scope: &Scope,
) -> Result<Vec<TupleSortKey>, String> {
let mut out = Vec::with_capacity(keys.len());
for key in keys {
let by_name = items.iter().position(|c| c.name == key.text);
let column = match (by_name, &key.target) {
(Some(i), _) => i,
(None, SortTarget::Expr(e)) => {
let compiled = ProjExpr::Value(compile_expr(e, embedder, params, scope)?);
items
.iter()
.position(|c| c.expr == compiled)
.ok_or_else(|| unknown_column(&key.text, items))?
}
(
None,
SortTarget::Agg {
func,
arg,
distinct,
},
) => {
let arg = arg
.as_ref()
.map(|e| compile_expr(e, embedder, params, scope))
.transpose()?;
let compiled = ProjExpr::Agg(Agg {
func: *func,
arg,
distinct: *distinct,
});
items
.iter()
.position(|c| c.expr == compiled)
.ok_or_else(|| unknown_column(&key.text, items))?
}
(None, SortTarget::Name(_)) => return Err(unknown_column(&key.text, items)),
};
out.push(TupleSortKey {
column,
descending: key.descending,
});
}
Ok(out)
}
fn unknown_column(text: &str, items: &[ProjItem]) -> String {
let names: Vec<&str> = items.iter().map(|c| c.name.as_str()).collect();
format!(
"ORDER BY `{text}` must name a returned column: {}",
names.join(", ")
)
}
const DEFAULT_VECTOR_PROPERTY: &str = "embedding";
fn vector_property(property: Option<&str>) -> String {
property.unwrap_or(DEFAULT_VECTOR_PROPERTY).to_string()
}
fn compile_source(
src: &QuerySource,
embedder: Option<&dyn Embedder>,
params: &crate::Params,
) -> Result<Source, String> {
let label = src.first.label.clone();
Ok(match &src.kind {
SourceKind::Match => match label {
Some(l) => Source::ScanLabel(l),
None => Source::ScanAll,
},
SourceKind::Search {
property,
query,
metric,
k,
} => Source::VectorTopK {
label,
property: vector_property(property.as_deref()),
query: resolve(query, embedder)?,
metric: *metric,
k: *k,
},
SourceKind::Keyword { property, query, k } => Source::KeywordTopK {
label: label.ok_or_else(|| {
"a keyword SEARCH needs a label — the BM25 index is declared on \
(label, property), e.g. SEARCH (d:Doc) ON body MATCHING \"…\""
.to_string()
})?,
property: property.clone().ok_or_else(|| {
"a keyword SEARCH needs ON <property> — keyword properties follow \
no convention, so there is nothing to default to. \
(Only NEAR defaults, to `embedding`.)"
.to_string()
})?,
query: query.clone(),
k: *k,
},
SourceKind::Hybrid(h) => Source::Hybrid(Box::new(compile_hybrid(h, label, embedder)?)),
SourceKind::Call(c) => Source::Algo {
label,
algo: compile_algo(c, params)?,
},
})
}
fn compile_hybrid(
h: &HybridClause,
label: Option<String>,
embedder: Option<&dyn Embedder>,
) -> Result<HybridSpec, String> {
if h.vector.is_none() && h.keyword.is_none() {
return Err(
"HYBRID needs at least one of a VECTOR or KEYWORD channel; GRAPH only \
boosts what those find"
.to_string(),
);
}
if h.keyword.is_some() && label.is_none() {
return Err(
"the HYBRID KEYWORD channel needs a label — the BM25 index is declared \
on (label, property), e.g. HYBRID (d:Doc) KEYWORD ON body MATCHING \"…\""
.to_string(),
);
}
let mut weights = HybridWeights::default();
let vector = match &h.vector {
Some(v) => {
if let Some(w) = v.weight {
weights.vector = w;
}
Some(VectorChannel {
property: vector_property(v.property.as_deref()),
query: resolve(&v.query, embedder)?,
metric: v.metric,
})
}
None => None,
};
let keyword = match &h.keyword {
Some(k) => {
if let Some(w) = k.weight {
weights.keyword = w;
}
Some(KeywordChannel {
property: k.property.clone().ok_or_else(|| {
"the HYBRID KEYWORD channel needs ON <property> — keyword \
properties follow no convention, so there is nothing to \
default to. (Only NEAR defaults, to `embedding`.)"
.to_string()
})?,
query: k.query.clone(),
})
}
None => None,
};
let graph = h.graph.as_ref().map(|g| {
if let Some(w) = g.weight {
weights.graph = w;
}
GraphChannel {
hops: g.hops,
decay: g.decay.unwrap_or(DEFAULT_GRAPH_DECAY),
seeds: g.seeds.unwrap_or(DEFAULT_GRAPH_SEEDS) as usize,
}
});
Ok(HybridSpec {
label,
vector,
keyword,
graph,
weights,
candidates: h.candidates.unwrap_or(DEFAULT_CANDIDATES) as usize,
k: h.k.unwrap_or(DEFAULT_HYBRID_K) as usize,
})
}
const DEFAULT_GRAPH_SEEDS: u64 = 10;
const DEFAULT_GRAPH_DECAY: f32 = 0.5;
const DEFAULT_CANDIDATES: u64 = 100;
const DEFAULT_HYBRID_K: u64 = 10;
fn compile_algo(c: &CallClause, params: &crate::Params) -> Result<Algo, String> {
let name = c.name.to_ascii_lowercase();
let mut args = AlgoArgs::new(c, params)?;
let algo = match name.as_str() {
"pagerank" => {
let d = PageRankOptions::default();
Algo::PageRank {
damping: args.float("damping")?.unwrap_or(d.damping),
max_iters: args
.u32("iterations")?
.or(args.u32("max_iters")?)
.unwrap_or(d.max_iters),
tolerance: args.float("tolerance")?.unwrap_or(d.tolerance),
}
}
"components" | "connected_components" => Algo::ConnectedComponents,
"louvain" => {
let d = LouvainOptions::default();
Algo::Louvain {
max_levels: args.u32("max_levels")?.unwrap_or(d.max_levels),
min_gain: args.float("min_gain")?.unwrap_or(d.min_gain),
}
}
"shortest_path" => Algo::ShortestPath {
from: args.node_ref("from")?,
to: args.node_ref("to")?,
dir: match args
.string("dir")?
.map(|s| s.to_ascii_lowercase())
.as_deref()
{
None | Some("out") => Dir::Out,
Some("in") => Dir::In,
Some("both") => Dir::Both,
Some(other) => {
return Err(format!(
"shortest_path dir must be out, in or both, not `{other}`"
));
}
},
weight: args.string("weight")?,
},
other => {
return Err(format!(
"unknown algorithm `{other}`; expected pagerank, components, \
shortest_path or louvain"
));
}
};
args.finish(&name)?;
Ok(algo)
}
struct AlgoArgs {
values: Vec<(String, PropValue)>,
}
impl AlgoArgs {
fn new(c: &CallClause, params: &crate::Params) -> Result<Self, String> {
let mut values = Vec::with_capacity(c.args.len());
for (name, val) in &c.args {
let v = match val {
Val::Lit(v) => v.clone(),
Val::Param(p) => crate::resolve_param(params, p)?,
};
values.push((name.to_ascii_lowercase(), v));
}
Ok(Self { values })
}
fn take(&mut self, name: &str) -> Option<PropValue> {
let pos = self.values.iter().position(|(n, _)| n == name)?;
Some(self.values.remove(pos).1)
}
fn float(&mut self, name: &str) -> Result<Option<f64>, String> {
match self.take(name) {
None => Ok(None),
Some(PropValue::Float(f)) => Ok(Some(f)),
Some(PropValue::Int(n)) => Ok(Some(n as f64)),
Some(other) => Err(format!("`{name}` must be a number, got {other:?}")),
}
}
fn int(&mut self, name: &str) -> Result<Option<i64>, String> {
match self.take(name) {
None => Ok(None),
Some(PropValue::Int(n)) if n >= 0 => Ok(Some(n)),
Some(other) => Err(format!(
"`{name}` must be a non-negative whole number, got {other:?}"
)),
}
}
fn u32(&mut self, name: &str) -> Result<Option<u32>, String> {
match self.int(name)? {
None => Ok(None),
Some(n) => u32::try_from(n)
.map(Some)
.map_err(|_| format!("`{name}` must be at most {}, got {n}", u32::MAX)),
}
}
fn string(&mut self, name: &str) -> Result<Option<String>, String> {
match self.take(name) {
None => Ok(None),
Some(PropValue::Str(s)) => Ok(Some(s)),
Some(other) => Err(format!("`{name}` must be a string, got {other:?}")),
}
}
fn node_ref(&mut self, name: &str) -> Result<NodeRef, String> {
match self.take(name) {
Some(PropValue::Str(s)) => Ok(NodeRef::Key(s)),
Some(PropValue::Int(n)) if n >= 0 => Ok(NodeRef::Id(NodeId(n as u64))),
Some(other) => Err(format!(
"`{name}` must be an external key (a string) or a node id, got {other:?}"
)),
None => Err(format!("shortest_path needs a `{name}` argument")),
}
}
fn finish(self, algo: &str) -> Result<(), String> {
match self.values.first() {
None => Ok(()),
Some((name, _)) => Err(format!("`{algo}` has no argument named `{name}`")),
}
}
}
fn key_seek_keys(
e: &PExpr,
var: Option<&str>,
params: &crate::Params,
) -> Option<Result<Vec<String>, String>> {
fn as_key(e: &PExpr, params: &crate::Params) -> Option<Result<String, String>> {
match e {
PExpr::Lit(PropValue::Str(s)) => Some(Ok(s.clone())),
PExpr::Param(name) => Some(match crate::resolve_param(params, name) {
Ok(PropValue::Str(s)) => Ok(s),
Ok(other) => Err(format!("`key()` compares against a string, got {other:?}")),
Err(e) => Err(e),
}),
_ => None,
}
}
let is_key = |e: &PExpr| matches!(e, PExpr::ExternalKey { var: v } if Some(v.as_str()) == var);
let values: Vec<&PExpr> = match e {
PExpr::Compare {
op: CmpOp::Eq,
lhs,
rhs,
} if is_key(lhs) => vec![rhs],
PExpr::Compare {
op: CmpOp::Eq,
lhs,
rhs,
} if is_key(rhs) => vec![lhs],
PExpr::In { lhs, list } if is_key(lhs) && !list.is_empty() => list.iter().collect(),
_ => return None,
};
let mut keys = Vec::with_capacity(values.len());
for v in values {
match as_key(v, params)? {
Ok(k) => keys.push(k),
Err(e) => return Some(Err(e)),
}
}
Some(Ok(keys))
}
fn split_and(e: PExpr) -> Vec<PExpr> {
use dr_strange_core::compute::expr::LogicOp;
match e {
PExpr::Logic {
op: LogicOp::And,
lhs,
rhs,
} => {
let mut v = split_and(*lhs);
v.extend(split_and(*rhs));
v
}
other => vec![other],
}
}
fn reads_row_channel(e: &PExpr) -> bool {
match e {
PExpr::Score | PExpr::Hops => true,
PExpr::Lit(_)
| PExpr::Param(_)
| PExpr::Prop { .. }
| PExpr::HasLabel { .. }
| PExpr::ExternalKey { .. }
| PExpr::Similarity { .. }
| PExpr::Distance { .. } => false,
PExpr::In { lhs, list } => reads_row_channel(lhs) || list.iter().any(reads_row_channel),
PExpr::InValue { lhs, haystack } => reads_row_channel(lhs) || reads_row_channel(haystack),
PExpr::IsNull(x) | PExpr::Not(x) | PExpr::Neg(x) => reads_row_channel(x),
PExpr::StringMatch { lhs, rhs, .. }
| PExpr::Compare { lhs, rhs, .. }
| PExpr::Logic { lhs, rhs, .. }
| PExpr::Arith { lhs, rhs, .. } => reads_row_channel(lhs) || reads_row_channel(rhs),
}
}
fn referenced_vars(e: &PExpr) -> BTreeSet<String> {
fn go(e: &PExpr, out: &mut BTreeSet<String>) {
match e {
PExpr::Prop { var, .. }
| PExpr::HasLabel { var, .. }
| PExpr::ExternalKey { var }
| PExpr::Similarity { var, .. }
| PExpr::Distance { var, .. } => {
out.insert(var.clone());
}
PExpr::Lit(_) | PExpr::Param(_) | PExpr::Score | PExpr::Hops => {}
PExpr::In { lhs, list } => {
go(lhs, out);
for e in list {
go(e, out);
}
}
PExpr::InValue { lhs, haystack } => {
go(lhs, out);
go(haystack, out);
}
PExpr::StringMatch { lhs, rhs, .. } => {
go(lhs, out);
go(rhs, out);
}
PExpr::IsNull(x) | PExpr::Not(x) | PExpr::Neg(x) => go(x, out),
PExpr::Compare { lhs, rhs, .. }
| PExpr::Logic { lhs, rhs, .. }
| PExpr::Arith { lhs, rhs, .. } => {
go(lhs, out);
go(rhs, out);
}
}
}
let mut out = BTreeSet::new();
go(e, &mut out);
out
}
struct Scope<'a, 'b> {
var_slot: &'a AHashMap<&'b str, usize>,
current: usize,
}
impl Scope<'_, '_> {
fn read(&self, var: &str, inner: Expr) -> Result<Expr, String> {
if var.is_empty() {
return Ok(inner);
}
let slot = *self
.var_slot
.get(var)
.ok_or_else(|| format!("unknown variable `{var}`"))?;
Ok(match slot == self.current {
true => inner,
false => Expr::At {
binding: Binding::Node(slot as u32),
inner: Box::new(inner),
},
})
}
}
fn compile_expr(
e: &PExpr,
embedder: Option<&dyn Embedder>,
params: &crate::Params,
scope: &Scope,
) -> Result<Expr, String> {
let sub = |x: &PExpr| compile_expr(x, embedder, params, scope).map(Box::new);
Ok(match e {
PExpr::Lit(v) => Expr::Literal(v.clone()),
PExpr::Param(name) => Expr::Literal(crate::resolve_param(params, name)?),
PExpr::Prop { var, key } => scope.read(var, Expr::Property(key.clone()))?,
PExpr::HasLabel { var, label } => scope.read(var, Expr::HasLabel(label.clone()))?,
PExpr::ExternalKey { var } => scope.read(var, Expr::ExternalKey)?,
PExpr::In { lhs, list } => {
let lhs = compile_expr(lhs, embedder, params, scope)?;
let mut eqs = Vec::with_capacity(list.len());
for item in list {
eqs.push(Expr::Compare {
op: CmpOp::Eq,
lhs: Box::new(lhs.clone()),
rhs: sub(item)?,
});
}
balanced(LogicOp::Or, eqs).unwrap_or(Expr::Literal(PropValue::Bool(false)))
}
PExpr::InValue { lhs, haystack } => Expr::In {
needle: sub(lhs)?,
haystack: sub(haystack)?,
},
PExpr::StringMatch { op, lhs, rhs } => Expr::StringMatch {
op: *op,
lhs: sub(lhs)?,
rhs: sub(rhs)?,
},
PExpr::IsNull(x) => Expr::IsNull(sub(x)?),
PExpr::Not(x) => null_guarded(Expr::Not(sub(x)?), &[x], embedder, params, scope)?,
PExpr::Neg(x) => match compile_expr(x, embedder, params, scope)? {
Expr::Literal(PropValue::Int(n)) => Expr::Literal(PropValue::Int(-n)),
Expr::Literal(PropValue::Float(f)) => Expr::Literal(PropValue::Float(-f)),
other => Expr::Arith {
op: ArithOp::Sub,
lhs: Box::new(Expr::Literal(PropValue::Int(0))),
rhs: Box::new(other),
},
},
PExpr::Compare { op, lhs, rhs } => {
let cmp = Expr::Compare {
op: *op,
lhs: sub(lhs)?,
rhs: sub(rhs)?,
};
match op {
CmpOp::Ne => null_guarded(cmp, &[lhs, rhs], embedder, params, scope)?,
CmpOp::Eq
if !is_non_null_constant(lhs, params) && !is_non_null_constant(rhs, params) =>
{
null_guarded(cmp, &[lhs, rhs], embedder, params, scope)?
}
_ => cmp,
}
}
PExpr::Logic { op, lhs, rhs } => Expr::Logic {
op: *op,
lhs: sub(lhs)?,
rhs: sub(rhs)?,
},
PExpr::Arith { op, lhs, rhs } => Expr::Arith {
op: *op,
lhs: sub(lhs)?,
rhs: sub(rhs)?,
},
PExpr::Score => Expr::Score,
PExpr::Hops => Expr::Hops,
PExpr::Similarity {
var,
property,
query,
metric,
} => scope.read(
var,
Expr::Similarity {
property: property.clone(),
query: resolve(query, embedder)?,
metric: *metric,
},
)?,
PExpr::Distance {
var,
property,
query,
metric,
} => scope.read(
var,
Expr::Distance {
property: property.clone(),
query: resolve(query, embedder)?,
metric: *metric,
},
)?,
})
}
fn is_non_null_constant(e: &PExpr, params: &crate::Params) -> bool {
match e {
PExpr::Lit(v) => !matches!(v, PropValue::Null),
PExpr::Param(name) => !matches!(
crate::resolve_param(params, name),
Ok(PropValue::Null) | Err(_)
),
_ => false,
}
}
fn nullable_leaves<'e>(roots: &[&'e PExpr]) -> Vec<&'e PExpr> {
fn go<'e>(e: &'e PExpr, out: &mut Vec<&'e PExpr>) {
match e {
PExpr::Prop { .. }
| PExpr::ExternalKey { .. }
| PExpr::Score
| PExpr::Similarity { .. }
| PExpr::Distance { .. }
| PExpr::Lit(PropValue::Null)
| PExpr::Param(_) => {
if !out.contains(&e) {
out.push(e);
}
}
PExpr::Lit(_) | PExpr::HasLabel { .. } | PExpr::Hops | PExpr::IsNull(_) => {}
PExpr::In { lhs, list } => {
go(lhs, out);
for e in list {
go(e, out);
}
}
PExpr::InValue { lhs, haystack } => {
go(lhs, out);
go(haystack, out);
}
PExpr::StringMatch { lhs, rhs, .. }
| PExpr::Compare { lhs, rhs, .. }
| PExpr::Logic { lhs, rhs, .. }
| PExpr::Arith { lhs, rhs, .. } => {
go(lhs, out);
go(rhs, out);
}
PExpr::Not(x) | PExpr::Neg(x) => go(x, out),
}
}
let mut out = Vec::new();
for r in roots {
go(r, &mut out);
}
out
}
fn null_guarded(
pred: Expr,
roots: &[&PExpr],
embedder: Option<&dyn Embedder>,
params: &crate::Params,
scope: &Scope,
) -> Result<Expr, String> {
let mut present = Vec::new();
for leaf in nullable_leaves(roots) {
present.push(compile_expr(leaf, embedder, params, scope)?.is_null().not());
}
Ok(match balanced(LogicOp::And, present) {
None => pred,
Some(guards) => guards.and(pred),
})
}
fn balanced(op: LogicOp, mut items: Vec<Expr>) -> Option<Expr> {
while items.len() > 1 {
let mut next = Vec::with_capacity(items.len().div_ceil(2));
let mut it = items.into_iter();
while let Some(lhs) = it.next() {
next.push(match it.next() {
Some(rhs) => Expr::Logic {
op,
lhs: Box::new(lhs),
rhs: Box::new(rhs),
},
None => lhs,
});
}
items = next;
}
items.pop()
}