use crate::value::IndexValue;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ValType {
Vector,
I64,
F64,
Str,
}
impl ValType {
pub fn tag(self) -> &'static str {
match self {
ValType::I64 => "i64",
ValType::F64 => "f64",
ValType::Str => "str",
ValType::Vector => "vector",
}
}
pub fn parse(raw: &[u8]) -> Option<ValType> {
if raw.eq_ignore_ascii_case(b"i64") {
Some(ValType::I64)
} else if raw.eq_ignore_ascii_case(b"f64") {
Some(ValType::F64)
} else if raw.eq_ignore_ascii_case(b"str") {
Some(ValType::Str)
} else if raw.eq_ignore_ascii_case(b"vector") {
Some(ValType::Vector)
} else {
None
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IndexKind {
Range,
Unique,
Text,
Ann,
Agg,
}
impl IndexKind {
pub fn tag(self) -> &'static str {
match self {
IndexKind::Range => "range",
IndexKind::Unique => "unique",
IndexKind::Text => "text",
IndexKind::Ann => "ann",
IndexKind::Agg => "agg",
}
}
pub fn parse(raw: &[u8]) -> Option<IndexKind> {
if raw.eq_ignore_ascii_case(b"range") {
Some(IndexKind::Range)
} else if raw.eq_ignore_ascii_case(b"unique") {
Some(IndexKind::Unique)
} else if raw.eq_ignore_ascii_case(b"text") {
Some(IndexKind::Text)
} else if raw.eq_ignore_ascii_case(b"ann") {
Some(IndexKind::Ann)
} else if raw.eq_ignore_ascii_case(b"agg") {
Some(IndexKind::Agg)
} else {
None
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IndexState {
Building,
Ready,
FailedOverBudget,
}
#[derive(Debug, Clone, PartialEq)]
pub struct FieldSpec {
pub name: Vec<u8>,
pub weight: f32,
}
impl FieldSpec {
pub fn new(name: impl Into<Vec<u8>>) -> FieldSpec {
FieldSpec { name: name.into(), weight: 1.0 }
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct ValueSpec {
pub name: Vec<u8>,
pub ty: ValType,
}
impl ValueSpec {
pub fn new(name: impl Into<Vec<u8>>) -> ValueSpec {
ValueSpec { name: name.into(), ty: ValType::Str }
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct IndexSpec {
pub name: Vec<u8>,
pub prefix: Vec<u8>,
pub fields: Vec<FieldSpec>,
pub ty: ValType,
pub kind: IndexKind,
pub max_bytes: u64,
pub ann: Option<AnnSpec>,
pub group_by: Option<Vec<u8>>,
pub with_positions: bool,
pub values: Vec<ValueSpec>,
pub composite: Option<Vec<crate::composite::CompositeCol>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AnnSpec {
pub dim: u32,
pub distance: u8,
pub m: u16,
pub ef: u16,
}
pub const MAX_INDEXES: usize = 64;
#[derive(Debug, Clone, Default)]
pub struct Catalog {
pub(crate) specs: Vec<(IndexSpec, IndexState)>,
}
pub type RowInputs = (Vec<(Vec<u8>, f32)>, Vec<Option<Vec<u8>>>);
impl IndexSpec {
pub fn read_row(
&self,
mut get: impl FnMut(&[u8]) -> Option<Vec<u8>>,
) -> RowInputs {
let mut fields = Vec::with_capacity(self.fields.len());
for f in &self.fields {
if let Some(raw) = get(&f.name) {
fields.push((raw, f.weight));
}
}
let values = self.values.iter().map(|v| get(&v.name)).collect();
(fields, values)
}
pub fn field(&self) -> &[u8] {
self.fields.first().map_or(&[][..], |f| f.name.as_slice())
}
pub fn single_field(
name: Vec<u8>,
prefix: Vec<u8>,
field: Vec<u8>,
ty: ValType,
kind: IndexKind,
) -> IndexSpec {
IndexSpec {
name,
prefix,
fields: vec![FieldSpec::new(field)],
ty,
kind,
max_bytes: 0,
ann: None,
group_by: None,
with_positions: false,
values: Vec::new(),
composite: None,
}
}
}
impl Catalog {
pub fn new() -> Self {
Self::default()
}
pub fn create(&mut self, spec: IndexSpec) -> Result<(), &'static str> {
if self.specs.len() >= MAX_INDEXES {
return Err("ERR index limit reached (64)");
}
if self.specs.iter().any(|(s, _)| s.name == spec.name) {
return Err("ERR index already exists");
}
if spec.fields.is_empty() {
return Err("ERR index needs at least one field");
}
if spec.fields.len() > 1 && spec.kind != IndexKind::Text {
return Err("ERR only KIND text indexes several fields");
}
if spec.with_positions && spec.kind != IndexKind::Text {
return Err("ERR WITH POSITIONS requires KIND text");
}
if !spec.values.is_empty()
&& !matches!(spec.kind, IndexKind::Text | IndexKind::Range | IndexKind::Unique)
{
return Err("ERR VALUES requires KIND text|range|unique");
}
crate::composite::composite_guard(&spec)?;
self.specs.push((spec, IndexState::Building));
Ok(())
}
pub fn drop_index(&mut self, name: &[u8]) -> bool {
let before = self.specs.len();
self.specs.retain(|(s, _)| s.name != name);
self.specs.len() != before
}
pub fn set_state(&mut self, name: &[u8], state: IndexState) -> bool {
for (s, st) in &mut self.specs {
if s.name == name {
*st = state;
return true;
}
}
false
}
pub fn get(&self, name: &[u8]) -> Option<(&IndexSpec, IndexState)> {
self.specs.iter().find(|(s, _)| s.name == name).map(|(s, st)| (s, *st))
}
pub fn iter(&self) -> impl Iterator<Item = (&IndexSpec, IndexState)> {
self.specs.iter().map(|(s, st)| (s, *st))
}
pub fn len(&self) -> usize {
self.specs.len()
}
pub fn is_empty(&self) -> bool {
self.specs.is_empty()
}
pub fn matching<'a>(&'a self, key: &'a [u8]) -> impl Iterator<Item = (&'a IndexSpec, IndexState)> {
self.specs
.iter()
.filter(move |(s, _)| key.starts_with(&s.prefix))
.map(|(s, st)| (s, *st))
}
pub fn coerce(spec: &IndexSpec, raw: &[u8]) -> Option<IndexValue> {
IndexValue::coerce(spec.ty, raw)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn spec(name: &str, prefix: &str) -> IndexSpec {
IndexSpec {
name: name.into(),
prefix: prefix.into(),
fields: vec![FieldSpec::new(b"age".to_vec())],
ty: ValType::I64,
kind: IndexKind::Range,
ann: None,
max_bytes: 0,
group_by: None,
with_positions: false,
values: Vec::new(),
composite: None,
}
}
#[test]
fn create_drop_match_lifecycle() {
let mut c = Catalog::new();
c.create(spec("a", "user:")).unwrap();
c.create(spec("b", "sess:")).unwrap();
assert!(c.create(spec("a", "x:")).is_err(), "dup name");
assert_eq!(c.matching(b"user:42").count(), 1);
assert_eq!(c.matching(b"other:1").count(), 0);
assert_eq!(c.get(b"a").unwrap().1, IndexState::Building);
assert!(c.set_state(b"a", IndexState::Ready));
assert_eq!(c.get(b"a").unwrap().1, IndexState::Ready);
assert!(c.drop_index(b"b"));
assert!(!c.drop_index(b"b"));
assert_eq!(c.len(), 1);
}
#[test]
fn sidecar_roundtrip_with_escapes() {
let mut c = Catalog::new();
let mut s = spec("weird", "pre\tfix:");
s.fields = vec![FieldSpec::new(b"f%\n".to_vec())];
s.max_bytes = 1024;
c.create(s).unwrap();
let text = c.to_sidecar();
let c2 = Catalog::from_sidecar(&text).unwrap();
let (got, st) = c2.get(b"weird").unwrap();
assert_eq!(got.prefix, b"pre\tfix:".to_vec());
assert_eq!(got.field(), b"f%\n");
assert_eq!(got.max_bytes, 1024);
assert_eq!(st, IndexState::Building, "boot loads as Building");
assert!(Catalog::from_sidecar("bogus").is_none());
}
#[test]
fn only_text_indexes_accept_several_fields() {
let two = || {
vec![
FieldSpec::new(b"title".to_vec()),
FieldSpec::new(b"body".to_vec()),
]
};
let mut range = spec("multi-range", "p:");
range.fields = two();
assert!(Catalog::new().create(range).is_err(), "range must refuse two fields");
let mut text = spec("multi-text", "p:");
text.kind = IndexKind::Text;
text.fields = two();
assert!(Catalog::new().create(text).is_ok(), "text must accept them");
}
#[test]
fn an_index_needs_at_least_one_field() {
let mut s = spec("nofields", "p:");
s.fields.clear();
assert!(Catalog::new().create(s).is_err());
}
#[test]
fn cap_enforced() {
let mut c = Catalog::new();
for i in 0..MAX_INDEXES {
c.create(spec(&format!("i{i}"), "p:")).unwrap();
}
assert!(c.create(spec("over", "p:")).is_err());
}
}