use std::time::{Duration, Instant};
use yo_common::geo;
use yo_common::num::parse_f64;
use yo_common::{Result, parse_i64};
use yo_search::explain::{Note, Why};
use yo_search::field::{self, Algo, Coords, Kind, Tag, Text, Vector, Width};
use yo_search::follow::Errors;
use yo_search::index::{Definition, Sieve, Source};
use yo_search::query::parse::{BAD_POINT, BAD_RADIUS};
use yo_search::query::{self, Ask, Bad, Circle, Mask, Node, Pair, Range, What, Yield};
use yo_search::score::Scorer;
use yo_search::sorted::{self, Sorted};
use yo_search::spell::{self, Lists};
use yo_search::summary::{self, Trim, Wanted, Wrap};
use yo_search::tags::Tags;
use yo_search::walk;
use yo_search::{Clash, English, Field, Index, Registry};
use yo_shape::Metric;
use super::Server;
use super::args::{self, Args};
use super::indexing;
use super::table::Spec;
use crate::reply::Out;
mod aggregate;
mod config;
pub(super) mod cursor;
mod debug;
pub(super) mod docs;
mod hybrid;
mod profile;
use aggregate::{Pipe, Reads, Shape, apply, group, keeps, piped, sorts, windows};
use cursor::{Asks, Kept, Made};
pub(super) use hybrid::hybrid;
#[derive(Default)]
pub(super) struct Watch {
pub(super) parsing: Duration,
pub(super) creating: Duration,
pub(super) walking: Duration,
pub(super) ran: Option<walk::Ran>,
pub(super) steps: Vec<(Vec<u8>, usize)>,
}
const LANGUAGES: &[&[u8]] = &[
b"arabic",
b"armenian",
b"chinese",
b"danish",
b"dutch",
b"english",
b"finnish",
b"french",
b"german",
b"greek",
b"hungarian",
b"indonesian",
b"irish",
b"italian",
b"lithuanian",
b"nepali",
b"norwegian",
b"portuguese",
b"romanian",
b"russian",
b"serbian",
b"spanish",
b"swedish",
b"tamil",
b"turkish",
b"yiddish",
];
const PHONETICS: &[&[u8]] = &[b"dm:en", b"dm:fr", b"dm:pt", b"dm:es"];
struct Fail<'a> {
head: &'static str,
word: &'a [u8],
tail: &'static str,
whole: Option<Vec<u8>>,
}
impl<'a> Fail<'a> {
const fn plain(head: &'static str) -> Fail<'a> {
Fail {
head,
word: b"",
tail: "",
whole: None,
}
}
const fn naming(head: &'static str, word: &'a [u8]) -> Fail<'a> {
Fail {
head,
word,
tail: "",
whole: None,
}
}
const fn about(head: &'static str, word: &'a [u8], tail: &'static str) -> Fail<'a> {
Fail {
head,
word,
tail,
whole: None,
}
}
const fn said(whole: Vec<u8>) -> Fail<'a> {
Fail {
head: "",
word: b"",
tail: "",
whole: Some(whole),
}
}
fn write(&self, out: &mut Out) {
match &self.whole {
Some(line) => out.error(line),
None => out.error_about(self.head.as_bytes(), self.word, self.tail.as_bytes()),
}
}
}
type Answer<'a> = core::result::Result<(), Fail<'a>>;
const EXISTS: &str = "SEARCH_INDEX_EXISTS Index already exists";
const NOT_DB_ZERO: &str = "Cannot create index on db != 0";
const ALIAS_EXISTS: &str = "SEARCH_INDEX_EXISTS Alias already exists";
const MISSING: &str = "SEARCH_INDEX_NOT_FOUND Index not found: ";
const CONFLICT: &str = "SEARCH_ALIAS_CONFLICT Alias conflicts with an existing index name";
const NOT_MINE: &str = "SEARCH_INDEX_NOT_FOUND Alias does not belong to provided spec";
const NO_ALIAS: &str = "Alias does not exist";
const NO_FIELD: &str = "SEARCH_ATTR_BAD No such field";
const NOT_A_TAG: &str = "SEARCH_ATTR_BAD Not a tag field";
const NO_TARGET: &str = "SEARCH_INDEX_NOT_FOUND Unknown index name (or name is an alias itself)";
const NO_SCHEMA: &str = "SEARCH_PARSE_ARGS No schema found";
const NO_FIELDS: &str = "SEARCH_PARSE_ARGS Fields arguments are missing";
const AFTER_ALTER: &str = "ALTER must be followed by SCHEMA";
const ALTER_ACTION: &str = "Unknown action passed to ALTER SCHEMA";
const UNKNOWN: &str = "SEARCH_ARG_UNRECOGNIZED Unknown argument `";
const UNKNOWN_BARE: &str = "SEARCH_ARG_UNRECOGNIZED Unknown argument";
const NO_TYPE: &str = "SEARCH_PARSE_ARGS Field `";
const NO_TYPE_END: &str = "` does not have a type";
const BAD_TYPE: &str = "SEARCH_PARSE_ARGS Invalid field type for field `";
const DUPLICATE: &str = "SEARCH_QUERY_BAD Duplicate field in schema - ";
const RULE: &str = "SEARCH_ADD_ARGS Invalid rule type: ";
const LANGUAGE: &str = "SEARCH_ADD_ARGS Invalid language";
const SCORE: &str = "SEARCH_ADD_ARGS Invalid score";
const BOTH: &str =
"SEARCH_PARSE_ARGS 'Field cannot be defined with both `NOINDEX` and `INDEXMISSING` `";
const SEPARATOR: &str = "SEARCH_PARSE_ARGS Tag separator must be a single character. Got `%s`";
const MATCHER: &str = "SEARCH_QUERY_BAD Matcher Format: <2 chars algorithm>:<2 chars language>. Support algorithms: double metaphone (dm). Supported languages: English (en), French (fr), Portuguese (pt) and Spanish (es)";
const AS_ARG: &str = "SEARCH_PARSE_ARGS AS requires an argument";
const TOO_MANY_TEXT: &str = "SEARCH_QUERY_BAD MAXTEXTFIELDS cannot be used with NOFIELDS";
const BAD_ARGS: &str = "SEARCH_PARSE_ARGS Bad arguments for ";
const NOT_ENOUGH: &str =
"SEARCH_PARSE_ARGS Bad arguments for vector similarity: not enough arguments";
const ALGO_WORD: &str = "vector similarity algorithm";
const COUNT_WORD: &str = "vector similarity number of parameters";
const NO_TRAINING: &str =
"SEARCH_PARSE_ARGS TRAINING_THRESHOLD is irrelevant when compression was not requested";
const NO_REDUCE: &str =
"SEARCH_PARSE_ARGS REDUCE is irrelevant when compression is not of type LeanVec";
const SMALL_TRAINING: &str =
"SEARCH_PARSE_ARGS Invalid TRAINING_THRESHOLD: cannot be lower than DEFAULT_BLOCK_SIZE (1024)";
const ARITY: &str = "ERR wrong number of arguments for '_";
const ARITY_PLAIN: &str = "ERR wrong number of arguments for '";
const ARITY_END: &str = "' command";
const NOT_A_NUMBER: &str = ": Could not convert argument to expected type";
const OUT_OF_RANGE: &str = ": Value is outside acceptable bounds";
const NOT_THERE: &str = ": Expected an argument, but none provided";
const UNKNOWN_WORD: &str = ": Unknown argument";
const V_NOT_A_NUMBER: &str = "`: Could not convert argument to expected type";
const V_OUT_OF_RANGE: &str = "`: Value is outside acceptable bounds";
const V_NOT_THERE: &str = "`: Expected an argument, but none provided";
const V_UNKNOWN: &str = "`: Unknown argument";
fn bad(what: &'static str, why: &'static str) -> Fail<'static> {
Fail::about(BAD_ARGS, what.as_bytes(), why)
}
pub(super) struct Fill<'a> {
pub(super) name: &'a [u8],
pub(super) obeys: bool,
}
pub(super) enum After<'a> {
Scan(Fill<'a>),
Sweep(Vec<Box<[u8]>>),
}
pub(super) fn execute<'a>(
server: &Server,
reg: &mut Registry,
db: usize,
spec: &Spec,
args: Args<'a>,
out: &mut Out,
) -> Result<Option<After<'a>>> {
if spec.name == "FT.CONFIG" {
config::run(reg, args, out)?;
return Ok(None);
}
if spec.name == "_FT.DEBUG" {
debug::run(reg, args, out)?;
return Ok(None);
}
let mut made = None;
let done = match spec.name {
"FT.CREATE" => create(reg, db, args, out, false).map(|next| made = next),
"FT._CREATEIFNX" => create(reg, db, args, out, true).map(|next| made = next),
"FT.ALTER" => alter(reg, args, out, false),
"FT._ALTERIFNX" => alter(reg, args, out, true),
"FT.DROPINDEX" => drop_index(reg, spec, args, out, false, true).map(|next| made = next),
"FT._DROPINDEXIFX" => drop_index(reg, spec, args, out, true, true).map(|next| made = next),
"FT.DROP" => drop_index(reg, spec, args, out, false, false).map(|next| made = next),
"FT._DROPIFX" => drop_index(reg, spec, args, out, true, false).map(|next| made = next),
"FT.INFO" => info(server, reg, args, out),
"FT._LIST" => list(reg, spec, args, out),
"FT.ALIASADD" => alias_add(reg, args, out, false),
"FT._ALIASADDIFNX" => alias_add(reg, args, out, true),
"FT.ALIASDEL" => alias_del(reg, args, out, false),
"FT._ALIASDELIFX" => alias_del(reg, args, out, true),
"FT.ALIASUPDATE" => alias_update(reg, args, out),
"FT.ALIASLIST" => alias_list(reg, args, out),
"FT.EXPLAIN" => explain(reg, args, out, false),
"FT.EXPLAINCLI" => explain(reg, args, out, true),
"FT.TAGVALS" => tag_vals(reg, args, out),
"FT.DICTADD" => dict_add(reg, args, out),
"FT.DICTDEL" => dict_del(reg, args, out),
"FT.DICTDUMP" => dict_dump(reg, args, out),
"FT.SYNUPDATE" => syn_update(reg, args, out).map(|next| made = next),
"FT.SYNDUMP" => syn_dump(reg, args, out),
"FT.SPELLCHECK" => spellcheck(reg, args, out),
other => unreachable!("{other} is not a search command"),
};
if let Err(f) = done {
f.write(out);
return Ok(None);
}
if matches!(
spec.name,
"FT.DROPINDEX" | "FT._DROPINDEXIFX" | "FT.DROP" | "FT._DROPIFX"
) {
server.cursors.lock().strays(reg);
}
Ok(made)
}
fn create<'a>(
reg: &mut Registry,
db: usize,
args: Args<'a>,
out: &mut Out,
ifnx: bool,
) -> core::result::Result<Option<After<'a>>, Fail<'a>> {
let name = args.get(1);
if ifnx && reg.named(name).is_some() {
out.ok();
return Ok(None);
}
if db != 0 {
return Err(Fail::plain(NOT_DB_ZERO));
}
if reg.named(name).is_some() {
return Err(Fail::plain(EXISTS));
}
let (definition, at) = definition(args, 2)?;
let mut schema = Vec::new();
fields(args, at, &mut schema)?;
let _ = reg.create(Index::new(name, definition, schema));
out.ok();
Ok(Some(After::Scan(Fill { name, obeys: true })))
}
fn definition(args: Args<'_>, from: usize) -> core::result::Result<(Definition, usize), Fail<'_>> {
let mut d = Definition::default();
let mut prefixes: Option<Vec<Box<[u8]>>> = None;
let mut at = from;
loop {
let Some(a) = args.opt(at) else {
return Err(Fail::plain(NO_SCHEMA));
};
at += 1;
if args::is(a, b"schema") {
break;
} else if args::is(a, b"on") {
let v = value(args, &mut at, "ON")?;
d.on = if args::is(v, b"hash") {
Source::Hash
} else if args::is(v, b"json") {
Source::Json
} else {
return Err(Fail::naming(RULE, v));
};
} else if args::is(a, b"prefix") {
let n = count(args, &mut at, "PREFIX")?;
let mut list = Vec::with_capacity(n);
for _ in 0..n {
list.push(value(args, &mut at, "PREFIX")?.into());
}
prefixes = Some(list);
} else if args::is(a, b"filter") {
let src = value(args, &mut at, "FILTER")?;
Sieve::parse(src).map_err(Fail::said)?;
d.filter = Some(src.into());
} else if args::is(a, b"language") {
let v = value(args, &mut at, "LANGUAGE")?;
d.language = Some(language(v).ok_or(Fail::plain(LANGUAGE))?.into());
} else if args::is(a, b"language_field") {
d.language_field = Some(value(args, &mut at, "LANGUAGE_FIELD")?.into());
} else if args::is(a, b"score") {
let v = value(args, &mut at, "SCORE")?;
d.score = match parse_f64(v) {
Some(s) if (0.0..=1.0).contains(&s) => s,
_ => return Err(Fail::plain(SCORE)),
};
} else if args::is(a, b"score_field") {
d.score_field = Some(value(args, &mut at, "SCORE_FIELD")?.into());
} else if args::is(a, b"payload_field") {
d.payload_field = Some(value(args, &mut at, "PAYLOAD_FIELD")?.into());
} else if args::is(a, b"temporary") {
d.temporary = Some(count(args, &mut at, "TEMPORARY")? as u64);
} else if args::is(a, b"stopwords") {
let n = count(args, &mut at, "STOPWORDS")?;
let mut list = Vec::with_capacity(n);
for _ in 0..n {
list.push(value(args, &mut at, "STOPWORDS")?.into());
}
d.stopwords = Some(list);
} else if args::is(a, b"maxtextfields") {
d.options.maxtextfields = true;
} else if args::is(a, b"nooffsets") {
d.options.nooffsets = true;
} else if args::is(a, b"nohl") {
d.options.nohl = true;
} else if args::is(a, b"nofields") {
d.options.nofields = true;
} else if args::is(a, b"nofreqs") {
d.options.nofreqs = true;
} else if args::is(a, b"skipinitialscan") {
d.skip_initial_scan = true;
} else {
return Err(Fail::about(UNKNOWN, a, "`"));
}
}
if d.options.maxtextfields && d.options.nofields {
return Err(Fail::plain(TOO_MANY_TEXT));
}
if let Some(list) = prefixes
&& !list.is_empty()
{
d.prefixes = list;
}
Ok((d, at))
}
fn language(v: &[u8]) -> Option<&'static [u8]> {
LANGUAGES.iter().copied().find(|l| args::is(v, l))
}
fn value<'a>(
args: Args<'a>,
at: &mut usize,
what: &'static str,
) -> core::result::Result<&'a [u8], Fail<'a>> {
let v = args.opt(*at).ok_or_else(|| bad(what, NOT_THERE))?;
*at += 1;
Ok(v)
}
fn count<'a>(
args: Args<'a>,
at: &mut usize,
what: &'static str,
) -> core::result::Result<usize, Fail<'a>> {
let v = value(args, at, what)?;
let n = parse_i64(v).ok_or_else(|| bad(what, NOT_A_NUMBER))?;
usize::try_from(n).map_err(|_| bad(what, OUT_OF_RANGE))
}
fn fields<'a>(args: Args<'a>, from: usize, into: &mut Vec<Field>) -> Answer<'a> {
if args.opt(from).is_none() {
return Err(Fail::plain(NO_FIELDS));
}
let mut at = from;
while at < args.len() {
let f = one_field(args, &mut at, into)?;
into.push(f);
}
Ok(())
}
fn one_field<'a>(
args: Args<'a>,
at: &mut usize,
have: &[Field],
) -> core::result::Result<Field, Fail<'a>> {
let identifier = args.get(*at);
*at += 1;
let attribute = if args.opt(*at).is_some_and(|a| args::is(a, b"as")) {
*at += 1;
let named = args.opt(*at).ok_or(Fail::plain(AS_ARG))?;
*at += 1;
named
} else {
identifier
};
let kind = kind(args, at, attribute)?;
let mut f = Field::new(identifier, kind).named(attribute);
if have.iter().any(|o| o.attribute == f.attribute) {
return Err(Fail::naming(DUPLICATE, attribute));
}
while let Some(a) = args.opt(*at) {
if args::is(a, b"withsuffixtrie") && f.kind.takes_empty() {
f.suffix_trie = true;
} else if args::is(a, b"indexempty") && f.kind.takes_empty() {
f.index_empty = true;
} else if args::is(a, b"indexmissing") {
f.index_missing = true;
} else if !type_option(args, at, &mut f)? {
break;
}
*at += 1;
}
while let Some(a) = args.opt(*at) {
if args::is(a, b"sortable") {
f.sortable = true;
*at += 1;
if args.opt(*at).is_some_and(|n| args::is(n, b"unf")) {
f.unf = true;
*at += 1;
}
} else if args::is(a, b"noindex") {
f.noindex = true;
*at += 1;
} else {
break;
}
}
if f.noindex && f.index_missing {
return Err(Fail::about(BOTH, attribute, "` '"));
}
Ok(f)
}
fn kind<'a>(
args: Args<'a>,
at: &mut usize,
attribute: &'a [u8],
) -> core::result::Result<Kind, Fail<'a>> {
let Some(t) = args.opt(*at) else {
return Err(Fail::about(NO_TYPE, attribute, NO_TYPE_END));
};
*at += 1;
if args::is(t, b"text") {
Ok(Kind::Text(Text::default()))
} else if args::is(t, b"tag") {
Ok(Kind::Tag(Tag::default()))
} else if args::is(t, b"numeric") {
Ok(Kind::Numeric)
} else if args::is(t, b"geo") {
Ok(Kind::Geo)
} else if args::is(t, b"geoshape") {
let coords = match args.opt(*at) {
Some(c) if args::is(c, b"flat") => {
*at += 1;
Coords::Flat
}
Some(c) if args::is(c, b"spherical") => {
*at += 1;
Coords::Spherical
}
_ => Coords::Spherical,
};
Ok(Kind::GeoShape(coords))
} else if args::is(t, b"vector") {
vector(args, at).map(Kind::Vector)
} else {
Err(Fail::about(BAD_TYPE, attribute, "`"))
}
}
fn type_option<'a>(
args: Args<'a>,
at: &mut usize,
f: &mut Field,
) -> core::result::Result<bool, Fail<'a>> {
let a = args.get(*at);
match &mut f.kind {
Kind::Text(t) => {
if args::is(a, b"nostem") {
t.nostem = true;
} else if args::is(a, b"weight") {
let mut next = *at + 1;
let v = value(args, &mut next, "weight")?;
t.weight = parse_f64(v).ok_or_else(|| bad("weight", NOT_A_NUMBER))?;
*at = next - 1;
} else if args::is(a, b"phonetic") {
let mut next = *at + 1;
let v = value(args, &mut next, "PHONETIC")?;
if !PHONETICS.iter().any(|p| args::is(v, p)) {
return Err(Fail::plain(MATCHER));
}
t.phonetic = Some(v.into());
*at = next - 1;
} else {
return Ok(false);
}
Ok(true)
}
Kind::Tag(t) => {
if args::is(a, b"casesensitive") {
t.casesensitive = true;
} else if args::is(a, b"separator") {
let mut next = *at + 1;
let v = value(args, &mut next, "SEPARATOR")?;
let [c] = v else {
return Err(Fail::naming(SEPARATOR, v));
};
t.separator = *c;
*at = next - 1;
} else {
return Ok(false);
}
Ok(true)
}
_ => Ok(false),
}
}
fn vector<'a>(args: Args<'a>, at: &mut usize) -> core::result::Result<Vector, Fail<'a>> {
let a = args.opt(*at).ok_or_else(|| bad(ALGO_WORD, NOT_THERE))?;
*at += 1;
let (algo, label) = if args::is(a, b"flat") {
(Algo::Flat, "FLAT")
} else if args::is(a, b"hnsw") {
(Algo::Hnsw, "HNSW")
} else if args::is(a, b"svs-vamana") {
(Algo::Svs, "SVS-VAMANA")
} else {
return Err(bad(ALGO_WORD, UNKNOWN_WORD));
};
let c = args.opt(*at).ok_or_else(|| bad(COUNT_WORD, NOT_THERE))?;
*at += 1;
let n = parse_i64(c).ok_or_else(|| bad(COUNT_WORD, NOT_A_NUMBER))?;
let words = usize::try_from(n).map_err(|_| bad(COUNT_WORD, OUT_OF_RANGE))?;
if args.len() - *at < words {
return Err(Fail::plain(NOT_ENOUGH));
}
let mut width = None;
let mut dim = None;
let mut metric = None;
let mut training: Option<u64> = None;
let mut reduce = false;
let mut v = Vector::new(algo, Width::Float32, 0, Metric::L2);
let mut left = words;
while left > 0 {
let key = args.get(*at);
*at += 1;
left -= 1;
if left == 0 {
return Err(vbad(label, key, V_NOT_THERE));
}
let val = args.get(*at);
*at += 1;
left -= 1;
if args::is(key, b"type") {
width = Some(self::width(val).ok_or_else(|| vbad(label, key, V_UNKNOWN))?);
} else if args::is(key, b"dim") {
let d = parse_i64(val).ok_or_else(|| vbad(label, key, V_NOT_A_NUMBER))?;
if d <= 0 {
return Err(vbad(label, key, V_OUT_OF_RANGE));
}
dim = Some(d as u64);
} else if args::is(key, b"distance_metric") {
metric = Some(self::metric(val).ok_or_else(|| vbad(label, key, V_UNKNOWN))?);
} else if args::is(key, b"initial_cap") && algo != Algo::Svs {
v.initial_cap = Some(whole(label, key, val)?);
} else if args::is(key, b"block_size") && algo == Algo::Flat {
v.block_size = Some(whole(label, key, val)?);
} else if args::is(key, b"m") && algo == Algo::Hnsw {
v.m = whole(label, key, val)?;
} else if args::is(key, b"ef_construction") && algo == Algo::Hnsw {
v.ef_construction = whole(label, key, val)?;
} else if args::is(key, b"ef_runtime") && algo == Algo::Hnsw {
v.ef_runtime = whole(label, key, val)?;
} else if args::is(key, b"epsilon") && algo != Algo::Flat {
v.epsilon = parse_f64(val).ok_or_else(|| vbad(label, key, V_NOT_A_NUMBER))?;
} else if args::is(key, b"graph_max_degree") && algo == Algo::Svs {
v.graph_max_degree = whole(label, key, val)?;
} else if args::is(key, b"construction_window_size") && algo == Algo::Svs {
v.construction_window = whole(label, key, val)?;
} else if args::is(key, b"search_window_size") && algo == Algo::Svs {
let _ = whole(label, key, val)?;
} else if args::is(key, b"compression") && algo == Algo::Svs {
let c = compression(val).ok_or_else(|| vbad(label, key, V_UNKNOWN))?;
v.compression = Some(c.as_bytes().into());
} else if args::is(key, b"training_threshold") && algo == Algo::Svs {
training = Some(whole(label, key, val)?);
} else if args::is(key, b"reduce") && algo == Algo::Svs {
let _ = whole(label, key, val)?;
reduce = true;
} else {
return Err(unwanted(label, key));
}
}
if let Some(t) = training {
if v.compression.is_none() {
return Err(Fail::plain(NO_TRAINING));
}
if t < field::MIN_TRAINING {
return Err(Fail::plain(SMALL_TRAINING));
}
v.training_threshold = Some(t);
}
if reduce && !v.compression.as_deref().is_some_and(is_leanvec) {
return Err(Fail::plain(NO_REDUCE));
}
v.width = width.ok_or_else(|| missing(label, "TYPE"))?;
v.dim = dim.ok_or_else(|| missing(label, "DIM"))?;
v.metric = metric.ok_or_else(|| missing(label, "DISTANCE_METRIC"))?;
Ok(v)
}
const COMPRESSIONS: &[&str] = &[
"LVQ8",
"LVQ4",
"LVQ4x4",
"LVQ4x8",
"LeanVec4x8",
"LeanVec8x8",
];
fn compression(v: &[u8]) -> Option<&'static str> {
COMPRESSIONS
.iter()
.copied()
.find(|c| args::is(v, c.as_bytes()))
}
fn is_leanvec(c: &[u8]) -> bool {
c.starts_with(b"LeanVec")
}
fn unwanted<'a>(label: &'static str, key: &'a [u8]) -> Fail<'a> {
Fail::naming(
match label {
"FLAT" => "SEARCH_PARSE_ARGS Bad arguments for algorithm FLAT: ",
"HNSW" => "SEARCH_PARSE_ARGS Bad arguments for algorithm HNSW: ",
_ => "SEARCH_PARSE_ARGS Bad arguments for algorithm SVS-VAMANA: ",
},
key,
)
}
fn vbad<'a>(label: &'static str, key: &'a [u8], why: &'static str) -> Fail<'a> {
Fail::about(
match label {
"FLAT" => "SEARCH_PARSE_ARGS Bad arguments for vector similarity FLAT index `",
"HNSW" => "SEARCH_PARSE_ARGS Bad arguments for vector similarity HNSW index `",
_ => "SEARCH_PARSE_ARGS Bad arguments for vector similarity SVS-VAMANA index `",
},
key,
why,
)
}
fn missing(label: &'static str, what: &'static str) -> Fail<'static> {
Fail::about(
match label {
"FLAT" => {
"SEARCH_PARSE_ARGS Missing mandatory parameter: cannot create FLAT index without specifying "
}
"HNSW" => {
"SEARCH_PARSE_ARGS Missing mandatory parameter: cannot create HNSW index without specifying "
}
_ => {
"SEARCH_PARSE_ARGS Missing mandatory parameter: cannot create SVS-VAMANA index without specifying "
}
},
what.as_bytes(),
" argument",
)
}
fn whole<'a>(
label: &'static str,
key: &'a [u8],
val: &[u8],
) -> core::result::Result<u64, Fail<'a>> {
let n = parse_i64(val).ok_or_else(|| vbad(label, key, V_NOT_A_NUMBER))?;
u64::try_from(n).map_err(|_| vbad(label, key, V_OUT_OF_RANGE))
}
fn width(v: &[u8]) -> Option<Width> {
[
Width::Int8,
Width::Uint8,
Width::Float16,
Width::BFloat16,
Width::Float32,
Width::Float64,
]
.into_iter()
.find(|w| args::is(v, w.token().as_bytes()))
}
fn metric(v: &[u8]) -> Option<Metric> {
if args::is(v, b"l2") {
Some(Metric::L2)
} else if args::is(v, b"ip") {
Some(Metric::Ip)
} else if args::is(v, b"cosine") {
Some(Metric::Cosine)
} else {
None
}
}
fn alter<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out, ifnx: bool) -> Answer<'a> {
let name = args.get(1);
reg.touch(name);
let mut at = 2;
if args
.opt(at)
.is_some_and(|a| args::is(a, b"skipinitialscan"))
{
at += 1;
}
if !args.opt(at).is_some_and(|a| args::is(a, b"schema")) {
return Err(Fail::plain(AFTER_ALTER));
}
at += 1;
if !args.opt(at).is_some_and(|a| args::is(a, b"add")) {
return Err(Fail::plain(ALTER_ACTION));
}
at += 1;
let Some(index) = reg.get(name) else {
return Err(Fail::naming(MISSING, name));
};
let mut schema = index.schema.clone();
match fields(args, at, &mut schema) {
Ok(()) => {}
Err(f) if ifnx && f.head == DUPLICATE => {
out.ok();
return Ok(());
}
Err(f) => return Err(f),
}
if let Some(index) = reg.get_mut(name) {
index.schema = schema;
}
out.ok();
Ok(())
}
fn drop_index<'a>(
reg: &mut Registry,
spec: &Spec,
args: Args<'a>,
out: &mut Out,
ifx: bool,
newer: bool,
) -> core::result::Result<Option<After<'a>>, Fail<'a>> {
if args.len() < 2 {
return Err(Fail::about(ARITY_PLAIN, spec.name.as_bytes(), ARITY_END));
}
if args.len() > 3 {
return Err(Fail::about(ARITY, spec.name.as_bytes(), ARITY_END));
}
let name = args.get(1);
if !reg.touch(name) {
if ifx {
out.ok();
return Ok(None);
}
return Err(Fail::naming(MISSING, name));
}
let word = match newer {
true => &b"dd"[..],
false => &b"keepdocs"[..],
};
let said = match args.opt(2) {
Some(a) if args::is(a, word) => true,
Some(_) => return Err(Fail::plain(UNKNOWN_BARE)),
None => false,
};
let index = reg.drop(name);
out.ok();
if said != newer {
return Ok(None);
}
Ok(index.ok().map(|index| {
After::Sweep(
index
.held
.docs
.all()
.map(|(_, doc)| doc.key.clone())
.collect(),
)
}))
}
fn list<'a>(reg: &Registry, spec: &Spec, args: Args<'a>, out: &mut Out) -> Answer<'a> {
if args.len() > 2 {
return Err(Fail::about(ARITY_PLAIN, spec.name.as_bytes(), ARITY_END));
}
out.set(reg.len());
for index in reg.iter() {
word(out, &index.name);
}
Ok(())
}
fn alias_add<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out, ifnx: bool) -> Answer<'a> {
let alias = args.get(1);
let name = args.get(2);
if reg.named(name).is_some() {
reg.touch(name);
}
if ifnx
&& let Some(at) = reg.target(alias)
&& reg.named(name).is_some_and(|i| *i.name == *at)
{
out.ok();
return Ok(());
}
match reg.alias(alias, name) {
Ok(()) => out.ok(),
Err(Clash::IsIndex) => return Err(Fail::plain(CONFLICT)),
Err(Clash::Aliased) => return Err(Fail::plain(ALIAS_EXISTS)),
Err(Clash::IsAlias) => return Err(Fail::plain(NO_TARGET)),
Err(_) if ifnx => return Err(Fail::plain(NO_TARGET)),
Err(_) => return Err(Fail::naming(MISSING, name)),
}
Ok(())
}
fn alias_update<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out) -> Answer<'a> {
let alias = args.get(1);
let name = args.get(2);
let old: Option<Box<[u8]>> = reg.target(alias).map(Into::into);
match reg.realias(alias, name) {
Ok(()) => {
if let Some(old) = old {
reg.touch(&old);
}
reg.touch(name);
out.ok();
}
Err(Clash::IsIndex) => return Err(Fail::plain(NOT_MINE)),
Err(Clash::IsAlias) => return Err(Fail::plain(NO_TARGET)),
Err(_) => return Err(Fail::naming(MISSING, name)),
}
Ok(())
}
fn alias_del<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out, ifx: bool) -> Answer<'a> {
let alias = args.get(1);
if ifx {
reg.touch(alias);
}
reg.touch(alias);
if reg.named(alias).is_some() {
return Err(Fail::plain(NOT_MINE));
}
match reg.unalias(alias) {
Ok(()) => out.ok(),
Err(_) if ifx => out.ok(),
Err(_) => return Err(Fail::plain(NO_ALIAS)),
}
Ok(())
}
fn alias_list<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out) -> Answer<'a> {
let name = args.get(1);
if reg.named(name).is_none() {
return Err(Fail::naming(MISSING, name));
}
let real: Box<[u8]> = name.into();
reg.open(&real);
let n = reg.aliases_of(&real).count();
out.set(n);
for alias in reg.aliases_of(&real) {
out.bulk(alias);
}
Ok(())
}
fn tag_vals<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out) -> Answer<'a> {
let name = args.get(1);
let attribute = args.get(2);
let Some(index) = reg.open(name) else {
return Err(Fail::naming(MISSING, name));
};
match index.field(attribute) {
None => return Err(Fail::plain(NO_FIELD)),
Some(field) if !matches!(field.kind, Kind::Tag(_)) => {
return Err(Fail::plain(NOT_A_TAG));
}
Some(_) => {}
}
let values = index.held.values(attribute);
out.set(values.map_or(0, Tags::len));
for (value, _) in values.into_iter().flat_map(Tags::all) {
out.bulk(value);
}
Ok(())
}
fn dict_add<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out) -> Answer<'a> {
let name = args.get(1);
let terms: Vec<&[u8]> = (2..args.len()).map(|i| args.get(i)).collect();
out.int(reg.dicts.add(name, &terms) as i64);
Ok(())
}
fn dict_del<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out) -> Answer<'a> {
let name = args.get(1);
let terms: Vec<&[u8]> = (2..args.len()).map(|i| args.get(i)).collect();
out.int(reg.dicts.del(name, &terms) as i64);
Ok(())
}
fn dict_dump<'a>(reg: &Registry, args: Args<'a>, out: &mut Out) -> Answer<'a> {
let name = args.get(1);
out.set(reg.dicts.len(name));
for term in reg.dicts.dump(name) {
out.bulk(term);
}
Ok(())
}
fn syn_update<'a>(
reg: &mut Registry,
args: Args<'a>,
out: &mut Out,
) -> core::result::Result<Option<After<'a>>, Fail<'a>> {
let name = args.get(1);
let group = args.get(2);
let skip = args.opt(3).is_some_and(|a| args::is(a, b"skipinitialscan"));
let from = if skip { 4 } else { 3 };
let terms: Vec<&[u8]> = (from..args.len()).map(|i| args.get(i)).collect();
let Some(index) = reg.open(name) else {
return Err(Fail::naming(MISSING, name));
};
index.synonyms.update(group, &terms);
out.ok();
Ok((!skip).then_some(After::Scan(Fill { name, obeys: false })))
}
fn syn_dump<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out) -> Answer<'a> {
let name = args.get(1);
let Some(index) = reg.open(name) else {
return Err(Fail::naming(MISSING, name));
};
out.map(index.synonyms.len());
for (term, ids) in index.synonyms.dump() {
out.bulk(term);
out.array(ids.len());
for id in ids {
out.bulk(id);
}
}
Ok(())
}
const TERM: &[u8] = b"TERM";
const RESULTS: &[u8] = b"results";
const BAD_DISTANCE: &str = "bad distance given, distance must be a natural number between 1 to 4";
const NEED_DISTANCE: &str = "DISTANCE arg is given but no DISTANCE comes after";
const NEED_TERMS: &str = "TERM arg is given but no TERM params comes after";
const BAD_TERMS: &str = "bad format, exclude/include operation was not given";
const NO_DICT: &str = "Dict does not exist: ";
struct Spelling<'a> {
distance: u8,
fixed: bool,
include: Vec<&'a [u8]>,
exclude: Vec<&'a [u8]>,
asked: Asked<'a>,
}
fn spellcheck<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out) -> Answer<'a> {
let name = args.get(1);
let query = args.get(2);
if reg.open(name).is_none() {
return Err(Fail::naming(MISSING, name));
}
let asked = match spelling(args, 3) {
Ok(asked) => asked,
Err(text) => {
out.error(&text);
return Ok(());
}
};
for named in asked.include.iter().chain(&asked.exclude) {
if reg.dicts.len(named) == 0 {
out.error(&line(NO_DICT, named, ""));
return Ok(());
}
}
let lists = Lists {
include: asked
.include
.iter()
.flat_map(|n| reg.dicts.dump(n))
.collect(),
exclude: asked
.exclude
.iter()
.flat_map(|n| reg.dicts.dump(n))
.collect(),
};
let index = reg.named(name).expect("an index that was just opened");
let ask = Ask {
dialect: asked.asked.dialect,
params: &asked.asked.params,
verbatim: false,
stopwords: false,
};
let node = match query::parse(query, index, &ask) {
Ok(node) => node,
Err(bad) => {
out.error(&refused(&bad));
return Ok(());
}
};
let found = spell::check(index, &node, asked.distance, &lists);
if out.proto().is_resp3() {
out.map(1);
out.bulk(RESULTS);
out.map(found.len());
for checked in &found {
out.bulk(&checked.word);
out.array(checked.guesses.len());
for guess in &checked.guesses {
out.map(1);
out.bulk(&guess.term);
out.double(guess.score);
}
}
return Ok(());
}
out.array(found.len());
for checked in &found {
out.array(3);
out.bulk(TERM);
out.bulk(&checked.word);
out.array(checked.guesses.len());
for guess in &checked.guesses {
out.array(2);
out.double(guess.score);
out.bulk(&guess.term);
}
}
Ok(())
}
fn spelling(args: Args<'_>, from: usize) -> core::result::Result<Spelling<'_>, Vec<u8>> {
let mut spelling = Spelling {
distance: spell::NEAREST,
fixed: false,
include: Vec::new(),
exclude: Vec::new(),
asked: Asked::default(),
};
let mut at = from;
while at < args.len() {
let word = args.get(at);
if args::is(word, b"DISTANCE") {
let Some(value) = args.opt(at + 1) else {
return Err(NEED_DISTANCE.as_bytes().to_vec());
};
let Some(distance) = parse_i64(value)
.filter(|d| (i64::from(spell::NEAREST)..=i64::from(spell::FURTHEST)).contains(d))
else {
return Err(BAD_DISTANCE.as_bytes().to_vec());
};
if !spelling.fixed {
spelling.distance = u8::try_from(distance).unwrap_or(spell::NEAREST);
spelling.fixed = true;
}
at += 2;
continue;
}
if args::is(word, b"TERMS") {
let (Some(how), Some(named)) = (args.opt(at + 1), args.opt(at + 2)) else {
return Err(NEED_TERMS.as_bytes().to_vec());
};
if args::is(how, b"INCLUDE") {
spelling.include.push(named);
} else if args::is(how, b"EXCLUDE") {
spelling.exclude.push(named);
} else {
return Err(BAD_TERMS.as_bytes().to_vec());
}
at += 3;
continue;
}
if args::is(word, b"DIALECT") {
let Some(value) = args.opt(at + 1) else {
return Err(line(NEED_ARG, b"DIALECT", ""));
};
let Some(dialect) = parse_i64(value).filter(|d| (1..=i64::from(NEWEST)).contains(d))
else {
return Err(BAD_DIALECT.as_bytes().to_vec());
};
spelling.asked.dialect = u8::try_from(dialect).unwrap_or(1);
at += 2;
continue;
}
if args::is(word, b"PARAMS") {
at = params(args, at, &mut spelling.asked)?;
continue;
}
at += 1;
}
Ok(spelling)
}
fn word(out: &mut Out, s: &[u8]) {
if s.contains(&b'\r') || s.contains(&b'\n') {
out.bulk(s);
} else {
out.simple(s);
}
}
fn pair(out: &mut Out, name: &str, value: &[u8]) {
out.simple(name.as_bytes());
word(out, value);
}
fn number(out: &mut Out, name: &str, value: f64) {
out.simple(name.as_bytes());
out.double(value);
}
fn tally(out: &mut Out, name: &str, value: u64) {
out.simple(name.as_bytes());
out.uint(value);
}
fn info<'a>(server: &Server, reg: &mut Registry, args: Args<'a>, out: &mut Out) -> Answer<'a> {
let name = args.get(1);
let Some(index) = reg.open(name) else {
return Err(Fail::naming(MISSING, name));
};
let (whole, mine) = cursor::stats(server, &index.name);
let idle = whole;
let d = &index.definition;
let stopwords = d.stopwords.as_ref();
out.map(34 + usize::from(stopwords.is_some()));
pair(out, "index_name", &index.name);
out.simple(b"index_options");
let tokens = d.options.tokens();
out.array(tokens.len());
for t in &tokens {
out.simple(t.as_bytes());
}
out.simple(b"index_definition");
let mut pairs = 4;
pairs += usize::from(d.filter.is_some());
pairs += usize::from(d.language.is_some());
pairs += usize::from(d.language_field.is_some());
pairs += usize::from(d.score_field.is_some());
pairs += usize::from(d.payload_field.is_some());
out.map(pairs);
pair(out, "key_type", d.on.token().as_bytes());
out.simple(b"prefixes");
out.array(d.prefixes.len());
for p in &d.prefixes {
word(out, p);
}
if let Some(f) = &d.filter {
pair(out, "filter", f);
}
if let Some(l) = &d.language {
pair(out, "default_language", l);
}
if let Some(l) = &d.language_field {
pair(out, "language_field", l);
}
number(out, "default_score", d.score);
if let Some(s) = &d.score_field {
pair(out, "score_field", s);
}
if let Some(p) = &d.payload_field {
pair(out, "payload_field", p);
}
pair(out, "indexes_all", b"false");
out.simple(b"attributes");
out.array(index.schema.len());
for f in &index.schema {
attribute(out, f);
}
let docs = index.held.docs.len() as u64;
let records = index.held.records();
tally(out, "num_docs", docs);
tally(out, "max_doc_id", u64::from(index.held.docs.last()));
tally(out, "num_terms", index.held.words() as u64);
tally(out, "num_records", records);
number(out, "inverted_sz_mb", 0.0);
number(out, "vector_index_sz_mb", 0.0);
tally(out, "total_inverted_index_blocks", 0);
number(out, "offset_vectors_sz_mb", 0.0);
number(out, "doc_table_size_mb", 0.0);
number(out, "sortable_values_size_mb", 0.0);
number(out, "key_table_size_mb", 0.0);
number(out, "tag_overhead_sz_mb", 0.0);
number(out, "text_overhead_sz_mb", 0.0);
number(out, "total_index_memory_sz_mb", 0.0);
number(out, "geoshapes_sz_mb", 0.0);
number(
out,
"records_per_doc_avg",
f64::from(records as f32 / docs as f32),
);
number(out, "bytes_per_record_avg", f64::NAN);
number(out, "offsets_per_term_avg", f64::NAN);
number(out, "offset_bits_per_record_avg", f64::NAN);
tally(
out,
"hash_indexing_failures",
index.trouble.whole().failures(),
);
number(out, "total_indexing_time", 0.0);
tally(out, "indexing", 0);
number(out, "percent_indexed", 1.0);
tally(out, "number_of_uses", index.uses);
tally(out, "cleaning", 0);
out.simple(b"gc_stats");
out.map(7);
number(out, "bytes_collected", 0.0);
number(out, "total_ms_run", 0.0);
number(out, "total_cycles", 0.0);
number(out, "average_cycle_time_ms", f64::NAN);
number(out, "last_run_time_ms", 0.0);
number(out, "gc_numeric_trees_missed", 0.0);
number(out, "gc_blocks_denied", 0.0);
out.simple(b"cursor_stats");
out.map(4);
tally(out, "global_idle", idle);
tally(out, "global_total", whole);
tally(out, "index_capacity", CURSOR_CAPACITY);
tally(out, "index_total", mine);
if let Some(list) = stopwords {
out.simple(b"stopwords_list");
out.array(list.len());
for w in list {
out.bulk(w);
}
}
out.simple(b"dialect_stats");
out.map(4);
for n in 1..=4 {
out.simple(match n {
1 => b"dialect_1".as_slice(),
2 => b"dialect_2",
3 => b"dialect_3",
_ => b"dialect_4",
});
out.uint(0);
}
out.simple(b"Index Errors");
out.map(4);
errors(out, index.trouble.whole());
out.simple(b"background indexing status");
out.simple(b"OK");
out.simple(b"field statistics");
out.array(index.schema.len());
for f in &index.schema {
statistics(out, f, index.trouble.field(&f.attribute));
}
Ok(())
}
const CURSOR_CAPACITY: u64 = 128;
fn attribute(out: &mut Out, f: &Field) {
let mut flags: Vec<&str> = Vec::new();
if f.sortable {
flags.push("SORTABLE");
}
if f.is_unf() {
flags.push("UNF");
}
if let Kind::Text(t) = &f.kind
&& t.nostem
{
flags.push("NOSTEM");
}
if let Kind::Tag(t) = &f.kind
&& t.casesensitive
{
flags.push("CASESENSITIVE");
}
if f.suffix_trie {
flags.push("WITHSUFFIXTRIE");
}
if f.index_empty {
flags.push("INDEXEMPTY");
}
if f.index_missing {
flags.push("INDEXMISSING");
}
if f.noindex {
flags.push("NOINDEX");
}
let pairs = 3 + match &f.kind {
Kind::Text(_) | Kind::Tag(_) => 1,
Kind::GeoShape(_) => 1,
Kind::Vector(v) => match v.algo {
Algo::Flat => 4,
Algo::Hnsw => 7,
Algo::Svs => 7 + usize::from(v.compression.is_some()),
},
Kind::Numeric | Kind::Geo => 0,
};
if out.proto().is_resp3() {
out.map(pairs + 1);
} else {
out.array(pairs * 2 + flags.len());
}
pair(out, "identifier", &f.identifier);
pair(out, "attribute", &f.attribute);
pair(out, "type", f.kind.token().as_bytes());
match &f.kind {
Kind::Text(t) => number(out, "WEIGHT", t.weight),
Kind::Tag(t) => {
out.simple(b"SEPARATOR");
word(out, &[t.separator]);
}
Kind::GeoShape(c) => pair(out, "coord_system", c.token().as_bytes()),
Kind::Vector(v) => {
pair(out, "algorithm", v.algo.token().as_bytes());
pair(out, "data_type", v.width.token().as_bytes());
tally(out, "dim", v.dim);
pair(out, "distance_metric", v.metric_token().as_bytes());
match v.algo {
Algo::Flat => {}
Algo::Hnsw => {
tally(out, "M", v.m);
tally(out, "ef_construction", v.ef_construction);
tally(out, "ef_runtime", v.ef_runtime);
}
Algo::Svs => {
tally(out, "graph_max_degree", v.graph_max_degree);
tally(out, "construction_window_size", v.construction_window);
pair(
out,
"compression",
v.compression
.as_deref()
.unwrap_or(field::NO_COMPRESSION.as_bytes()),
);
if v.compression.is_some() {
tally(
out,
"training_threshold",
v.training_threshold.unwrap_or(field::TRAINING_THRESHOLD),
);
}
}
}
}
Kind::Numeric | Kind::Geo => {}
}
if out.proto().is_resp3() {
out.simple(b"flags");
out.array(flags.len());
}
for flag in &flags {
out.simple(flag.as_bytes());
}
}
fn errors(out: &mut Out, e: &Errors) {
tally(out, "indexing failures", e.failures());
out.simple(b"last indexing error");
out.simple(e.sentence());
out.simple(b"last indexing error key");
out.bulk(e.about());
}
fn statistics(out: &mut Out, f: &Field, e: &Errors) {
let vector = matches!(f.kind, Kind::Vector(_));
out.map(3 + if vector { 4 } else { 0 });
pair(out, "identifier", &f.identifier);
pair(out, "attribute", &f.attribute);
out.simple(b"Index Errors");
out.map(3);
errors(out, e);
if vector {
tally(out, "memory", 0);
tally(out, "marked_deleted", 0);
tally(out, "direct_hnsw_insertions", 0);
tally(out, "flat_buffer_size", 0);
}
}
const NEWEST: u8 = 4;
const BAD_DIALECT: &str = "SEARCH_PARSE_ARGS DIALECT requires a non negative integer >=1 and <= 4";
const NEED_ARG: &str = "SEARCH_PARSE_ARGS Need an argument for ";
const ODD_PARAMS: &str = "SEARCH_ADD_ARGS Parameters must be specified in PARAM VALUE pairs";
const NOT_MAIN: &str = "` at position ";
const NOT_MAIN_END: &str = " for <main>";
const MOST: i64 = 1_000_000;
const LIMIT_TWO: &str = "SEARCH_PARSE_ARGS LIMIT requires two arguments";
const LIMIT_NUMBERS: &str = "SEARCH_PARSE_ARGS LIMIT needs two numeric arguments";
const LIMIT_OVER: &str = "SEARCH_LIMIT_OVER LIMIT exceeds maximum of 1000000";
const LIMIT_START: &str = "SEARCH_LIMIT_OVER The `offset` of the LIMIT must be 0 when `num` is 0";
const TIMEOUT_ARG: &str = "SEARCH_PARSE_ARGS Need argument for TIMEOUT";
const TIMEOUT_NUMBER: &str = "SEARCH_PARSE_ARGS TIMEOUT requires a non negative integer";
const NO_SCORER: &str = "SEARCH_QUERY_BAD No such scorer ";
const NO_LANGUAGE: &str = "SEARCH_QUERY_BAD No such language";
const LOW_RANGE: &str = "SEARCH_PARSE_ARGS Bad lower range: ";
const HIGH_RANGE: &str = "SEARCH_PARSE_ARGS Bad upper range: ";
const BACKWARDS: &str = "SEARCH_SYNTAX Invalid numeric range (min > max): @";
const FILTER_THREE: &str = "SEARCH_PARSE_ARGS FILTER requires 3 arguments";
const GEO_FIVE: &str = "SEARCH_PARSE_ARGS GEOFILTER requires 5 arguments";
const GEO_UNIT: &str = "SEARCH_PARSE_ARGS Unknown distance unit ";
const SYNTAX: &str = "SEARCH_SYNTAX ";
const NEED_NAME: &str = "SEARCH_PARSE_ARGS RETURN path AS name - must be accompanied with NAME";
const NEED_LOAD_NAME: &str = "SEARCH_PARSE_ARGS LOAD path AS name - must be accompanied with NAME";
const LOAD_COUNT: &str =
"SEARCH_PARSE_ARGS Bad arguments for LOAD: Expected number of fields or `*`";
const LOAD_BOUNDS: &str =
"SEARCH_PARSE_ARGS Bad arguments for LOAD: Value is outside acceptable bounds";
const LOAD_SHORT: &str =
"SEARCH_PARSE_ARGS Bad arguments for LOAD: Expected an argument, but none provided";
const NOT_HERE: &str = " is not supported on FT.AGGREGATE";
const GROUP_SHORT: &str =
"SEARCH_PARSE_ARGS Bad arguments for GROUPBY: Expected an argument, but none provided";
const GROUP_COUNT: &str =
"SEARCH_PARSE_ARGS Bad arguments for GROUPBY: Could not convert argument to expected type";
const NO_AT: &str = "SEARCH_PARSE_ARGS Bad arguments for GROUPBY: Unknown property `";
const NO_AT_MID: &str = "`. Did you mean `@";
const NO_AT_END: &str = "`?";
const SORT_SHORT: &str =
"SEARCH_PARSE_ARGS Bad arguments for SORTBY: Expected an argument, but none provided";
const SORT_COUNT: &str =
"SEARCH_PARSE_ARGS Bad arguments for SORTBY: Could not convert argument to expected type";
const SORT_BOUNDS: &str =
"SEARCH_PARSE_ARGS Bad arguments for SORTBY: Value is outside acceptable bounds";
const MAX_COUNT: &str =
"SEARCH_PARSE_ARGS Bad arguments for MAX: Could not convert argument to expected type";
const SORT_WAY: &str = "SEARCH_PARSE_ARGS MISSING ASC or DESC after sort field (";
const SORT_WAY_END: &str = ")";
const SORT_PROP: &str = "SEARCH_PROP_NOT_FOUND Property `";
const SORT_PROP_END: &str = "` not loaded nor in schema";
const SORT_TWICE: &str = "SEARCH_PARSE_ARGS Multiple SORTBY steps are not allowed. Sort multiple fields in a single step";
const SORT_ONCE: &str = "SEARCH_PARSE_ARGS Multiple SORTBY steps are not allowed";
const SORT_BARE: &str = "SEARCH_PARSE_ARGS Bad SORTBY arguments";
const SORT_MAX: &str = "SEARCH_PARSE_ARGS SORTBY MAX is not supported by FT.SEARCH";
const SCORE_ALONE: &str = "SEARCH_PARSE_ARGS EXPLAINSCORE must be accompanied with WITHSCORES";
const NO_PROPERTY: &str = "SEARCH_PROP_NOT_FOUND No such property `";
const NOT_LOADED: &str = "SEARCH_PROP_NOT_FOUND Property not loaded nor in pipeline: `";
const QUOTE_END: &str = "`";
const DUPLICATE_PROP: &str = "SEARCH_FIELD_DUP Property `";
const DUPLICATE_END: &str = "` specified more than once";
const LOAD_LATE: &str = "SEARCH_QUERY_BAD LOAD cannot be applied after projectors or reducers";
const STEP_SHORT: &str =
"SEARCH_PARSE_ARGS Bad arguments for APPLY/FILTER: Expected an argument, but none provided";
const AS_SHORT: &str = "SEARCH_PARSE_ARGS AS needs argument";
const REDUCE_BARE: &str = "SEARCH_PARSE_ARGS Bad arguments for REDUCE: SUCCESS";
const NO_REDUCER: &str = "SEARCH_REDUCER_NOT_FOUND No such reducer: ";
const MISSING_ARGS: &str = "SEARCH_PARSE_ARGS Missing arguments for ";
const COUNT_ONLY: &str = "SEARCH_ATTR_BAD Count accepts 0 values only";
const PERCENTAGE: &str = "SEARCH_PARSE_ARGS Percentage must be between 0.0 and 1.0";
const RESOLUTION: &str = "SEARCH_PARSE_ARGS Invalid resolution";
const SAMPLE_BIG: &str = "SEARCH_PARSE_ARGS Sample size too large";
const SAMPLE_SIZE: &str = "SEARCH_PARSE_ARGS Bad arguments for <sample size>";
const RESOLUTION_ARG: &str = "SEARCH_PARSE_ARGS Bad arguments for <resolution>";
const MOST_SAMPLE: i64 = 1000;
type Returns<'a> = Vec<(Box<[u8]>, &'a [u8])>;
struct Rows<'a> {
content: bool,
scores: bool,
payloads: bool,
offset: usize,
count: usize,
ret: Option<Returns<'a>>,
infields: Option<Vec<&'a [u8]>>,
inkeys: Option<Vec<&'a [u8]>>,
filters: Vec<(&'a [u8], f64, f64)>,
circles: Vec<Circle>,
scorer: Scorer,
payload: Option<&'a [u8]>,
explaining: bool,
slop: Option<i64>,
inorder: bool,
sort: Option<(&'a [u8], bool)>,
sorting: Option<Sorting>,
sortkeys: bool,
summarize: Option<(Vec<&'a [u8]>, Trim)>,
highlight: Option<(Vec<&'a [u8]>, Wrap)>,
upfront: Vec<&'a [u8]>,
texts: Vec<Box<[u8]>>,
stops: Option<Vec<Box<[u8]>>>,
wanted: Wanted,
distance: Vec<Yield>,
nearest: Option<Box<[u8]>>,
}
struct Sorting {
field: Field,
desc: bool,
slot: Option<usize>,
distance: bool,
}
impl Rows<'_> {
fn loading(&self) -> bool {
self.content && !self.ret.as_ref().is_some_and(Vec::is_empty)
}
}
impl Rows<'_> {
fn found(&self) -> Shows {
Shows {
fields: self.loading(),
scores: self.scores,
payloads: self.payloads,
sortkeys: self.sortkeys,
..Shows::default()
}
}
}
impl Asked<'_> {
fn rolls(&self) -> Shows {
Shows {
fields: self.rows.content,
scores: self.rows.scores,
payloads: self.rows.payloads,
sortkeys: self.pipe.sortkeys,
addscores: self.pipe.addscores,
}
}
}
impl Default for Rows<'_> {
fn default() -> Rows<'static> {
Rows {
content: true,
scores: false,
payloads: false,
offset: 0,
count: 10,
ret: None,
infields: None,
inkeys: None,
filters: Vec::new(),
circles: Vec::new(),
scorer: Scorer::default_scorer(),
payload: None,
explaining: false,
slop: None,
inorder: false,
sort: None,
sorting: None,
sortkeys: false,
summarize: None,
highlight: None,
upfront: Vec::new(),
texts: Vec::new(),
stops: None,
wanted: Wanted::default(),
distance: Vec::new(),
nearest: None,
}
}
}
#[derive(Clone, Copy, Default)]
struct Shows {
fields: bool,
scores: bool,
payloads: bool,
sortkeys: bool,
addscores: bool,
}
struct Asked<'a> {
dialect: u8,
params: Vec<Pair>,
verbatim: bool,
stopwords: bool,
rows: Rows<'a>,
pipe: Pipe<'a>,
cursor: Option<Asks>,
}
impl Default for Asked<'_> {
fn default() -> Asked<'static> {
Asked {
dialect: 1,
params: Vec::new(),
verbatim: false,
stopwords: true,
rows: Rows::default(),
pipe: Pipe::default(),
cursor: None,
}
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Mode {
Search,
Explain,
Aggregate,
}
#[derive(Clone, Copy)]
enum Bind<'a> {
Reading,
Binding(&'a [Yield]),
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Order {
Ranked,
Forwards,
Backwards,
}
fn buffered(asked: &Asked<'_>) -> bool {
asked.pipe.addscores && asked.rows.scorer.settles()
}
const IGNORED: &[(&[u8], usize)] = &[(b"WITHSORTKEYS", 0), (b"ADDSCORES", 0)];
fn options<'a>(
args: Args<'a>,
from: usize,
mode: Mode,
index: &Index,
bind: Bind<'_>,
) -> core::result::Result<Asked<'a>, Vec<u8>> {
let main = mode == Mode::Search;
let mut asked = Asked::default();
if let Bind::Binding(held) = bind {
asked.pipe.binding = true;
for want in held {
asked
.pipe
.base
.push((want.name.clone(), aggregate::Reads::Distance));
}
}
if mode == Mode::Aggregate {
asked.rows.count = usize::MAX;
}
let mut at = from;
while at < args.len() {
let word = args.get(at);
if mode != Mode::Search
&& let Some(next) = step(args, at, &mut asked, index)?
{
at = next;
continue;
}
if args::is(word, b"DIALECT") {
let Some(value) = args.opt(at + 1) else {
return Err(line(NEED_ARG, b"DIALECT", ""));
};
let Some(dialect) = parse_i64(value).filter(|d| (1..=i64::from(NEWEST)).contains(d))
else {
return Err(BAD_DIALECT.as_bytes().to_vec());
};
asked.dialect = u8::try_from(dialect).unwrap_or(1);
at += 2;
continue;
}
if args::is(word, b"PARAMS") {
at = params(args, at, &mut asked)?;
continue;
}
if let Some(next) = plan(args, at, &mut asked, mode)? {
at = next;
continue;
}
if !asked.pipe.stepped {
if args::is(word, b"VERBATIM") {
asked.verbatim = true;
at += 1;
continue;
}
if args::is(word, b"NOSTOPWORDS") {
asked.stopwords = false;
at += 1;
continue;
}
if mode == Mode::Aggregate
&& let Some(next) = extra(args, at, &mut asked)?
{
at = next;
continue;
}
if let Some(next) = row(args, at, &mut asked.rows, main, index)? {
at = next;
continue;
}
if mode == Mode::Explain
&& let Some((_, takes)) = IGNORED.iter().find(|(k, _)| args::is(word, k))
{
if args.opt(at + takes).is_none() {
return Err(line(BAD_ARGS, word, NOT_THERE));
}
at += takes + 1;
continue;
}
}
return Err(unknown(word, at - from + 1));
}
if asked.pipe.pending {
return Err(NEED_LOAD_NAME.as_bytes().to_vec());
}
if asked.rows.explaining && !asked.rows.scores {
return Err(SCORE_ALONE.as_bytes().to_vec());
}
aggregate::most(&mut asked);
Ok(asked)
}
fn settle<'a>(
asked: &mut Asked<'a>,
index: &Index,
node: &Node,
) -> core::result::Result<(), Vec<u8>> {
asked.rows.distance = query::yields(node);
let Some((field, desc)) = asked.rows.sort else {
return Ok(());
};
if asked.rows.distance.iter().any(|held| *held.name == *field) {
asked.rows.sorting = Some(Sorting {
field: Field::new(field, Kind::Numeric),
desc,
slot: None,
distance: true,
});
return Ok(());
}
let Some(known) = index.field(field) else {
return Err(line(SORT_PROP, field, SORT_PROP_END));
};
asked.rows.sorting = Some(Sorting {
field: known.clone(),
desc,
slot: index.slot(field),
distance: false,
});
Ok(())
}
fn plan<'a>(
args: Args<'a>,
at: usize,
asked: &mut Asked<'a>,
mode: Mode,
) -> core::result::Result<Option<usize>, Vec<u8>> {
let rows = &mut asked.rows;
let word = args.get(at);
if args::is(word, b"LIMIT") {
let (Some(offset), Some(count)) = (args.opt(at + 1), args.opt(at + 2)) else {
return Err(LIMIT_TWO.as_bytes().to_vec());
};
let (Some(offset), Some(count)) = (counted(offset), counted(count)) else {
return Err(LIMIT_NUMBERS.as_bytes().to_vec());
};
if count > MOST && mode == Mode::Search {
return Err(LIMIT_OVER.as_bytes().to_vec());
}
if count == 0 && offset != 0 {
return Err(LIMIT_START.as_bytes().to_vec());
}
rows.offset = usize::try_from(offset).unwrap_or(0);
rows.count = usize::try_from(count).unwrap_or(0);
if mode == Mode::Aggregate {
let (offset, count) = (rows.offset, rows.count);
windows(asked, offset, count);
}
return Ok(Some(at + 3));
}
if args::is(word, b"WITHCURSOR") {
let mut asks = asked.cursor.unwrap_or_default();
let mut next = at + 1;
while next < args.len() {
let word = args.get(next);
let count = args::is(word, b"COUNT");
if !count && !args::is(word, b"MAXIDLE") {
break;
}
let name: &[u8] = match count {
true => b"COUNT",
false => b"MAXIDLE",
};
let Some(value) = args.opt(next + 1) else {
return Err(line(BAD_ARGS, name, NOT_THERE));
};
let Some(number) = parse_i64(value) else {
return Err(line(BAD_ARGS, name, NOT_A_NUMBER));
};
if !(1..=cursor::MOST).contains(&number) {
return Err(line(BAD_ARGS, name, OUT_OF_RANGE));
}
match count {
true => asks.count = usize::try_from(number).unwrap_or(usize::MAX),
false => asks.idle = u64::try_from(number).unwrap_or(u64::MAX),
}
next += 2;
}
asked.cursor = Some(asks);
return Ok(Some(next));
}
if args::is(word, b"TIMEOUT") {
let Some(value) = args.opt(at + 1) else {
return Err(TIMEOUT_ARG.as_bytes().to_vec());
};
if counted(value).is_none() {
return Err(TIMEOUT_NUMBER.as_bytes().to_vec());
}
return Ok(Some(at + 2));
}
Ok(None)
}
fn step<'a>(
args: Args<'a>,
at: usize,
asked: &mut Asked<'a>,
index: &Index,
) -> core::result::Result<Option<usize>, Vec<u8>> {
if args::is(args.get(at), b"GROUPBY") {
return group(args, at, asked, index).map(Some);
}
if args::is(args.get(at), b"APPLY") {
return apply(args, at, asked, index).map(Some);
}
if args::is(args.get(at), b"FILTER") {
return keeps(args, at, asked, index).map(Some);
}
if args::is(args.get(at), b"SORTBY") {
return sorts(args, at, asked, index).map(Some);
}
if !args::is(args.get(at), b"LOAD") {
return Ok(None);
}
if asked.pipe.stage.is_some() {
return Err(LOAD_LATE.as_bytes().to_vec());
}
let Some(count) = args.opt(at + 1) else {
return Err(LOAD_SHORT.as_bytes().to_vec());
};
asked.pipe.stepped = true;
if count == b"*" {
asked.pipe.all = true;
asked.pipe.loader = true;
asked.pipe.loads = true;
return Ok(Some(at + 2));
}
let Some(count) = parse_i64(count) else {
return Err(LOAD_COUNT.as_bytes().to_vec());
};
let Ok(count) = usize::try_from(count) else {
return Err(LOAD_BOUNDS.as_bytes().to_vec());
};
asked.pipe.loader |= count > 0;
let mut at = at + 2;
let end = at + count;
while at < end {
let Some(path) = args.opt(at) else {
return Err(LOAD_SHORT.as_bytes().to_vec());
};
at += 1;
let field = path.strip_prefix(b"@").unwrap_or(path);
let name = match at < end && args::is(args.get(at), b"AS") {
false => field,
true => {
at += 1;
match args.opt(at).filter(|_| at < end) {
None => {
asked.pipe.pending = true;
break;
}
Some(name) => {
at += 1;
name
}
}
}
};
asked.pipe.load.push((field, name));
asked.pipe.loads |= !index.field(field).is_some_and(|held| held.sortable);
if !asked.pipe.base.iter().any(|(held, _)| **held == *name) {
asked
.pipe
.base
.push((name.into(), Reads::Field(field.into(), holds(index, field))));
}
}
Ok(Some(at))
}
fn extra<'a>(
args: Args<'a>,
at: usize,
asked: &mut Asked<'a>,
) -> core::result::Result<Option<usize>, Vec<u8>> {
let word = args.get(at);
if args::is(word, b"ADDSCORES") {
asked.pipe.addscores = true;
return Ok(Some(at + 1));
}
if args::is(word, b"WITHSORTKEYS") {
asked.pipe.sortkeys = true;
return Ok(Some(at + 1));
}
for name in [b"RETURN".as_slice(), b"SUMMARIZE", b"HIGHLIGHT"] {
if args::is(word, name) {
return Err(line("SEARCH_PARSE_ARGS ", name, NOT_HERE));
}
}
Ok(None)
}
fn row<'a>(
args: Args<'a>,
at: usize,
rows: &mut Rows<'a>,
main: bool,
index: &Index,
) -> core::result::Result<Option<usize>, Vec<u8>> {
let word = args.get(at);
if args::is(word, b"NOCONTENT") {
rows.content = false;
return Ok(Some(at + 1));
}
if args::is(word, b"WITHSCORES") {
rows.scores = true;
return Ok(Some(at + 1));
}
if args::is(word, b"WITHPAYLOADS") {
rows.payloads = true;
return Ok(Some(at + 1));
}
if args::is(word, b"EXPLAINSCORE") {
rows.explaining = true;
return Ok(Some(at + 1));
}
if args::is(word, b"INORDER") {
rows.inorder = true;
return Ok(Some(at + 1));
}
if args::is(word, b"SLOP") {
let Some(value) = args.opt(at + 1) else {
return Err(line(BAD_ARGS, word, NOT_THERE));
};
let Some(slop) = parse_i64(value) else {
return Err(line(BAD_ARGS, word, NOT_A_NUMBER));
};
rows.slop = Some(slop);
return Ok(Some(at + 2));
}
if args::is(word, b"SCORER") {
let Some(name) = args.opt(at + 1) else {
return Err(line(BAD_ARGS, word, NOT_THERE));
};
let Some(scorer) = Scorer::named(name) else {
return Err(line(NO_SCORER, name, ""));
};
rows.scorer = scorer;
return Ok(Some(at + 2));
}
if args::is(word, b"LANGUAGE") {
let Some(name) = args.opt(at + 1) else {
return Err(line(BAD_ARGS, word, NOT_THERE));
};
if !LANGUAGES.iter().any(|known| args::is(name, known)) {
return Err(NO_LANGUAGE.as_bytes().to_vec());
}
return Ok(Some(at + 2));
}
if args::is(word, b"EXPANDER") || args::is(word, b"PAYLOAD") {
let Some(value) = args.opt(at + 1) else {
return Err(line(BAD_ARGS, word, NOT_THERE));
};
if args::is(word, b"PAYLOAD") {
rows.payload = Some(value);
}
return Ok(Some(at + 2));
}
if args::is(word, b"RETURN") {
return returned(args, at, rows, index).map(Some);
}
if args::is(word, b"SUMMARIZE") || args::is(word, b"HIGHLIGHT") {
return marking(args, at, rows, index).map(Some);
}
if main && args::is(word, b"SORTBY") {
return sorting(args, at, rows).map(Some);
}
if main && args::is(word, b"WITHSORTKEYS") {
rows.sortkeys = true;
return Ok(Some(at + 1));
}
if args::is(word, b"INFIELDS") {
let (names, next) = names(args, at)?;
rows.infields = Some(names);
return Ok(Some(next));
}
if args::is(word, b"INKEYS") {
let (names, next) = names(args, at)?;
rows.inkeys = Some(names);
return Ok(Some(next));
}
if main && args::is(word, b"FILTER") {
return filter(args, at, rows, index).map(Some);
}
if main && args::is(word, b"GEOFILTER") {
return geofilter(args, at, rows).map(Some);
}
Ok(None)
}
fn sorting<'a>(
args: Args<'a>,
at: usize,
rows: &mut Rows<'a>,
) -> core::result::Result<usize, Vec<u8>> {
if rows.sort.is_some() {
return Err(SORT_ONCE.as_bytes().to_vec());
}
let Some(field) = args.opt(at + 1) else {
return Err(SORT_BARE.as_bytes().to_vec());
};
let mut next = at + 2;
let mut desc = false;
if let Some(word) = args.opt(next)
&& (args::is(word, b"ASC") || args::is(word, b"DESC"))
{
desc = args::is(word, b"DESC");
next += 1;
}
if args.opt(next).is_some_and(|word| args::is(word, b"MAX")) {
return Err(SORT_MAX.as_bytes().to_vec());
}
rows.sort = Some((field, desc));
Ok(next)
}
fn counted(value: &[u8]) -> Option<i64> {
parse_i64(value).filter(|n| *n >= 0)
}
fn returned<'a>(
args: Args<'a>,
at: usize,
rows: &mut Rows<'a>,
index: &Index,
) -> core::result::Result<usize, Vec<u8>> {
let Some(count) = args.opt(at + 1) else {
return Err(line(BAD_ARGS, b"RETURN", NOT_THERE));
};
let Some(count) = parse_i64(count).filter(|n| *n >= 0) else {
return Err(line(BAD_ARGS, b"RETURN", NOT_A_NUMBER));
};
let count = usize::try_from(count).unwrap_or(0);
let mut want = Vec::new();
let mut step = 0;
while step < count {
let Some(field) = args.opt(at + 2 + step) else {
return Err(line(BAD_ARGS, b"RETURN", NOT_THERE));
};
step += 1;
if step < count && args.opt(at + 2 + step).is_some_and(|w| args::is(w, b"AS")) {
if step + 1 >= count {
return Err(NEED_NAME.as_bytes().to_vec());
}
let Some(name) = args.opt(at + 3 + step) else {
return Err(line(BAD_ARGS, b"RETURN", NOT_THERE));
};
want.push((reads(index, field), name));
step += 2;
continue;
}
want.push((reads(index, field), field));
}
rows.ret = Some(want);
Ok(at + 2 + count)
}
fn marking<'a>(
args: Args<'a>,
at: usize,
rows: &mut Rows<'a>,
index: &Index,
) -> core::result::Result<usize, Vec<u8>> {
let word = args.get(at);
let summary = args::is(word, b"SUMMARIZE");
let bad = || line(BAD_ARGS, word, "");
let mut at = at + 1;
let mut fields: Vec<&'a [u8]> = Vec::new();
let mut trim = Trim::default();
let mut wrap = Wrap::default();
while let Some(next) = args.opt(at) {
if args::is(next, b"FIELDS") {
let Some(count) = args.opt(at + 1).and_then(parse_i64).filter(|n| *n >= 0) else {
return Err(bad());
};
let count = usize::try_from(count).unwrap_or(0);
for step in 0..count {
let Some(name) = args.opt(at + 2 + step) else {
return Err(bad());
};
if !index.schema.iter().any(|f| *f.attribute == *name) {
return Err(line(NO_PROPERTY, name, QUOTE_END));
}
fields.push(name);
}
at += 2 + count;
continue;
}
if summary && (args::is(next, b"FRAGS") || args::is(next, b"LEN")) {
let Some(value) = args.opt(at + 1).and_then(parse_i64).filter(|n| *n >= 0) else {
return Err(bad());
};
let value = usize::try_from(value).unwrap_or(0);
match args::is(next, b"FRAGS") {
true => trim.frags = value,
false => trim.len = value,
}
at += 2;
continue;
}
if summary && args::is(next, b"SEPARATOR") {
let Some(value) = args.opt(at + 1) else {
return Err(bad());
};
trim.separator = value.into();
at += 2;
continue;
}
if !summary && args::is(next, b"TAGS") {
let (Some(open), Some(close)) = (args.opt(at + 1), args.opt(at + 2)) else {
return Err(bad());
};
wrap.open = open.into();
wrap.close = close.into();
at += 3;
continue;
}
break;
}
for name in &fields {
if !rows.upfront.contains(name) {
rows.upfront.push(name);
}
}
if rows.texts.is_empty() {
rows.texts = index
.schema
.iter()
.filter(|field| matches!(field.kind, Kind::Text(_)))
.map(|field| field.attribute.clone())
.collect();
rows.stops = index.definition.stopwords.clone();
}
match summary {
true => rows.summarize = Some((fields, trim)),
false => rows.highlight = Some((fields, wrap)),
}
Ok(at)
}
fn reads(index: &Index, name: &[u8]) -> Box<[u8]> {
index
.field(name)
.map_or_else(|| name.into(), |field| field.identifier.clone())
}
fn names<'a>(args: Args<'a>, at: usize) -> core::result::Result<(Vec<&'a [u8]>, usize), Vec<u8>> {
let word = args.get(at);
let Some(count) = args.opt(at + 1) else {
return Err(line(BAD_ARGS, word, NOT_THERE));
};
let Some(count) = parse_i64(count).filter(|n| *n >= 0) else {
return Err(line(BAD_ARGS, word, NOT_A_NUMBER));
};
let count = usize::try_from(count).unwrap_or(0);
let mut out = Vec::with_capacity(count);
for step in 0..count {
let Some(name) = args.opt(at + 2 + step) else {
return Err(line(BAD_ARGS, word, NOT_THERE));
};
out.push(name);
}
Ok((out, at + 2 + count))
}
fn filter<'a>(
args: Args<'a>,
at: usize,
rows: &mut Rows<'a>,
index: &Index,
) -> core::result::Result<usize, Vec<u8>> {
let (Some(field), Some(min), Some(max)) =
(args.opt(at + 1), args.opt(at + 2), args.opt(at + 3))
else {
return Err(FILTER_THREE.as_bytes().to_vec());
};
let Some(low) = ends(min) else {
return Err(line(LOW_RANGE, min, ""));
};
let Some(high) = ends(max) else {
return Err(line(HIGH_RANGE, max, ""));
};
if low > high && numeric(index, field) {
return Err(backwards(field, low, high));
}
rows.filters.push((field, low, high));
Ok(at + 4)
}
fn geofilter<'a>(
args: Args<'a>,
at: usize,
rows: &mut Rows<'a>,
) -> core::result::Result<usize, Vec<u8>> {
let (Some(field), Some(lon), Some(lat), Some(radius), Some(unit)) = (
args.opt(at + 1),
args.opt(at + 2),
args.opt(at + 3),
args.opt(at + 4),
args.opt(at + 5),
) else {
return Err(GEO_FIVE.as_bytes().to_vec());
};
let mut ends = [0.0; 3];
for (slot, (raw, what)) in ends.iter_mut().zip([
(lon, &b"<lon>"[..]),
(lat, &b"<lat>"[..]),
(radius, &b"<radius>"[..]),
]) {
let Some(value) = coord(raw) else {
return Err(line(BAD_ARGS, what, NOT_A_NUMBER));
};
*slot = value;
}
if geo::Unit::parse(unit).is_none() {
return Err(line(GEO_UNIT, unit, ""));
}
if !ends[0].is_nan() && !ends[1].is_nan() && !geo::in_range(ends[0], ends[1]) {
return Err(line(SYNTAX, BAD_POINT.as_bytes(), ""));
}
if ends[2] <= 0.0 {
return Err(line(SYNTAX, BAD_RADIUS.as_bytes(), ""));
}
rows.circles.push(Circle {
field: field.into(),
lon: ends[0],
lat: ends[1],
radius: ends[2],
unit: unit.into(),
});
Ok(at + 6)
}
fn coord(raw: &[u8]) -> Option<f64> {
let text = core::str::from_utf8(raw).ok()?;
let value: f64 = text.trim_ascii_start().parse().ok()?;
(!value.is_infinite()).then_some(value)
}
fn numeric(index: &Index, field: &[u8]) -> bool {
index
.schema
.iter()
.any(|f| *f.attribute == *field && matches!(f.kind, Kind::Numeric))
}
fn ends(value: &[u8]) -> Option<f64> {
match value {
b"+inf" | b"inf" | b"INF" | b"+INF" => Some(f64::INFINITY),
b"-inf" | b"-INF" => Some(f64::NEG_INFINITY),
_ => parse_f64(value),
}
}
fn backwards(field: &[u8], low: f64, high: f64) -> Vec<u8> {
let mut out = BACKWARDS.as_bytes().to_vec();
out.extend_from_slice(field);
out.extend_from_slice(b":[");
out.extend_from_slice(format!("{low:.6} {high:.6}").as_bytes());
out.push(b']');
out
}
fn params(
args: Args<'_>,
at: usize,
asked: &mut Asked<'_>,
) -> core::result::Result<usize, Vec<u8>> {
let Some(count) = args.opt(at + 1) else {
return Err(line(BAD_ARGS, b"PARAMS", NOT_THERE));
};
let Some(count) = parse_i64(count) else {
return Err(line(BAD_ARGS, b"PARAMS", NOT_A_NUMBER));
};
let Ok(count) = usize::try_from(count) else {
return Err(line(BAD_ARGS, b"PARAMS", OUT_OF_RANGE));
};
if at + 2 + count > args.len() {
return Err(line(BAD_ARGS, b"PARAMS", NOT_THERE));
}
if count == 0 || count % 2 != 0 {
return Err(ODD_PARAMS.as_bytes().to_vec());
}
for step in 0..count / 2 {
let name = args.opt(at + 2 + step * 2);
let value = args.opt(at + 3 + step * 2);
let (Some(name), Some(value)) = (name, value) else {
return Err(line(BAD_ARGS, b"PARAMS", NOT_THERE));
};
asked.params.push((name.into(), value.into()));
}
Ok(at + 2 + count)
}
fn line(head: &str, word: &[u8], tail: &str) -> Vec<u8> {
let mut out = head.as_bytes().to_vec();
out.extend_from_slice(word);
out.extend_from_slice(tail.as_bytes());
out
}
fn unknown(word: &[u8], position: usize) -> Vec<u8> {
let mut out = UNKNOWN.as_bytes().to_vec();
out.extend_from_slice(word);
out.extend_from_slice(NOT_MAIN.as_bytes());
out.extend_from_slice(position.to_string().as_bytes());
out.extend_from_slice(NOT_MAIN_END.as_bytes());
out
}
fn refused(bad: &Bad) -> Vec<u8> {
match bad {
Bad::Syntax { at, near } => spot("SEARCH_SYNTAX Syntax error at offset ", *at, near),
Bad::Unknown { at, near } => named(
"SEARCH_SYNTAX Unknown field at offset ",
*at,
near.as_deref(),
),
Bad::Wrong { kind, at, near } => {
let head = format!("SEARCH_SYNTAX Expected a {kind} field at offset ");
named(&head, *at, near.as_deref())
}
Bad::Attribute(name) => line("SEARCH_OPTION_INVALID Invalid attribute ", name, ""),
Bad::Blob { got, want } => format!(
"SEARCH_QUERY_BAD Error parsing vector similarity query: query vector blob size ({got}) does not match index's expected size ({want})."
)
.into_bytes(),
Bad::Large => b"SEARCH_QUERY_BAD Error parsing vector similarity query: \
query KNN K parameter is too large, must not exceed 288230376151711744"
.to_vec(),
Bad::Count { name, value } => {
let mut out = b"SEARCH_NUMERIC_VALUE_INVALID Invalid numeric value (".to_vec();
out.extend_from_slice(value);
out.extend_from_slice(b") for parameter `");
out.extend_from_slice(name);
out.push(b'`');
out
}
Bad::Radius(radius) => {
let mut out =
b"SEARCH_QUERY_BAD Error parsing vector similarity query: negative radius ("
.to_vec();
out.extend_from_slice(radius);
out.extend_from_slice(b") given in a range query");
out
}
Bad::Value { name, value } => {
let mut out = b"SEARCH_SYNTAX Invalid value (".to_vec();
out.extend_from_slice(value);
out.extend_from_slice(b") for `");
out.extend_from_slice(name);
out.push(b'`');
out
}
Bad::Missing(name) => {
let mut out = b"SEARCH_PARAM_NOT_FOUND Parameter not found `".to_vec();
out.extend_from_slice(name);
out.push(b'`');
out
}
Bad::Taken(name) => {
let mut out = b"SEARCH_INDEX_EXISTS Property `".to_vec();
out.extend_from_slice(name);
out.extend_from_slice(b"` already exists in schema");
out
}
Bad::Option(why) => {
let mut out = why.code().as_bytes().to_vec();
out.push(b' ');
out.extend_from_slice(why.words().as_bytes());
out.extend_from_slice(b" (Error parsing vector similarity parameters)");
out
}
Bad::Twice(first, second) => {
let mut out =
b"SEARCH_FIELD_DUP Distance field was specified twice for vector query: ".to_vec();
out.extend_from_slice(first);
out.extend_from_slice(b" and ");
out.extend_from_slice(second);
out
}
Bad::Plain(text) => line("SEARCH_SYNTAX ", text.as_bytes(), ""),
Bad::Refused(text) => line("SEARCH_QUERY_BAD ", text.as_bytes(), ""),
}
}
fn named(head: &str, at: usize, near: Option<&[u8]>) -> Vec<u8> {
let Some(near) = near else {
let mut out = head.as_bytes().to_vec();
out.extend_from_slice(at.to_string().as_bytes());
return out;
};
spot(head, at, near)
}
fn spot(head: &str, at: usize, near: &[u8]) -> Vec<u8> {
let mut out = head.as_bytes().to_vec();
out.extend_from_slice(at.to_string().as_bytes());
out.extend_from_slice(b" near ");
out.extend_from_slice(near);
out
}
fn explain<'a>(reg: &mut Registry, args: Args<'a>, out: &mut Out, cli: bool) -> Answer<'a> {
let name = args.get(1);
let query = args.get(2);
let Some(index) = reg.open(name) else {
return Err(Fail::naming(MISSING, name));
};
let asked = match options(args, 3, Mode::Explain, index, Bind::Binding(&[])) {
Ok(asked) => asked,
Err(text) => {
out.error(&text);
return Ok(());
}
};
let ask = Ask {
dialect: asked.dialect,
params: &asked.params,
verbatim: asked.verbatim,
stopwords: asked.stopwords,
};
let node = match query::parse(query, index, &ask) {
Ok(node) => node,
Err(bad) => {
out.error(&refused(&bad));
return Ok(());
}
};
let printed = query::explain(&node, index);
if !cli {
out.bulk(&printed);
return Ok(());
}
let lines = query::explain::lines(&printed);
out.array(lines.len());
for line in lines {
out.simple(line);
}
Ok(())
}
type Built<'a> = (&'a Row, Option<Vec<(&'a [u8], &'a [u8])>>);
type Rolled<'a> = (&'a Row, Vec<(&'a [u8], &'a [u8])>);
type Named = (Box<[u8]>, Box<[u8]>);
type Pairs = Vec<Named>;
type Away = (Box<[u8]>, f64);
type Scored<'a> = (walk::Hit<'a>, f64, Vec<Away>);
#[derive(Clone)]
struct Row {
key: Box<[u8]>,
score: f64,
payload: Option<Box<[u8]>>,
note: Option<Note>,
sort: Option<Sorted>,
dists: Vec<Away>,
}
impl Row {
fn away(&self, name: &[u8]) -> Option<f64> {
self.dists
.iter()
.find(|(held, _)| **held == *name)
.map(|(_, away)| *away)
}
}
pub(super) fn find(server: &Server, db: usize, args: Args<'_>, out: &mut Out) -> Result<()> {
searched(server, db, args, 2, None, out)
}
pub(super) fn profiled(server: &Server, db: usize, args: Args<'_>, out: &mut Out) -> Result<()> {
profile::run(server, db, args, out)
}
pub(super) fn searched(
server: &Server,
db: usize,
args: Args<'_>,
at: usize,
watch: Option<&mut Watch>,
out: &mut Out,
) -> Result<()> {
let mut watch = watch;
let name = args.get(1);
let query = args.get(at);
let clock = Instant::now();
let mut asked;
let mut canon: Box<[u8]> = Box::default();
let (total, rows) = {
let mut reg = server.search.lock();
let Some(index) = reg.open(name) else {
Fail::naming(MISSING, name).write(out);
return Ok(());
};
asked = match options(args, at + 1, Mode::Search, index, Bind::Binding(&[])) {
Ok(asked) => asked,
Err(text) => {
out.error(&text);
return Ok(());
}
};
let ask = Ask {
dialect: asked.dialect,
params: &asked.params,
verbatim: asked.verbatim,
stopwords: asked.stopwords,
};
let node = match query::parse(query, index, &ask) {
Ok(node) => node,
Err(bad) => {
out.error(&refused(&bad));
return Ok(());
}
};
if let Err(text) = settle(&mut asked, index, &node) {
out.error(&text);
return Ok(());
}
if asked.cursor.is_some() {
canon = index.name.clone();
if let Err(fail) = cursor::room(server, &canon, 1) {
fail.write(out);
return Ok(());
}
}
if asked.rows.summarize.is_some() || asked.rows.highlight.is_some() {
asked.rows.wanted = Wanted::of(&node);
}
if let Some(watch) = watch.as_deref_mut() {
watch.parsing = clock.elapsed();
}
let shaped = shape(node, index, &asked.rows);
if let Some(watch) = watch.as_deref_mut() {
watch.creating = clock.elapsed() - watch.parsing;
}
let (total, rows, ran) = gather(
index,
shaped,
&asked.rows,
Order::Ranked,
false,
watch.is_some(),
);
if let Some(watch) = watch.as_deref_mut()
&& let Some((ran, spent)) = ran
{
watch.ran = Some(ran);
watch.walking = spent;
}
(total, rows)
};
if let Some(watch) = watch {
watch.steps = searching(&asked.rows, total, rows.len());
}
write(server, db, total, &rows, &asked, &canon, out);
Ok(())
}
fn searching(rows: &Rows<'_>, total: usize, window: usize) -> Vec<(Vec<u8>, usize)> {
let mut out = vec![(b"Index".to_vec(), total)];
if !rows.distance.is_empty() {
out.push((b"Metrics Applier".to_vec(), total));
}
if rows.count == 0 {
out.push((b"Counter".to_vec(), 1));
return out;
}
if rows.sorting.is_none() || rows.scores {
out.push((b"Scorer".to_vec(), total));
}
out.push((b"Sorter".to_vec(), window));
if rows.content {
out.push((b"Loader".to_vec(), window));
if rows.summarize.is_some() || rows.highlight.is_some() {
out.push((b"Highlighter".to_vec(), window));
}
}
out
}
pub(super) fn roll(server: &Server, db: usize, args: Args<'_>, out: &mut Out) -> Result<()> {
aggregated(server, db, args, 2, None, out)
}
pub(super) fn aggregated(
server: &Server,
db: usize,
args: Args<'_>,
at: usize,
watch: Option<&mut Watch>,
out: &mut Out,
) -> Result<()> {
let mut watch = watch;
let name = args.get(1);
let query = args.get(at);
let clock = Instant::now();
let asked;
let mut canon: Box<[u8]> = Box::default();
let (total, rows) = {
let mut reg = server.search.lock();
let Some(index) = reg.open(name) else {
Fail::naming(MISSING, name).write(out);
return Ok(());
};
let read = match options(args, at + 1, Mode::Aggregate, index, Bind::Reading) {
Ok(read) => read,
Err(text) => {
out.error(&text);
return Ok(());
}
};
let ask = Ask {
dialect: read.dialect,
params: &read.params,
verbatim: read.verbatim,
stopwords: read.stopwords,
};
let node = match query::parse(query, index, &ask) {
Ok(node) => node,
Err(bad) => {
out.error(&refused(&bad));
return Ok(());
}
};
let held = query::yields(&node);
asked = match options(args, at + 1, Mode::Aggregate, index, Bind::Binding(&held)) {
Ok(mut asked) => {
asked.rows.nearest = held
.iter()
.find(|want| want.ordered)
.map(|want| want.name.clone());
asked.rows.distance = held;
asked
}
Err(text) => {
out.error(&text);
return Ok(());
}
};
if asked.cursor.is_some() {
canon = index.name.clone();
if let Err(fail) = cursor::room(server, &canon, 1) {
fail.write(out);
return Ok(());
}
}
let order = match buffered(&asked) {
true => Order::Backwards,
false => Order::Forwards,
};
if let Some(watch) = watch.as_deref_mut() {
watch.parsing = clock.elapsed();
}
let shaped = shape(node, index, &asked.rows);
if let Some(watch) = watch.as_deref_mut() {
watch.creating = clock.elapsed() - watch.parsing;
}
let (total, rows, ran) = gather(
index,
shaped,
&asked.rows,
order,
!asked.pipe.steps.is_empty(),
watch.is_some(),
);
if let Some(watch) = watch.as_deref_mut() {
if let Some((ran, spent)) = ran {
watch.ran = Some(ran);
watch.walking = spent;
}
watch.steps.push((b"Index".to_vec(), total));
if !asked.rows.distance.is_empty() {
watch.steps.push((b"Metrics Applier".to_vec(), total));
}
}
(total, rows)
};
rolled(server, db, total, &rows, &asked, &canon, watch, out);
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn rolled(
server: &Server,
db: usize,
total: usize,
rows: &[Row],
asked: &Asked<'_>,
index: &[u8],
watch: Option<&mut Watch>,
out: &mut Out,
) {
if !asked.pipe.steps.is_empty() {
piped(server, db, total, rows, asked, index, watch, out);
return;
}
let pipe = &asked.pipe;
let want = &asked.rows;
let held: Vec<Option<indexing::Document>> = match pipe.loader {
true => rows
.iter()
.map(|row| indexing::read(&server.dbs[db], &row.key))
.collect(),
false => Vec::new(),
};
let mut built: Vec<Rolled<'_>> = Vec::with_capacity(rows.len());
let mut lost = 0;
for (at, row) in rows.iter().enumerate() {
if pipe.loader {
let Some(doc) = held.get(at).and_then(Option::as_ref) else {
lost += 1;
continue;
};
built.push((row, props(doc, pipe)));
} else {
built.push((row, Vec::new()));
}
}
let total = total - lost;
if let Some(watch) = watch {
if pipe.addscores {
watch.steps.push((b"Scorer".to_vec(), built.len()));
}
if pipe.loads {
watch.steps.push((b"Loader".to_vec(), built.len()));
}
}
if let Some(asks) = asked.cursor {
let made = Made::Rolled(
built
.into_iter()
.map(|(row, props)| (row.clone(), owned(props)))
.collect(),
);
let kept = Kept {
made,
walk: vec![true; total],
shows: asked.rolls(),
total,
whole: want.count == 0 || buffered(asked),
loader: pipe.loader,
offset: want.offset,
window: want.count,
};
cursor::open(server, index, kept, asks, out);
return;
}
let wide = out.proto().is_resp3();
let whole = want.count == 0 || buffered(asked);
let count = match whole {
true => total,
false => {
let reached = match wide {
true => built.len(),
false => 1,
};
want.offset.saturating_add(reached).min(total)
}
};
rolls(count, &built, asked.rolls(), out);
}
fn rolls(count: usize, built: &[Rolled<'_>], shows: Shows, out: &mut Out) {
if out.proto().is_resp3() {
rolled_deep(count, built, shows, out);
return;
}
let per = usize::from(shows.scores)
+ usize::from(shows.payloads)
+ usize::from(shows.sortkeys)
+ usize::from(shows.fields);
out.array(1 + built.len() * per);
out.int(count as i64);
for (row, fields) in built {
if shows.scores {
out.double(row.score);
}
if shows.payloads {
match &row.payload {
Some(payload) => out.bulk(payload),
None => out.nil(),
}
}
if shows.sortkeys {
out.nil();
}
if shows.fields {
out.map(fields.len() + usize::from(shows.addscores) + row.dists.len());
spaced(row, out);
if shows.addscores {
out.bulk(b"__score");
out.bulk(twelve(row.score).as_bytes());
}
for (field, value) in fields {
out.bulk(field);
out.bulk(value);
}
}
}
}
fn spaced(row: &Row, out: &mut Out) {
for (name, away) in &row.dists {
out.bulk(name);
out.bulk(twelve(*away).as_bytes());
}
}
fn rolled_deep(count: usize, built: &[Rolled<'_>], shows: Shows, out: &mut Out) {
out.map(5);
out.simple(b"attributes");
out.array(0);
out.simple(b"format");
out.simple(b"STRING");
out.simple(b"results");
out.array(built.len());
for (row, fields) in built {
out.map(
1 + usize::from(shows.scores)
+ usize::from(shows.payloads)
+ usize::from(shows.sortkeys)
+ usize::from(shows.fields),
);
if shows.scores {
out.simple(b"score");
out.double(row.score);
}
if shows.payloads {
out.simple(b"payload");
match &row.payload {
Some(payload) => out.bulk(payload),
None => out.nil(),
}
}
if shows.sortkeys {
out.simple(b"sortkey");
out.nil();
}
if shows.fields {
out.simple(b"extra_attributes");
out.map(fields.len() + usize::from(shows.addscores) + row.dists.len());
spaced(row, out);
if shows.addscores {
out.bulk(b"__score");
out.bulk(twelve(row.score).as_bytes());
}
for (field, value) in fields {
out.bulk(field);
out.bulk(value);
}
}
out.simple(b"values");
out.array(0);
}
out.simple(b"total_results");
out.int(count as i64);
out.simple(b"warning");
out.array(0);
}
fn props<'d, 'w: 'd>(doc: &'d indexing::Document, pipe: &Pipe<'w>) -> Vec<(&'d [u8], &'d [u8])> {
let pairs = doc.pairs();
let mut out: Vec<(&[u8], &[u8])> = Vec::with_capacity(pipe.load.len());
for (field, name) in &pipe.load {
if out.iter().any(|(held, _)| held == name) {
continue;
}
let Some((_, value)) = pairs.iter().find(|(held, _)| held == field) else {
continue;
};
out.push((*name, *value));
}
if pipe.all {
for (field, value) in pairs {
if !out.iter().any(|(held, _)| *held == field) {
out.push((field, value));
}
}
}
out
}
use yo_search::expr::{seventeen, twelve};
fn holds(index: &Index, field: &[u8]) -> Shape {
let numeric = index
.field(field)
.is_some_and(|held| held.kind == Kind::Numeric);
match numeric {
true => Shape::Number,
false => Shape::Words,
}
}
fn shape(node: Node, index: &Index, rows: &Rows<'_>) -> Node {
let mut node = node;
if let Some(want) = &rows.infields {
if !want.is_empty() {
let mask = want
.iter()
.filter_map(|field| query::explain::bit(index, field))
.fold(0 as Mask, |mask, bit| mask | bit);
node.narrow(mask);
}
}
node.slop = rows.slop;
node.inorder = rows.inorder || node.inorder;
if rows.filters.is_empty() && rows.circles.is_empty() {
return node;
}
let mut under = match node.what {
What::Wildcard => Vec::new(),
_ => vec![node],
};
for (field, min, max) in &rows.filters {
under.push(Node::new(What::Numeric(Range {
field: (*field).into(),
min: *min,
max: *max,
min_open: false,
max_open: false,
})));
}
for circle in &rows.circles {
under.push(Node::new(What::Geo(circle.clone())));
}
Node::new(What::Intersect(under))
}
fn gather(
index: &Index,
node: Node,
rows: &Rows<'_>,
order: Order,
whole: bool,
profiling: bool,
) -> (usize, Vec<Row>, Option<(walk::Ran, Duration)>) {
let facts = index.held.facts();
let measure = rows.scorer.divides();
let mut ran = None;
let mut spent = Duration::ZERO;
let walked = if profiling {
let clock = Instant::now();
let (walked, tree) = walk::profiled(&index.held, &node, measure);
spent = clock.elapsed();
ran = Some(tree);
walked
} else if measure {
walk::spaced(&index.held, &node)
} else {
walk::run(&index.held, &node)
};
let asking: Vec<(&Yield, &yo_search::vecs::Vecs)> = rows
.distance
.iter()
.filter_map(|want| Some((want, index.held.vecs(&want.field)?)))
.collect();
let mut found: Vec<Scored<'_>> = walked
.into_iter()
.filter_map(|hit| {
let doc = index.held.docs.get(hit.id)?;
if let Some(keys) = &rows.inkeys
&& !keys.contains(&&*doc.key)
{
return None;
}
let score = rows
.scorer
.of(&facts, doc, &hit.found, rows.payload, hit.slop);
let dists = asking
.iter()
.filter_map(|(want, vecs)| {
let away = vecs.distance(hit.id, &want.asked)?;
Some((want.name.clone(), f64::from(away)))
})
.collect();
Some((hit, score, dists))
})
.collect();
let mut scores: Vec<f64> = found.iter().map(|(_, score, _)| *score).collect();
let best = scores.iter().copied().fold(0.0_f64, f64::max);
rows.scorer.settle(&mut scores);
for (row, score) in found.iter_mut().zip(&scores) {
row.1 = *score;
}
let by = rows
.sorting
.as_ref()
.filter(|by| by.slot.is_some() || by.distance);
let later = rows.sorting.is_some() && by.is_none();
match (by, order) {
(Some(by), _) if by.distance => {
found.sort_by(|a, b| {
let ids = match by.desc {
true => b.0.id.cmp(&a.0.id),
false => a.0.id.cmp(&b.0.id),
};
let held = |row: &Scored<'_>| {
row.2
.iter()
.find(|(name, _)| **name == *by.field.attribute)
.map(|(_, away)| Sorted::Number(*away))
};
sorted::order(held(a).as_ref(), held(b).as_ref(), by.desc).then(ids)
});
}
(Some(by), _) => {
let slot = by.slot.unwrap_or_default();
let held = |id: u32| index.held.docs.get(id).and_then(|doc| doc.sorted(slot));
found.sort_by(|a, b| {
let ids = match by.desc {
true => b.0.id.cmp(&a.0.id),
false => a.0.id.cmp(&b.0.id),
};
sorted::order(held(a.0.id), held(b.0.id), by.desc).then(ids)
});
}
(None, _) if later => found.sort_by_key(|(hit, _, _)| hit.id),
(None, Order::Ranked) => found.sort_by(|a, b| {
b.1.partial_cmp(&a.1)
.unwrap_or(core::cmp::Ordering::Equal)
.then(a.0.id.cmp(&b.0.id))
}),
(None, Order::Forwards) => match &rows.nearest {
Some(name) => {
let away = |row: &Scored<'_>| {
row.2
.iter()
.find(|(held, _)| **held == **name)
.map_or(f64::INFINITY, |(_, away)| *away)
};
found.sort_by(|a, b| {
away(a)
.partial_cmp(&away(b))
.unwrap_or(core::cmp::Ordering::Equal)
.then(a.0.id.cmp(&b.0.id))
});
}
None => found.sort_by_key(|(hit, _, _)| hit.id),
},
(None, Order::Backwards) => {
found.sort_by_key(|(hit, _, _)| core::cmp::Reverse(hit.id));
}
}
let total = found.len();
let (offset, count) = match whole || later {
true => (0, usize::MAX),
false => (rows.offset, rows.count),
};
let why = Why::new(rows.scorer, facts).about(rows.payload).over(best);
let window = found
.into_iter()
.skip(offset)
.take(count)
.filter_map(|(hit, score, dists)| {
let doc = index.held.docs.get(hit.id)?;
let sort = match by {
Some(by) if by.distance => dists
.iter()
.find(|(name, _)| **name == *by.field.attribute)
.map(|(_, away)| Sorted::Number(*away)),
Some(by) => doc.sorted(by.slot.unwrap_or_default()).cloned(),
None => None,
};
Some(Row {
key: doc.key.clone(),
score,
payload: doc.payload.clone(),
sort,
note: rows.explaining.then(|| why.note(doc, &hit.found, hit.slop)),
dists,
})
})
.collect();
(total, window, ran.map(|ran| (ran, spent)))
}
type Held<'a> = (&'a Row, Option<Vec<(&'a [u8], &'a [u8])>>);
fn write(
server: &Server,
db: usize,
total: usize,
rows: &[Row],
asked: &Asked<'_>,
index: &[u8],
out: &mut Out,
) {
let want = &asked.rows;
let loading = want.loading();
let slow = want
.sorting
.as_ref()
.is_some_and(|by| by.slot.is_none() && !by.distance);
let mut held: Vec<Option<indexing::Document>> = if loading || slow {
rows.iter()
.map(|row| indexing::read(&server.dbs[db], &row.key))
.collect()
} else {
Vec::new()
};
let mut lost = 0;
let carried: Vec<Row>;
let mut rows = rows;
if let Some(by) = want
.sorting
.as_ref()
.filter(|by| by.slot.is_none() && !by.distance)
{
let mut pairs: Vec<(Row, Option<indexing::Document>)> = rows
.iter()
.cloned()
.zip(core::mem::take(&mut held))
.filter(|(_, doc)| doc.is_some())
.collect();
lost = rows.len() - pairs.len();
for (row, doc) in &mut pairs {
row.sort = doc
.as_ref()
.and_then(|doc| doc.held(&by.field.identifier))
.and_then(|raw| Sorted::read(&by.field, raw));
}
if by.desc {
pairs.reverse();
}
pairs.sort_by(|a, b| sorted::order(a.0.sort.as_ref(), b.0.sort.as_ref(), by.desc));
(carried, held) = pairs.into_iter().skip(want.offset).take(want.count).unzip();
rows = &carried;
}
let keys: Vec<Option<Vec<u8>>> = match want.sorting.is_some() {
true => rows.iter().map(|row| shown(row.sort.as_ref())).collect(),
false => Vec::new(),
};
let aways: Vec<Pairs> = rows
.iter()
.map(|row| {
row.dists
.iter()
.map(|(name, away)| (name.clone(), twelve(*away).into_bytes().into()))
.collect()
})
.collect();
let mut kept: Vec<Held<'_>> = Vec::with_capacity(rows.len());
for (at, row) in rows.iter().enumerate() {
if !loading {
kept.push((row, None));
continue;
}
let Some(doc) = held.get(at).and_then(Option::as_ref) else {
lost += 1;
continue;
};
let key = keys.get(at).and_then(Option::as_deref);
let away = aways.get(at).map_or(&[][..], Pairs::as_slice);
kept.push((row, Some(pick(doc, want, key, away))));
}
let redone = marked(&kept, want);
let mut built: Vec<Built<'_>> = kept
.iter()
.zip(&redone)
.map(|((row, fields), redone)| {
let fields = fields.as_ref().map(|fields| {
fields
.iter()
.zip(redone)
.map(|((name, value), redone)| match redone {
Some(redone) => (*name, &**redone),
None => (*name, *value),
})
.collect()
});
(*row, fields)
})
.collect();
shared(&mut built, &want.upfront);
let total = total - lost;
if let Some(asks) = asked.cursor {
let made = Made::Found(
built
.into_iter()
.map(|(row, fields)| (row.clone(), fields.map(owned)))
.collect(),
);
let kept = Kept {
made,
walk: Vec::new(),
shows: want.found(),
total,
whole: true,
loader: loading,
offset: want.offset,
window: want.count,
};
cursor::open(server, index, kept, asks, out);
return;
}
found(total, &built, want.found(), out);
}
fn marked(kept: &[Held<'_>], want: &Rows<'_>) -> Vec<Vec<Option<Vec<u8>>>> {
let widths = || {
kept.iter()
.map(|(_, fields)| vec![None; fields.as_ref().map_or(0, Vec::len)])
.collect()
};
if want.summarize.is_none() && want.highlight.is_none() {
return widths();
}
let mut english = English::new();
let nothing = Wanted::default();
let stops = want.stops.as_deref();
let wanted = &want.wanted;
kept.iter()
.map(|(_, fields)| {
fields
.iter()
.flatten()
.map(|(name, value)| {
if want.distance.iter().any(|held| *held.name == **name) {
return None;
}
let named = |list: &[&[u8]]| list.contains(name);
if !want.upfront.is_empty() && !named(&want.upfront) {
return None;
}
let text = want.texts.iter().any(|held| **held == **name);
let touches = |list: &Vec<&[u8]>| list.is_empty() || named(list);
let wrap = want
.highlight
.as_ref()
.filter(|(list, _)| touches(list))
.map(|(_, wrap)| wrap);
let seen = match text {
true => wanted,
false => ¬hing,
};
match want.summarize.as_ref().filter(|(list, _)| touches(list)) {
Some((_, trim)) => Some(summary::summarize(
value,
seen,
stops,
trim,
wrap,
&mut english,
)),
None => wrap
.map(|wrap| summary::highlight(value, seen, stops, wrap, &mut english)),
}
})
.collect()
})
.collect()
}
fn shared<'a>(built: &mut [Built<'a>], upfront: &[&'a [u8]]) {
let mut order: Vec<&[u8]> = upfront.to_vec();
for (_, fields) in built.iter() {
for (name, _) in fields.iter().flatten() {
if !order.contains(name) {
order.push(name);
}
}
}
for (_, fields) in built.iter_mut() {
let Some(fields) = fields else {
continue;
};
let mut out = Vec::with_capacity(fields.len());
for name in &order {
if let Some((_, value)) = fields.iter().find(|(held, _)| held == name) {
out.push((*name, *value));
}
}
*fields = out;
}
}
fn owned(pairs: Vec<(&[u8], &[u8])>) -> Pairs {
pairs
.into_iter()
.map(|(name, value)| (name.into(), value.into()))
.collect()
}
fn worth(row: &Row, out: &mut Out) {
let Some(note) = &row.note else {
out.double(row.score);
return;
};
out.array(2);
out.double(row.score);
reason(note, out);
}
fn reason(note: &Note, out: &mut Out) {
match note {
Note::Line(line) => out.bulk(line.as_bytes()),
Note::Under(line, under) => {
out.array(2);
out.bulk(line.as_bytes());
out.array(under.len());
for note in under {
reason(note, out);
}
}
}
}
fn found(total: usize, built: &[Built<'_>], shows: Shows, out: &mut Out) {
if out.proto().is_resp3() {
deep(total, built, shows, out);
return;
}
let per = 1
+ usize::from(shows.scores)
+ usize::from(shows.payloads)
+ usize::from(shows.sortkeys)
+ usize::from(shows.fields);
out.array(1 + built.len() * per);
out.int(total as i64);
for (row, fields) in built {
out.bulk(&row.key);
if shows.scores {
worth(row, out);
}
if shows.payloads {
match &row.payload {
Some(payload) => out.bulk(payload),
None => out.nil(),
}
}
if shows.sortkeys {
match sortkey(row.sort.as_ref()) {
Some(sorted) => out.bulk(&sorted),
None => out.nil(),
}
}
if let Some(fields) = fields {
out.map(fields.len());
for (field, value) in fields {
out.bulk(field);
out.bulk(value);
}
}
}
}
fn deep(total: usize, built: &[Built<'_>], shows: Shows, out: &mut Out) {
out.map(5);
out.simple(b"attributes");
out.array(0);
out.simple(b"format");
out.simple(b"STRING");
out.simple(b"results");
out.array(built.len());
for (row, fields) in built {
out.map(
2 + usize::from(shows.scores)
+ usize::from(shows.payloads)
+ usize::from(shows.sortkeys)
+ usize::from(fields.is_some()),
);
out.simple(b"id");
out.bulk(&row.key);
if shows.scores {
out.simple(b"score");
worth(row, out);
}
if shows.payloads {
out.simple(b"payload");
match &row.payload {
Some(payload) => out.bulk(payload),
None => out.nil(),
}
}
if shows.sortkeys {
out.simple(b"sortkey");
match sortkey(row.sort.as_ref()) {
Some(sorted) => out.bulk(&sorted),
None => out.nil(),
}
}
if let Some(fields) = fields {
out.simple(b"extra_attributes");
out.map(fields.len());
for (field, value) in fields {
out.bulk(field);
out.bulk(value);
}
}
out.simple(b"values");
out.array(0);
}
out.simple(b"total_results");
out.int(total as i64);
out.simple(b"warning");
out.array(0);
}
fn pick<'d, 'w: 'd>(
doc: &'d indexing::Document,
want: &'d Rows<'w>,
key: Option<&'d [u8]>,
away: &'d [Named],
) -> Vec<(&'d [u8], &'d [u8])> {
let pairs = doc.pairs();
let by = want.sorting.as_ref();
let Some(ret) = &want.ret else {
let sorted = by.zip(key);
if sorted.is_none() && away.is_empty() {
return pairs;
}
let mut out: Vec<(&[u8], &[u8])> = Vec::with_capacity(pairs.len() + away.len() + 1);
for (name, value) in away {
out.push((&**name, &**value));
}
if let Some((by, key)) = sorted
&& !out.iter().any(|(held, _)| *held == &*by.field.attribute)
{
out.push((&by.field.attribute, key));
}
for (name, value) in pairs {
match out.iter_mut().find(|(held, _)| *held == name) {
Some(held) => held.1 = value,
None => out.push((name, value)),
}
}
return out;
};
let mut out: Vec<(&[u8], &[u8])> = ret
.iter()
.filter_map(|(field, name)| {
if let Some((_, value)) = away.iter().find(|(held, _)| **held == **name) {
return Some((&**name, &**value));
}
let (_, value) = pairs.iter().find(|(held, _)| *held == &**field)?;
Some((&**name, *value))
})
.collect();
if let Some(by) = by
&& let Some(at) = out
.iter()
.position(|(name, _)| *name == &*by.field.attribute)
{
let held = out.remove(at);
out.insert(0, held);
}
let mut front: Vec<(&[u8], &[u8])> = Vec::with_capacity(away.len());
for (name, _) in away {
if let Some(at) = out.iter().position(|(held, _)| *held == &**name) {
front.push(out.remove(at));
}
}
front.append(&mut out);
front
}
fn shown(sort: Option<&Sorted>) -> Option<Vec<u8>> {
match sort? {
Sorted::Number(number) => Some(twelve(*number).into_bytes()),
Sorted::Text(text) => Some(text.to_vec()),
}
}
fn sortkey(sort: Option<&Sorted>) -> Option<Vec<u8>> {
let mut out = Vec::new();
match sort? {
Sorted::Number(number) => {
out.push(b'#');
out.extend_from_slice(seventeen(*number).as_bytes());
}
Sorted::Text(text) => {
out.push(b'$');
out.extend_from_slice(text);
}
}
Some(out)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_number_that_is_not_one_is_spelled_the_way_the_wire_spells_it() {
assert_eq!(twelve(f64::NAN), "nan");
assert_eq!(twelve(f64::INFINITY), "inf");
assert_eq!(twelve(f64::NEG_INFINITY), "-inf");
assert_eq!(twelve(0.0), "0");
assert_eq!(twelve(0.934_309_237_376_833_4), "0.934309237377");
}
#[test]
fn a_geofilter_reads_a_number_the_way_that_one_option_reads_one() {
assert_eq!(coord(b"0"), Some(0.0));
assert_eq!(coord(b"-0.1278"), Some(-0.1278));
assert_eq!(coord(b"1e2"), Some(100.0));
assert_eq!(coord(b" 5"), Some(5.0));
assert_eq!(coord(b"5 "), None);
assert!(coord(b"nan").is_some_and(f64::is_nan));
assert_eq!(coord(b"inf"), None);
assert_eq!(coord(b"-inf"), None);
assert_eq!(coord(b"1e400"), None);
assert_eq!(coord(b"x"), None);
assert_eq!(coord(b""), None);
}
#[test]
fn a_sort_key_is_written_five_digits_wider_than_the_row_it_sits_beside() {
use yo_search::expr::seventeen;
assert_eq!(seventeen(1.0 / 3.0), "0.33333333333333331");
assert_eq!(seventeen(-4.0), "-4");
assert_eq!(seventeen(5.5), "5.5");
assert_eq!(seventeen(1e16), "10000000000000000");
assert_eq!(seventeen(1e17), "1e+17");
assert_eq!(seventeen(0.000_01), "1.0000000000000001e-05");
assert_eq!(seventeen(0.000_1), "0.0001");
}
#[test]
fn a_sort_key_says_which_of_the_two_kinds_of_value_it_is() {
let number = Sorted::Number(1.0 / 3.0);
let text = Sorted::Text(b"four score".as_slice().into());
assert_eq!(
sortkey(Some(&number)).as_deref(),
Some(&b"#0.33333333333333331"[..])
);
assert_eq!(sortkey(Some(&text)).as_deref(), Some(&b"$four score"[..]));
assert_eq!(sortkey(None), None);
assert_eq!(
shown(Some(&number)).as_deref(),
Some(&b"0.333333333333"[..])
);
assert_eq!(shown(Some(&text)).as_deref(), Some(&b"four score"[..]));
assert_eq!(shown(None), None);
}
#[test]
fn every_row_of_a_window_is_written_in_the_order_the_window_shares() {
let row = Row {
key: Box::default(),
score: 0.0,
payload: None,
note: None,
sort: None,
dists: Vec::new(),
};
let mut built: Vec<Built<'_>> = vec![
(
&row,
Some(vec![(&b"n"[..], &b"0"[..]), (&b"p"[..], &b"a"[..])]),
),
(
&row,
Some(vec![
(&b"t"[..], &b"x"[..]),
(&b"p"[..], &b"b"[..]),
(&b"n"[..], &b"2"[..]),
]),
),
(&row, None),
];
shared(&mut built, &[]);
assert_eq!(
built[1].1.as_deref(),
Some(
&[
(&b"n"[..], &b"2"[..]),
(&b"p"[..], &b"b"[..]),
(&b"t"[..], &b"x"[..])
][..]
)
);
assert_eq!(built[0].1.as_ref().map(Vec::len), Some(2));
assert_eq!(built[2].1, None);
}
}