use std::ops::Bound;
use std::sync::Arc;
use super::IndexCandidate;
use crate::catalog::{FieldDefinition, Index, IndexDefinition};
use crate::expr::operator::MatchesOperator;
use crate::expr::with::With;
use crate::expr::{BinaryOperator, Cond, Idiom, Kind, KindLiteral, Part};
use crate::kvs::Direction;
use crate::val::{Number, Range, Value};
#[derive(Debug, Clone)]
pub(crate) struct IndexRef {
pub(crate) indexes: Arc<[IndexDefinition]>,
pub(crate) idx: usize,
}
impl IndexRef {
pub fn new(indexes: Arc<[IndexDefinition]>, idx: usize) -> Self {
Self {
indexes,
idx,
}
}
pub fn definition(&self) -> &IndexDefinition {
&self.indexes[self.idx]
}
pub fn is_unique(&self) -> bool {
matches!(self.definition().index, crate::catalog::Index::Uniq)
}
}
impl std::ops::Deref for IndexRef {
type Target = IndexDefinition;
fn deref(&self) -> &Self::Target {
self.definition()
}
}
impl std::hash::Hash for IndexRef {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.idx.hash(state);
}
}
impl PartialEq for IndexRef {
fn eq(&self, other: &Self) -> bool {
self.idx == other.idx
}
}
impl Eq for IndexRef {}
#[derive(Debug, Clone)]
pub enum AccessPath {
TableScan,
EmptyScan,
BTreeScan {
index_ref: IndexRef,
access: BTreeAccess,
direction: Direction,
},
FullTextSearch {
index_ref: IndexRef,
query: String,
operator: MatchesOperator,
},
KnnSearch {
index_ref: IndexRef,
vector: Vec<Number>,
k: u32,
ef: u32,
prefilter: Option<KnnPrefilterPlan>,
},
BitmapFusion {
root: BitmapPlan,
fallback: Option<Box<AccessPath>>,
},
Union {
paths: Vec<AccessPath>,
dedupe: bool,
},
}
impl AccessPath {
pub fn is_full_range_scan(&self) -> bool {
matches!(
self,
AccessPath::BTreeScan {
access: BTreeAccess::Range {
range: Range {
start: Bound::Unbounded,
end: Bound::Unbounded,
},
},
..
}
)
}
}
#[derive(Debug, Clone)]
pub struct KnnPrefilterPlan {
pub root: BitmapPlan,
pub residual: Option<Cond>,
pub uncovered_matches: bool,
}
#[derive(Debug, Clone)]
pub enum BitmapPlan {
BTree {
index_ref: IndexRef,
access: BTreeAccess,
},
FullText {
index_ref: IndexRef,
query: String,
operator: MatchesOperator,
},
Graph {
source: crate::val::RecordId,
direction: crate::expr::Dir,
edge_tables: Vec<surrealdb_strand::TableName>,
},
And(Vec<BitmapPlan>),
Or(Vec<BitmapPlan>),
AndNot {
base: Box<BitmapPlan>,
subtract: Box<BitmapPlan>,
},
}
#[derive(Debug, Clone)]
pub enum BTreeAccess {
Equality(Value),
Range {
range: Range,
},
Compound {
prefix: Vec<Value>,
range: Option<(BinaryOperator, Value)>,
},
FullText {
query: String,
operator: crate::expr::operator::MatchesOperator,
},
Knn {
vector: Vec<Number>,
k: u32,
ef: u32,
},
}
impl BTreeAccess {
pub(crate) fn describe(&self) -> String {
use surrealdb_types::ToSql;
match self {
BTreeAccess::Equality(v) => format!("= {}", v.to_sql()),
BTreeAccess::Range {
range,
} => {
let from_str = match range.start.as_ref() {
Bound::Included(x) => format!(">={}", x.to_sql()),
Bound::Excluded(x) => format!(">{}", x.to_sql()),
Bound::Unbounded => String::new(),
};
let to_str = match range.end.as_ref() {
Bound::Included(x) => format!("<={}", x.to_sql()),
Bound::Excluded(x) => format!("<{}", x.to_sql()),
Bound::Unbounded => String::new(),
};
format!("{from_str} {to_str}").trim().to_string()
}
BTreeAccess::Compound {
prefix,
range,
} => {
let prefix_str = prefix.iter().map(|v| v.to_sql()).collect::<Vec<_>>().join(", ");
if let Some((op, val)) = range {
let val_sql = val.to_sql();
format!("[{prefix_str}] {op:?} {val_sql}")
} else {
format!("[{prefix_str}]")
}
}
BTreeAccess::FullText {
query,
..
} => format!("@@ {query}"),
BTreeAccess::Knn {
k,
..
} => format!("knn {k}"),
}
}
}
pub(crate) fn access_fans_out(cols: &[Idiom], access: &BTreeAccess) -> bool {
let pinned = match access {
BTreeAccess::Compound {
prefix,
..
} => prefix.len(),
BTreeAccess::Equality(_) => 1,
_ => 0,
};
cols.iter().skip(pinned).any(|c| c.0.iter().any(|p| matches!(p, Part::All)))
}
pub(crate) fn with_array_columns_as_elements(
indexes: Arc<[IndexDefinition]>,
fields: &[FieldDefinition],
) -> ElementColumns {
let as_elements = |col: &Idiom| -> bool {
!col.0.iter().any(|p| matches!(p, Part::All | Part::Flatten))
&& declared_kind(col, fields).is_some_and(kind_holds_only_arrays)
};
let mut changed: Option<(Vec<IndexDefinition>, Vec<ColumnFlags>)> = None;
for (i, ix) in indexes.iter().enumerate() {
let flags: Box<[bool]> = if matches!(ix.index, Index::Idx | Index::Uniq) {
ix.cols.iter().map(as_elements).collect()
} else {
Box::default()
};
if flags.contains(&true) {
let (defs, masks) = changed.get_or_insert_with(|| {
(indexes[..i].to_vec(), indexes[..i].iter().map(|_| Box::default()).collect())
});
let mut cols = ix.cols.clone();
for (col, _) in cols.iter_mut().zip(flags.iter()).filter(|(_, rewrite)| **rewrite) {
col.0.push(Part::All);
}
defs.push(IndexDefinition {
cols,
..ix.clone()
});
masks.push(flags);
} else if let Some((defs, masks)) = changed.as_mut() {
defs.push(ix.clone());
masks.push(Box::default());
}
}
match changed {
Some((defs, masks)) => ElementColumns {
indexes: defs.into(),
rewritten: Some(masks.into()),
},
None => ElementColumns::unchanged(indexes),
}
}
pub(crate) type ColumnFlags = Box<[bool]>;
pub(crate) struct ElementColumns {
pub(crate) indexes: Arc<[IndexDefinition]>,
pub(crate) rewritten: Option<Arc<[ColumnFlags]>>,
}
impl ElementColumns {
pub(crate) fn unchanged(indexes: Arc<[IndexDefinition]>) -> Self {
Self {
indexes,
rewritten: None,
}
}
pub(crate) fn analyzer<'a>(
&self,
with: Option<&'a With>,
) -> crate::exec::index::analysis::IndexAnalyzer<'a> {
let analyzer =
crate::exec::index::analysis::IndexAnalyzer::new(Arc::clone(&self.indexes), with);
match &self.rewritten {
Some(rewritten) => analyzer.without_containment_on(Arc::clone(rewritten)),
None => analyzer,
}
}
}
pub(crate) fn may_have_array_columns(indexes: &[IndexDefinition]) -> bool {
indexes.iter().any(|ix| {
matches!(ix.index, Index::Idx | Index::Uniq)
&& ix.cols.iter().any(|c| !c.0.iter().any(|p| matches!(p, Part::All | Part::Flatten)))
})
}
pub(crate) fn without_btree_indexes(indexes: Arc<[IndexDefinition]>) -> Arc<[IndexDefinition]> {
if !indexes.iter().any(|ix| matches!(ix.index, Index::Idx | Index::Uniq)) {
return indexes;
}
indexes.iter().filter(|ix| !matches!(ix.index, Index::Idx | Index::Uniq)).cloned().collect()
}
fn declared_kind<'a>(col: &Idiom, fields: &'a [FieldDefinition]) -> Option<&'a Kind> {
if let Some(kind) =
fields.iter().find(|fd| &fd.name == col).and_then(|fd| fd.field_kind.as_ref())
{
return Some(kind);
}
(1..col.0.len()).rev().find_map(|n| {
let fd = fields.iter().find(|fd| fd.name.0.as_slice() == &col.0[..n])?;
col.0[n..].iter().try_fold(fd.field_kind.as_ref()?, |kind, part| match part {
Part::Field(name) => object_literal_field(kind, name.as_str()),
_ => None,
})
})
}
fn object_literal_field<'a>(kind: &'a Kind, name: &str) -> Option<&'a Kind> {
match kind {
Kind::Literal(KindLiteral::Object(fields)) => fields.get(name),
Kind::Either(kinds) => {
let mut objects = kinds.iter().filter(|k| !matches!(k, Kind::None | Kind::Null));
match (objects.next(), objects.next()) {
(Some(object), None) => object_literal_field(object, name),
_ => None,
}
}
_ => None,
}
}
fn kind_holds_only_arrays(kind: &Kind) -> bool {
fn arrays(kind: &Kind) -> Option<bool> {
match kind {
Kind::Array(..) | Kind::Literal(KindLiteral::Array(_)) => Some(true),
Kind::None | Kind::Null => Some(false),
Kind::Either(kinds) => kinds.iter().try_fold(false, |seen, k| Some(arrays(k)? || seen)),
_ => None,
}
}
arrays(kind) == Some(true)
}
pub(crate) fn scan_ordering(
index_ref: &IndexRef,
access: &BTreeAccess,
direction: Direction,
) -> Vec<crate::exec::ordering::SortProperty> {
use crate::exec::field_path::FieldPath;
use crate::exec::field_path_convert::field_path_from_idiom;
use crate::exec::operators::SortDirection;
use crate::exec::ordering::SortProperty;
let direction = match direction {
Direction::Forward => SortDirection::Asc,
Direction::Backward => SortDirection::Desc,
};
let property = |path| SortProperty {
path,
direction,
collate: false,
numeric: false,
};
let ix_def = index_ref.definition();
let pinned = match access {
BTreeAccess::Compound {
prefix,
..
} => prefix.len(),
BTreeAccess::Equality(_) => ix_def.cols.len(),
_ => 0,
};
let free = ix_def.cols.iter().skip(pinned);
let mut ordering: Vec<SortProperty> =
free.clone().map_while(|col| field_path_from_idiom(col).ok().map(property)).collect();
if !index_ref.is_unique() && !ix_def.cols.is_empty() && ordering.len() == free.len() {
ordering.push(property(FieldPath::field("id")));
}
ordering
}
pub fn select_access_path(
mut candidates: Vec<IndexCandidate>,
with_hints: Option<&With>,
direction: Direction,
) -> AccessPath {
if matches!(with_hints, Some(With::NoIndex)) {
return AccessPath::TableScan;
}
if let Some(pos) = candidates.iter().position(|c| matches!(c.access, BTreeAccess::Knn { .. })) {
if let Some(With::Index(names)) = with_hints {
tracing::warn!(
target: "surreal::index",
hinted = ?names,
"WITH INDEX hint overridden: a KNN operator can only be computed by its KnnScan",
);
}
return candidates.swap_remove(pos).to_access_path(direction);
}
if let Some(With::Index(names)) = with_hints {
if let Some(candidate) = find_hinted_index(&candidates, names) {
return candidate.to_access_path(direction);
}
tracing::warn!(
target: "surreal::index",
hinted = ?names,
candidates = ?candidates.iter().map(|c| c.index_ref.name.as_str()).collect::<Vec<_>>(),
"WITH INDEX hint did not match any analyzed candidate; falling back to best-effort plan",
);
}
if candidates.is_empty() {
return AccessPath::TableScan;
}
candidates.sort_by_key(|c| (std::cmp::Reverse(c.score()), c.index_ref.idx));
candidates
.into_iter()
.next()
.map(|c| c.to_access_path(direction))
.unwrap_or(AccessPath::TableScan)
}
fn find_hinted_index<'a>(
candidates: &'a [IndexCandidate],
names: &[String],
) -> Option<&'a IndexCandidate> {
for name in names {
if let Some(candidate) = candidates.iter().find(|c| &c.index_ref.name == name) {
return Some(candidate);
}
}
None
}
#[cfg(test)]
mod tests {
use std::str::FromStr;
use surrealdb_strand::Strand;
use surrealdb_types::ToSql;
use super::*;
use crate::catalog::{FullTextParams, Index, IndexId, Scoring};
use crate::exec::index::IndexCandidate;
use crate::exec::planner::util::strip_index_conditions;
use crate::expr::operator::{BooleanOperator, MatchesOperator};
use crate::expr::{Cond, Expr, Idiom};
fn idx_def(id: u32, name: &str, cols: &[&str], kind: Index) -> IndexDefinition {
IndexDefinition {
index_id: IndexId(id),
name: Strand::from(name),
table_name: "t".into(),
cols: cols.iter().map(|c| Idiom::from_str(c).expect("valid idiom")).collect(),
index: kind,
count_cond: None,
comment: None,
prepare_remove: false,
format_version: 1,
}
}
fn idx_basic(id: u32, name: &str, cols: &[&str]) -> IndexDefinition {
idx_def(id, name, cols, Index::Idx)
}
fn idx_uniq(id: u32, name: &str, cols: &[&str]) -> IndexDefinition {
idx_def(id, name, cols, Index::Uniq)
}
fn idx_ft(id: u32, name: &str, cols: &[&str]) -> IndexDefinition {
idx_def(
id,
name,
cols,
Index::FullText(FullTextParams {
analyzer: "simple".into(),
highlight: false,
scoring: Scoring::Bm {
k1: 1.2,
b: 0.75,
},
}),
)
}
fn refs(defs: Vec<IndexDefinition>) -> Arc<[IndexDefinition]> {
Arc::<[_]>::from(defs.into_boxed_slice())
}
fn index_ref(defs: Vec<IndexDefinition>, idx: usize) -> IndexRef {
IndexRef::new(refs(defs), idx)
}
fn candidate(defs: Vec<IndexDefinition>, idx: usize, access: BTreeAccess) -> IndexCandidate {
IndexCandidate::new(index_ref(defs, idx), access)
}
fn matches_op() -> MatchesOperator {
MatchesOperator {
rf: None,
operator: BooleanOperator::And,
}
}
fn num(n: i64) -> Value {
Value::from(n)
}
fn range(start: Bound<Value>, end: Bound<Value>) -> BTreeAccess {
BTreeAccess::Range {
range: Range {
start,
end,
},
}
}
fn parse_cond(snippet: &str) -> Cond {
let src = format!("SELECT * FROM t WHERE {snippet}");
let ast = crate::syn::parse(&src).expect("parse");
let mut exprs = ast.expressions;
assert_eq!(exprs.len(), 1, "expected one statement from {src:?}");
let top: crate::expr::TopLevelExpr = exprs.remove(0).into();
match top {
crate::expr::TopLevelExpr::Expr(Expr::Select(s)) => s.cond.expect("WHERE"),
other => panic!("expected SELECT, got {other:?}"),
}
}
fn residual(cond: &str, access: &BTreeAccess, cols: &[&str]) -> Option<String> {
let cols: Vec<Idiom> =
cols.iter().map(|c| Idiom::from_str(c).expect("valid idiom")).collect();
strip_index_conditions(&parse_cond(cond), access, &cols).map(|c| c.0.to_sql())
}
mod index_ref {
use super::*;
#[test]
fn resolves_the_definition_at_its_own_position() {
let r =
index_ref(vec![idx_basic(1, "ix_a", &["a"]), idx_basic(2, "ix_b", &["b", "c"])], 1);
assert_eq!(r.definition().name.as_str(), "ix_b");
assert_eq!(r.cols.len(), 2);
}
#[test]
fn is_unique_holds_only_for_the_uniq_kind() {
assert!(!index_ref(vec![idx_basic(1, "ix", &["a"])], 0).is_unique());
assert!(index_ref(vec![idx_uniq(1, "ix", &["a"])], 0).is_unique());
assert!(
!index_ref(vec![idx_ft(1, "ix", &["a"])], 0).is_unique(),
"a full-text index is not a unique b-tree"
);
}
#[test]
fn identity_is_the_position_alone() {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
fn hash(r: &IndexRef) -> u64 {
let mut h = DefaultHasher::new();
r.hash(&mut h);
h.finish()
}
let list = vec![idx_basic(1, "ix_a", &["a"]), idx_basic(2, "ix_b", &["b"])];
let first = index_ref(list.clone(), 0);
let second = index_ref(list, 1);
assert_ne!(first, second);
assert_ne!(hash(&first), hash(&second));
let other_list = index_ref(vec![idx_uniq(9, "unrelated", &["z"])], 0);
assert_eq!(first, other_list);
assert_eq!(hash(&first), hash(&other_list));
}
}
mod full_range {
use super::*;
fn btree(access: BTreeAccess) -> AccessPath {
AccessPath::BTreeScan {
index_ref: index_ref(vec![idx_basic(1, "ix_a", &["a"])], 0),
access,
direction: Direction::Forward,
}
}
#[test]
fn doubly_unbounded_btree_range_is_a_full_range_scan() {
assert!(btree(range(Bound::Unbounded, Bound::Unbounded)).is_full_range_scan());
}
#[test]
fn any_bound_makes_the_scan_selective() {
for access in [
range(Bound::Included(num(1)), Bound::Unbounded),
range(Bound::Excluded(num(1)), Bound::Unbounded),
range(Bound::Unbounded, Bound::Included(num(9))),
range(Bound::Unbounded, Bound::Excluded(num(9))),
range(Bound::Included(num(1)), Bound::Included(num(9))),
] {
assert!(
!btree(access.clone()).is_full_range_scan(),
"{} carries WHERE selectivity",
access.describe()
);
}
}
#[test]
fn other_shapes_are_never_full_range_scans() {
let ft = index_ref(vec![idx_ft(1, "ix_ft", &["body"])], 0);
let paths = vec![
AccessPath::TableScan,
AccessPath::EmptyScan,
btree(BTreeAccess::Equality(num(1))),
btree(BTreeAccess::Compound {
prefix: vec![],
range: None,
}),
AccessPath::FullTextSearch {
index_ref: ft.clone(),
query: "hello".to_owned(),
operator: matches_op(),
},
AccessPath::KnnSearch {
index_ref: ft,
vector: vec![Number::Int(1)],
k: 3,
ef: 10,
prefilter: None,
},
AccessPath::Union {
paths: vec![btree(range(Bound::Unbounded, Bound::Unbounded))],
dedupe: true,
},
];
for path in paths {
assert!(!path.is_full_range_scan(), "{path:?} is not a full-range b-tree scan");
}
}
}
mod describe {
use super::*;
#[test]
fn equality_renders_the_sql_literal() {
assert_eq!(BTreeAccess::Equality(num(5)).describe(), "= 5");
assert_eq!(BTreeAccess::Equality(Value::from("x")).describe(), "= 'x'");
assert_eq!(BTreeAccess::Equality(Value::None).describe(), "= NONE");
}
#[test]
fn every_range_bound_combination_renders() {
let cases = [
(Bound::Included(num(1)), Bound::Included(num(9)), ">=1 <=9"),
(Bound::Included(num(1)), Bound::Excluded(num(9)), ">=1 <9"),
(Bound::Excluded(num(1)), Bound::Included(num(9)), ">1 <=9"),
(Bound::Excluded(num(1)), Bound::Excluded(num(9)), ">1 <9"),
(Bound::Included(num(1)), Bound::Unbounded, ">=1"),
(Bound::Excluded(num(1)), Bound::Unbounded, ">1"),
(Bound::Unbounded, Bound::Included(num(9)), "<=9"),
(Bound::Unbounded, Bound::Excluded(num(9)), "<9"),
(Bound::Unbounded, Bound::Unbounded, ""),
];
for (start, end, expected) in cases {
assert_eq!(range(start, end).describe(), expected);
}
}
#[test]
fn compound_renders_the_prefix_and_any_range() {
let prefix = vec![num(1), Value::from("b")];
assert_eq!(
BTreeAccess::Compound {
prefix: prefix.clone(),
range: None,
}
.describe(),
"[1, 'b']"
);
assert_eq!(
BTreeAccess::Compound {
prefix,
range: Some((BinaryOperator::MoreThan, num(3))),
}
.describe(),
"[1, 'b'] MoreThan 3"
);
}
#[test]
fn fulltext_and_knn_render_their_own_shorthand() {
assert_eq!(
BTreeAccess::FullText {
query: "hello world".to_owned(),
operator: matches_op(),
}
.describe(),
"@@ hello world"
);
assert_eq!(
BTreeAccess::Knn {
vector: vec![Number::Int(1), Number::Int(2)],
k: 4,
ef: 40,
}
.describe(),
"knn 4"
);
}
}
mod selection {
use super::*;
fn defs() -> Vec<IndexDefinition> {
vec![idx_basic(1, "ix_a", &["a"]), idx_uniq(2, "ix_b", &["b"])]
}
fn scan_index(path: &AccessPath) -> &str {
match path {
AccessPath::BTreeScan {
index_ref,
..
} => index_ref.name.as_str(),
other => panic!("expected BTreeScan, got {other:?}"),
}
}
#[test]
fn noindex_hint_forces_a_table_scan() {
let mut empty = candidate(defs(), 1, BTreeAccess::Equality(num(1)));
empty.empty = true;
let path = select_access_path(vec![empty], Some(&With::NoIndex), Direction::Forward);
assert!(matches!(path, AccessPath::TableScan));
}
#[test]
fn named_hint_wins_over_a_better_scoring_candidate() {
let candidates = vec![
candidate(defs(), 0, range(Bound::Included(num(1)), Bound::Unbounded)),
candidate(defs(), 1, BTreeAccess::Equality(num(1))),
];
let with = With::Index(vec!["ix_a".to_owned()]);
let path = select_access_path(candidates, Some(&with), Direction::Forward);
assert_eq!(scan_index(&path), "ix_a");
}
#[test]
fn hint_name_order_decides_between_two_hinted_candidates() {
let candidates = vec![
candidate(defs(), 0, BTreeAccess::Equality(num(1))),
candidate(defs(), 1, BTreeAccess::Equality(num(1))),
];
let with = With::Index(vec!["ix_b".to_owned(), "ix_a".to_owned()]);
let path = select_access_path(candidates, Some(&with), Direction::Forward);
assert_eq!(scan_index(&path), "ix_b");
}
#[test]
fn unmatched_hint_falls_back_to_best_effort_selection() {
let candidates = vec![candidate(defs(), 1, BTreeAccess::Equality(num(1)))];
let with = With::Index(vec!["nonexistent".to_owned()]);
let path = select_access_path(candidates, Some(&with), Direction::Forward);
assert_eq!(scan_index(&path), "ix_b", "an unmatched hint does not veto the plan");
}
#[test]
fn no_candidates_is_a_table_scan() {
assert!(matches!(
select_access_path(vec![], None, Direction::Forward),
AccessPath::TableScan
));
}
#[test]
fn an_empty_candidate_short_circuits_to_empty_scan() {
let mut empty = candidate(defs(), 0, range(Bound::Included(num(1)), Bound::Unbounded));
empty.empty = true;
let candidates = vec![candidate(defs(), 1, BTreeAccess::Equality(num(1))), empty];
let path = select_access_path(candidates, None, Direction::Forward);
assert!(matches!(path, AccessPath::EmptyScan));
}
#[test]
fn a_score_tie_resolves_to_the_first_index_in_catalog_order() {
let defs = vec![idx_basic(1, "ix_a1", &["a"]), idx_basic(2, "ix_a2", &["a"])];
let candidates = vec![
candidate(defs.clone(), 0, BTreeAccess::Equality(num(1))),
candidate(defs, 1, BTreeAccess::Equality(num(1))),
];
let path = select_access_path(candidates, None, Direction::Forward);
assert_eq!(scan_index(&path), "ix_a1");
}
#[test]
fn a_score_tie_resolves_the_same_way_whatever_order_candidates_arrive_in() {
let defs = vec![idx_basic(1, "ix_a1", &["a"]), idx_basic(2, "ix_a2", &["a"])];
let reversed = vec![
candidate(defs.clone(), 1, BTreeAccess::Equality(num(1))),
candidate(defs, 0, BTreeAccess::Equality(num(1))),
];
let path = select_access_path(reversed, None, Direction::Forward);
assert_eq!(scan_index(&path), "ix_a1");
}
#[test]
fn the_requested_direction_reaches_the_btree_scan() {
let candidates = vec![candidate(defs(), 1, BTreeAccess::Equality(num(1)))];
let path = select_access_path(candidates, None, Direction::Backward);
match path {
AccessPath::BTreeScan {
direction,
..
} => assert_eq!(direction, Direction::Backward),
other => panic!("expected BTreeScan, got {other:?}"),
}
}
#[test]
fn specialised_access_shapes_get_their_own_path_kind() {
let ft_defs = vec![idx_ft(1, "ix_ft", &["body"])];
let ft = candidate(
ft_defs.clone(),
0,
BTreeAccess::FullText {
query: "hello".to_owned(),
operator: matches_op(),
},
);
assert!(matches!(
select_access_path(vec![ft], None, Direction::Forward),
AccessPath::FullTextSearch { .. }
));
let knn = candidate(
ft_defs,
0,
BTreeAccess::Knn {
vector: vec![Number::Int(1)],
k: 3,
ef: 10,
},
);
assert!(matches!(
select_access_path(vec![knn], None, Direction::Forward),
AccessPath::KnnSearch {
k: 3,
ef: 10,
..
}
));
}
}
mod residual {
use super::*;
#[test]
fn equality_consumes_its_own_leaf_and_leaves_the_rest() {
let access = BTreeAccess::Equality(num(5));
assert_eq!(residual("a = 5", &access, &["a"]), None);
assert_eq!(residual("a = 5 AND b = 1", &access, &["a"]), Some("b = 1".to_owned()));
}
#[test]
fn equality_on_another_value_or_column_is_retained() {
let access = BTreeAccess::Equality(num(5));
assert_eq!(residual("a = 6", &access, &["a"]), Some("a = 6".to_owned()));
assert_eq!(residual("b = 5", &access, &["a"]), Some("b = 5".to_owned()));
}
#[test]
fn range_consumes_only_the_leaf_its_bound_came_from() {
let access = range(Bound::Excluded(num(5)), Bound::Unbounded);
assert_eq!(residual("a > 5", &access, &["a"]), None);
assert_eq!(residual("a >= 5", &access, &["a"]), Some("a >= 5".to_owned()));
assert_eq!(residual("a > 6", &access, &["a"]), Some("a > 6".to_owned()));
}
#[test]
fn a_bounded_range_consumes_both_of_its_leaves() {
let access = range(Bound::Included(num(1)), Bound::Excluded(num(9)));
assert_eq!(residual("a >= 1 AND a < 9", &access, &["a"]), None);
}
#[test]
fn flipped_operand_order_is_still_consumed() {
let access = range(Bound::Excluded(num(5)), Bound::Unbounded);
assert_eq!(residual("5 < a", &access, &["a"]), None);
}
#[test]
fn compound_prefix_consumes_positional_equalities() {
let access = BTreeAccess::Compound {
prefix: vec![num(1), num(2)],
range: None,
};
assert_eq!(residual("a = 1 AND b = 2", &access, &["a", "b", "c"]), None);
assert_eq!(
residual("a = 1 AND b = 2 AND c = 3", &access, &["a", "b", "c"]),
Some("c = 3".to_owned()),
"no prefix value pins c"
);
}
#[test]
fn a_prefix_value_at_the_wrong_column_is_retained() {
let access = BTreeAccess::Compound {
prefix: vec![num(1), num(2)],
range: None,
};
assert_eq!(residual("a = 2", &access, &["a", "b"]), Some("a = 2".to_owned()));
assert_eq!(residual("b = 1", &access, &["a", "b"]), Some("b = 1".to_owned()));
}
#[test]
fn compound_range_is_consumed_only_on_the_column_after_the_prefix() {
let access = BTreeAccess::Compound {
prefix: vec![num(1)],
range: Some((BinaryOperator::MoreThan, num(2))),
};
assert_eq!(residual("a = 1 AND b > 2", &access, &["a", "b", "c"]), None);
assert_eq!(
residual("a = 1 AND c > 2", &access, &["a", "b", "c"]),
Some("c > 2".to_owned())
);
assert_eq!(
residual("a = 1 AND b >= 2", &access, &["a", "b", "c"]),
Some("b >= 2".to_owned())
);
}
#[test]
fn not_none_is_consumed_through_its_exclusive_none_encoding() {
let as_range = range(Bound::Excluded(Value::None), Bound::Unbounded);
assert_eq!(residual("a != NONE", &as_range, &["a"]), None);
let as_compound = BTreeAccess::Compound {
prefix: vec![num(1)],
range: Some((BinaryOperator::MoreThan, Value::None)),
};
assert_eq!(residual("a = 1 AND b != NONE", &as_compound, &["a", "b"]), None);
}
#[test]
fn a_leaf_under_or_is_never_consumed() {
let access = BTreeAccess::Equality(num(5));
assert_eq!(
residual("a = 5 OR b = 1", &access, &["a"]),
Some("a = 5 OR b = 1".to_owned())
);
}
#[test]
fn single_element_in_is_consumed_only_with_the_idiom_on_the_left() {
let access = BTreeAccess::Equality(num(5));
assert_eq!(residual("a IN [5]", &access, &["a"]), None);
assert_eq!(residual("[5] INSIDE a", &access, &["a"]), Some("[5] INSIDE a".to_owned()));
}
#[test]
fn containment_is_consumed_only_on_an_array_element_column() {
let access = BTreeAccess::Equality(Value::from("x"));
assert_eq!(residual("tags CONTAINS 'x'", &access, &["tags.*"]), None);
assert_eq!(
residual("tags CONTAINS 'x'", &access, &["tags"]),
Some("tags CONTAINS 'x'".to_owned())
);
}
#[test]
fn a_non_literal_operand_is_retained() {
let access = BTreeAccess::Equality(num(5));
assert_eq!(residual("a = $p", &access, &["a"]), Some("a = $p".to_owned()));
}
#[test]
fn fulltext_and_knn_shapes_consume_nothing() {
let ft = BTreeAccess::FullText {
query: "hello".to_owned(),
operator: matches_op(),
};
assert_eq!(residual("a = 5", &ft, &["a"]), Some("a = 5".to_owned()));
let knn = BTreeAccess::Knn {
vector: vec![Number::Int(1)],
k: 3,
ef: 10,
};
assert_eq!(residual("a = 5", &knn, &["a"]), Some("a = 5".to_owned()));
}
}
mod array_columns {
use super::*;
fn field(name: &str, kind: Option<Kind>) -> FieldDefinition {
FieldDefinition {
name: Idiom::from_str(name).expect("valid idiom"),
field_kind: kind,
..Default::default()
}
}
fn planned(cols: &[&str], kind: Index, fields: &[FieldDefinition]) -> Vec<String> {
let indexes: Arc<[IndexDefinition]> = vec![idx_def(1, "ix", cols, kind)].into();
with_array_columns_as_elements(indexes, fields).indexes[0]
.cols
.iter()
.map(|c| c.to_sql())
.collect()
}
fn array_of_strings() -> Kind {
Kind::Array(Box::new(Kind::String), None)
}
#[test]
fn a_column_whose_kind_holds_only_arrays_is_planned_as_its_elements() {
let kinds = [
array_of_strings(),
Kind::Either(vec![Kind::None, array_of_strings()]),
Kind::Either(vec![Kind::None, Kind::Null, array_of_strings()]),
Kind::Literal(KindLiteral::Array(vec![Kind::String])),
];
for kind in kinds {
let fields = [field("acl", Some(kind.clone()))];
assert_eq!(planned(&["acl"], Index::Idx, &fields), ["acl.*"], "{kind:?}");
assert_eq!(planned(&["acl"], Index::Uniq, &fields), ["acl.*"], "{kind:?}");
}
}
#[test]
fn only_the_columns_that_admit_an_array_change() {
let fields =
[field("acl", Some(array_of_strings())), field("name", Some(Kind::String))];
assert_eq!(planned(&["name", "acl"], Index::Idx, &fields), ["name", "acl.*"]);
}
#[test]
fn a_column_that_can_store_a_whole_value_is_taken_as_written() {
let kinds = [
Kind::String,
Kind::Either(vec![Kind::None, Kind::Int]),
Kind::Set(Box::new(Kind::String), None),
Kind::Object,
Kind::Any,
Kind::Either(vec![Kind::String, array_of_strings()]),
Kind::Either(vec![Kind::None, Kind::Null]),
];
for kind in kinds {
let fields = [field("acl", Some(kind.clone()))];
assert_eq!(planned(&["acl"], Index::Idx, &fields), ["acl"], "{kind:?}");
}
}
#[test]
fn an_undeclared_kind_is_taken_as_written() {
assert_eq!(planned(&["acl"], Index::Idx, &[field("acl", None)]), ["acl"]);
assert_eq!(planned(&["acl"], Index::Idx, &[]), ["acl"]);
}
#[test]
fn a_column_already_over_elements_or_flattened_is_taken_as_written() {
let fields = [field("acl", Some(array_of_strings()))];
let fields_flat = [field("acl…", Some(array_of_strings()))];
assert_eq!(planned(&["acl.*"], Index::Idx, &fields), ["acl.*"]);
assert_eq!(planned(&["acl…"], Index::Idx, &fields_flat), ["acl…"]);
}
#[test]
fn only_btree_indexes_are_rewritten() {
let fields = [field("acl", Some(array_of_strings()))];
assert_eq!(planned(&["acl"], Index::Count(None), &fields), ["acl"]);
}
#[test]
fn a_union_is_read_through_at_any_depth() {
let nested = Kind::Either(vec![
Kind::None,
Kind::Either(vec![array_of_strings(), Kind::Array(Box::new(Kind::Int), None)]),
]);
assert_eq!(planned(&["acl"], Index::Idx, &[field("acl", Some(nested))]), ["acl.*"]);
let nested_scalar = Kind::Either(vec![
array_of_strings(),
Kind::Either(vec![Kind::None, Kind::String]),
]);
assert_eq!(
planned(&["acl"], Index::Idx, &[field("acl", Some(nested_scalar))]),
["acl"]
);
}
#[test]
fn an_ancestor_object_literal_declares_the_column_kind() {
let object = |kind: Kind| {
Kind::Literal(KindLiteral::Object(
[(Strand::from("tags"), kind)].into_iter().collect(),
))
};
let declared = [field("metadata", Some(object(array_of_strings())))];
assert_eq!(planned(&["metadata.tags"], Index::Idx, &declared), ["metadata.tags.*"]);
let optional = [field(
"metadata",
Some(Kind::Either(vec![Kind::None, object(array_of_strings())])),
)];
assert_eq!(planned(&["metadata.tags"], Index::Idx, &optional), ["metadata.tags.*"]);
let scalar = [field("metadata", Some(object(Kind::String)))];
assert_eq!(planned(&["metadata.tags"], Index::Idx, &scalar), ["metadata.tags"]);
let opaque = [field("metadata", Some(Kind::Object))];
assert_eq!(planned(&["metadata.tags"], Index::Idx, &opaque), ["metadata.tags"]);
let both = [
field("metadata", Some(object(array_of_strings()))),
field("metadata.tags", Some(Kind::String)),
];
assert_eq!(planned(&["metadata.tags"], Index::Idx, &both), ["metadata.tags"]);
let untyped_child =
[field("metadata", Some(object(array_of_strings()))), field("metadata.tags", None)];
assert_eq!(
planned(&["metadata.tags"], Index::Idx, &untyped_child),
["metadata.tags.*"]
);
}
#[test]
fn a_column_through_an_element_path_is_taken_as_written() {
let fields = [field("items[*].tags", Some(array_of_strings()))];
assert_eq!(planned(&["items[*].tags"], Index::Idx, &fields), ["items.*.tags"]);
}
#[test]
fn a_partly_rewritten_list_keeps_every_index_in_place() {
let indexes: Arc<[IndexDefinition]> = vec![
idx_def(1, "a", &["name"], Index::Idx),
idx_def(2, "b", &["acl"], Index::Idx),
idx_def(3, "c", &["name"], Index::Uniq),
]
.into();
let fields =
[field("acl", Some(array_of_strings())), field("name", Some(Kind::String))];
let out = with_array_columns_as_elements(indexes, &fields);
let flags: Vec<Vec<bool>> =
out.rewritten.as_deref().expect("rewritten").iter().map(|f| f.to_vec()).collect();
assert_eq!(flags, [vec![], vec![true], vec![]]);
let cols: Vec<_> =
out.indexes.iter().map(|ix| (ix.name.as_str(), ix.cols[0].to_sql())).collect();
assert_eq!(
cols,
[("a", "name".to_owned()), ("b", "acl.*".to_owned()), ("c", "name".to_owned())]
);
}
#[test]
fn the_field_list_is_needed_only_for_a_btree_column_without_elements() {
let with = |cols: &[&str], kind: Index| {
may_have_array_columns(&[idx_def(1, "ix", cols, kind)])
};
assert!(with(&["acl"], Index::Idx));
assert!(with(&["acl.*", "name"], Index::Uniq));
assert!(!with(&["acl.*"], Index::Idx));
assert!(!with(&["acl"], Index::Count(None)));
}
#[test]
fn withholding_btree_indexes_keeps_the_others() {
let indexes: Arc<[IndexDefinition]> = vec![
idx_def(1, "a", &["acl"], Index::Idx),
idx_ft(2, "b", &["body"]),
idx_def(3, "c", &["name"], Index::Uniq),
]
.into();
let names: Vec<_> = without_btree_indexes(indexes)
.iter()
.map(|ix| ix.name.as_str().to_owned())
.collect();
assert_eq!(names, ["b"]);
}
#[test]
fn an_unchanged_list_is_returned_as_is() {
let indexes: Arc<[IndexDefinition]> =
vec![idx_def(1, "ix", &["name"], Index::Idx)].into();
let fields = [field("name", Some(Kind::String))];
let out = with_array_columns_as_elements(Arc::clone(&indexes), &fields);
assert!(Arc::ptr_eq(&indexes, &out.indexes));
assert!(out.rewritten.is_none());
}
#[test]
fn every_column_of_a_wide_index_is_rewritten() {
let names: Vec<String> = (0..70).map(|c| format!("c{c}")).collect();
let cols: Vec<&str> = names.iter().map(String::as_str).collect();
let fields: Vec<_> = names.iter().map(|n| field(n, Some(array_of_strings()))).collect();
let indexes: Arc<[IndexDefinition]> = vec![idx_def(1, "ix", &cols, Index::Idx)].into();
let out = with_array_columns_as_elements(indexes, &fields);
assert_eq!(out.indexes[0].cols[69].to_sql(), "c69.*");
let flags = &out.rewritten.as_ref().expect("rewritten")[0];
assert_eq!(flags.len(), 70);
assert!(flags.iter().all(|f| *f));
}
}
}