use std::collections::HashMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VectorAlgorithm {
Hnsw,
Flat,
}
impl std::fmt::Display for VectorAlgorithm {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
VectorAlgorithm::Hnsw => write!(f, "HNSW"),
VectorAlgorithm::Flat => write!(f, "FLAT"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DistanceMetric {
L2,
InnerProduct,
Cosine,
}
impl std::fmt::Display for DistanceMetric {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
DistanceMetric::L2 => write!(f, "L2"),
DistanceMetric::InnerProduct => write!(f, "IP"),
DistanceMetric::Cosine => write!(f, "COSINE"),
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct VectorParams {
pub algorithm: VectorAlgorithm,
pub dim: usize,
pub distance_metric: DistanceMetric,
pub hnsw_m: Option<usize>,
pub hnsw_ef_construction: Option<usize>,
}
impl VectorParams {
pub fn hnsw(dim: usize, distance_metric: DistanceMetric) -> Self {
Self {
algorithm: VectorAlgorithm::Hnsw,
dim,
distance_metric,
hnsw_m: None,
hnsw_ef_construction: None,
}
}
pub fn flat(dim: usize, distance_metric: DistanceMetric) -> Self {
Self {
algorithm: VectorAlgorithm::Flat,
dim,
distance_metric,
hnsw_m: None,
hnsw_ef_construction: None,
}
}
pub fn with_m(mut self, m: usize) -> Self {
self.hnsw_m = Some(m);
self
}
pub fn with_ef_construction(mut self, ef: usize) -> Self {
self.hnsw_ef_construction = Some(ef);
self
}
fn to_schema_args(&self) -> Vec<String> {
let mut args = vec![
"TYPE".to_string(),
"FLOAT32".to_string(),
"DIM".to_string(),
self.dim.to_string(),
"DISTANCE_METRIC".to_string(),
self.distance_metric.to_string(),
];
if self.algorithm == VectorAlgorithm::Hnsw {
if let Some(m) = self.hnsw_m {
args.push("M".to_string());
args.push(m.to_string());
}
if let Some(ef) = self.hnsw_ef_construction {
args.push("EF_CONSTRUCTION".to_string());
args.push(ef.to_string());
}
}
let mut result = vec![args.len().to_string()];
result.extend(args);
result
}
}
#[derive(Debug, Clone)]
pub struct SearchIndex {
pub name: String,
pub prefix: String,
pub fields: Vec<SearchField>,
}
impl SearchIndex {
pub fn new(name: impl Into<String>, prefix: impl Into<String>) -> Self {
Self {
name: name.into(),
prefix: prefix.into(),
fields: Vec::new(),
}
}
pub fn text(mut self, name: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: None,
field_type: SearchFieldType::Text,
sortable: false,
no_index: false,
});
self
}
pub fn text_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: Some(json_path.into()),
field_type: SearchFieldType::Text,
sortable: false,
no_index: false,
});
self
}
pub fn text_sortable(mut self, name: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: None,
field_type: SearchFieldType::Text,
sortable: true,
no_index: false,
});
self
}
pub fn text_sortable_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: Some(json_path.into()),
field_type: SearchFieldType::Text,
sortable: true,
no_index: false,
});
self
}
pub fn numeric(mut self, name: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: None,
field_type: SearchFieldType::Numeric,
sortable: false,
no_index: false,
});
self
}
pub fn numeric_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: Some(json_path.into()),
field_type: SearchFieldType::Numeric,
sortable: false,
no_index: false,
});
self
}
pub fn numeric_sortable(mut self, name: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: None,
field_type: SearchFieldType::Numeric,
sortable: true,
no_index: false,
});
self
}
pub fn numeric_sortable_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: Some(json_path.into()),
field_type: SearchFieldType::Numeric,
sortable: true,
no_index: false,
});
self
}
pub fn tag(mut self, name: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: None,
field_type: SearchFieldType::Tag,
sortable: false,
no_index: false,
});
self
}
pub fn tag_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: Some(json_path.into()),
field_type: SearchFieldType::Tag,
sortable: false,
no_index: false,
});
self
}
pub fn geo(mut self, name: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: None,
field_type: SearchFieldType::Geo,
sortable: false,
no_index: false,
});
self
}
pub fn geo_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: Some(json_path.into()),
field_type: SearchFieldType::Geo,
sortable: false,
no_index: false,
});
self
}
pub fn vector_hnsw(
mut self,
name: impl Into<String>,
dim: usize,
metric: DistanceMetric,
) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: None,
field_type: SearchFieldType::Vector(VectorParams::hnsw(dim, metric)),
sortable: false,
no_index: false,
});
self
}
pub fn vector_hnsw_at(
mut self,
name: impl Into<String>,
json_path: impl Into<String>,
dim: usize,
metric: DistanceMetric,
) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: Some(json_path.into()),
field_type: SearchFieldType::Vector(VectorParams::hnsw(dim, metric)),
sortable: false,
no_index: false,
});
self
}
pub fn vector_with_params(mut self, name: impl Into<String>, params: VectorParams) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: None,
field_type: SearchFieldType::Vector(params),
sortable: false,
no_index: false,
});
self
}
pub fn vector_flat(
mut self,
name: impl Into<String>,
dim: usize,
metric: DistanceMetric,
) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: None,
field_type: SearchFieldType::Vector(VectorParams::flat(dim, metric)),
sortable: false,
no_index: false,
});
self
}
pub fn vector_flat_at(
mut self,
name: impl Into<String>,
json_path: impl Into<String>,
dim: usize,
metric: DistanceMetric,
) -> Self {
self.fields.push(SearchField {
name: name.into(),
json_path: Some(json_path.into()),
field_type: SearchFieldType::Vector(VectorParams::flat(dim, metric)),
sortable: false,
no_index: false,
});
self
}
pub fn to_ft_create_args(&self) -> Vec<String> {
self.to_ft_create_args_with_prefix(None)
}
pub fn to_ft_create_args_with_prefix(&self, redis_prefix: Option<&str>) -> Vec<String> {
let prefix = redis_prefix.unwrap_or("");
let mut args = vec![
format!("{}idx:{}", prefix, self.name),
"ON".to_string(),
"JSON".to_string(),
"PREFIX".to_string(),
"1".to_string(),
format!("{}{}", prefix, self.prefix),
"SCHEMA".to_string(),
];
for field in &self.fields {
args.extend(field.to_schema_args());
}
args
}
}
#[derive(Debug, Clone)]
pub struct SearchField {
pub name: String,
pub json_path: Option<String>,
pub field_type: SearchFieldType,
pub sortable: bool,
pub no_index: bool,
}
impl SearchField {
fn to_schema_args(&self) -> Vec<String> {
let json_path = self
.json_path
.clone()
.unwrap_or_else(|| format!("$.payload.{}", self.name));
let mut args = vec![
json_path,
"AS".to_string(),
self.name.clone(),
];
match &self.field_type {
SearchFieldType::Vector(params) => {
args.push("VECTOR".to_string());
args.push(params.algorithm.to_string());
args.extend(params.to_schema_args());
}
_ => {
args.push(self.field_type.to_string());
}
}
if self.sortable {
args.push("SORTABLE".to_string());
}
if self.no_index {
args.push("NOINDEX".to_string());
}
args
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum SearchFieldType {
Text,
Numeric,
Tag,
Geo,
Vector(VectorParams),
}
impl std::fmt::Display for SearchFieldType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SearchFieldType::Text => write!(f, "TEXT"),
SearchFieldType::Numeric => write!(f, "NUMERIC"),
SearchFieldType::Tag => write!(f, "TAG"),
SearchFieldType::Geo => write!(f, "GEO"),
SearchFieldType::Vector(params) => write!(f, "VECTOR {}", params.algorithm),
}
}
}
pub struct IndexManager {
indexes: HashMap<String, SearchIndex>,
}
impl IndexManager {
pub fn new() -> Self {
Self {
indexes: HashMap::new(),
}
}
pub fn register(&mut self, index: SearchIndex) {
self.indexes.insert(index.name.clone(), index);
}
pub fn get(&self, name: &str) -> Option<&SearchIndex> {
self.indexes.get(name)
}
pub fn all(&self) -> impl Iterator<Item = &SearchIndex> {
self.indexes.values()
}
pub fn ft_create_args(&self, name: &str) -> Option<Vec<String>> {
self.indexes.get(name).map(|idx| idx.to_ft_create_args())
}
pub fn find_by_prefix(&self, key: &str) -> Option<&SearchIndex> {
self.indexes.values().find(|idx| key.starts_with(&idx.prefix))
}
}
impl Default for IndexManager {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_simple_index() {
let index = SearchIndex::new("users", "crdt:users:")
.text("name")
.text("email")
.numeric("age");
let args = index.to_ft_create_args();
assert_eq!(args[0], "idx:users");
assert_eq!(args[1], "ON");
assert_eq!(args[2], "JSON");
assert_eq!(args[3], "PREFIX");
assert_eq!(args[4], "1");
assert_eq!(args[5], "crdt:users:");
assert_eq!(args[6], "SCHEMA");
assert!(args.contains(&"$.payload.name".to_string()));
assert!(args.contains(&"name".to_string()));
assert!(args.contains(&"TEXT".to_string()));
}
#[test]
fn test_sortable_fields() {
let index = SearchIndex::new("users", "crdt:users:")
.text_sortable("name")
.numeric_sortable("age");
let args = index.to_ft_create_args();
let sortable_count = args.iter().filter(|a| *a == "SORTABLE").count();
assert_eq!(sortable_count, 2);
}
#[test]
fn test_tag_field() {
let index = SearchIndex::new("items", "crdt:items:").tag("tags");
let args = index.to_ft_create_args();
assert!(args.contains(&"TAG".to_string()));
}
#[test]
fn test_custom_json_path() {
let index =
SearchIndex::new("users", "crdt:users:").text_at("username", "$.profile.name");
let args = index.to_ft_create_args();
assert!(args.contains(&"$.profile.name".to_string()));
assert!(args.contains(&"username".to_string()));
}
#[test]
fn test_index_manager_register() {
let mut manager = IndexManager::new();
let index = SearchIndex::new("users", "crdt:users:")
.text("name")
.numeric("age");
manager.register(index);
assert!(manager.get("users").is_some());
assert!(manager.get("unknown").is_none());
}
#[test]
fn test_index_manager_find_by_prefix() {
let mut manager = IndexManager::new();
manager.register(SearchIndex::new("users", "crdt:users:"));
manager.register(SearchIndex::new("posts", "crdt:posts:"));
let found = manager.find_by_prefix("crdt:users:abc123");
assert!(found.is_some());
assert_eq!(found.unwrap().name, "users");
let found = manager.find_by_prefix("crdt:posts:xyz");
assert!(found.is_some());
assert_eq!(found.unwrap().name, "posts");
let not_found = manager.find_by_prefix("crdt:comments:1");
assert!(not_found.is_none());
}
#[test]
fn test_ft_create_full_command() {
let index = SearchIndex::new("users", "crdt:users:")
.text_sortable("name")
.text("email")
.numeric_sortable("age")
.tag("roles");
let args = index.to_ft_create_args();
assert_eq!(args[0], "idx:users");
assert_eq!(args[6], "SCHEMA");
let cmd = format!("FT.CREATE {}", args.join(" "));
assert!(cmd.contains("idx:users"));
assert!(cmd.contains("ON JSON"));
assert!(cmd.contains("PREFIX 1 crdt:users:"));
assert!(cmd.contains("$.payload.name AS name TEXT SORTABLE"));
assert!(cmd.contains("$.payload.email AS email TEXT"));
assert!(cmd.contains("$.payload.age AS age NUMERIC SORTABLE"));
assert!(cmd.contains("$.payload.roles AS roles TAG"));
}
#[test]
fn test_vector_hnsw_basic() {
let index = SearchIndex::new("docs", "crdt:docs:")
.text("title")
.vector_hnsw("embedding", 1536, DistanceMetric::Cosine);
let args = index.to_ft_create_args();
let cmd = format!("FT.CREATE {}", args.join(" "));
assert!(cmd.contains("$.payload.embedding AS embedding VECTOR HNSW"));
assert!(cmd.contains("TYPE FLOAT32"));
assert!(cmd.contains("DIM 1536"));
assert!(cmd.contains("DISTANCE_METRIC COSINE"));
}
#[test]
fn test_vector_hnsw_with_params() {
let params = VectorParams::hnsw(768, DistanceMetric::L2)
.with_m(32)
.with_ef_construction(400);
let index =
SearchIndex::new("embeddings", "crdt:embeddings:").vector_with_params("vec", params);
let args = index.to_ft_create_args();
let cmd = format!("FT.CREATE {}", args.join(" "));
assert!(cmd.contains("VECTOR HNSW"));
assert!(cmd.contains("DIM 768"));
assert!(cmd.contains("DISTANCE_METRIC L2"));
assert!(cmd.contains("M 32"));
assert!(cmd.contains("EF_CONSTRUCTION 400"));
}
#[test]
fn test_vector_flat() {
let index = SearchIndex::new("small", "crdt:small:")
.vector_flat("embedding", 384, DistanceMetric::InnerProduct);
let args = index.to_ft_create_args();
let cmd = format!("FT.CREATE {}", args.join(" "));
assert!(cmd.contains("VECTOR FLAT"));
assert!(cmd.contains("DIM 384"));
assert!(cmd.contains("DISTANCE_METRIC IP"));
}
#[test]
fn test_vector_custom_path() {
let index = SearchIndex::new("docs", "crdt:docs:").vector_hnsw_at(
"embedding",
"$.metadata.vector",
512,
DistanceMetric::Cosine,
);
let args = index.to_ft_create_args();
let cmd = format!("FT.CREATE {}", args.join(" "));
assert!(cmd.contains("$.metadata.vector AS embedding VECTOR HNSW"));
}
#[test]
fn test_vector_params_nargs() {
let params = VectorParams::hnsw(1536, DistanceMetric::Cosine);
let args = params.to_schema_args();
assert_eq!(args[0], "6");
assert_eq!(args[1], "TYPE");
assert_eq!(args[2], "FLOAT32");
assert_eq!(args[3], "DIM");
assert_eq!(args[4], "1536");
assert_eq!(args[5], "DISTANCE_METRIC");
assert_eq!(args[6], "COSINE");
}
#[test]
fn test_vector_params_nargs_with_hnsw_options() {
let params = VectorParams::hnsw(1536, DistanceMetric::Cosine)
.with_m(24)
.with_ef_construction(300);
let args = params.to_schema_args();
assert_eq!(args[0], "10");
assert!(args.contains(&"M".to_string()));
assert!(args.contains(&"24".to_string()));
assert!(args.contains(&"EF_CONSTRUCTION".to_string()));
assert!(args.contains(&"300".to_string()));
}
#[test]
fn test_mixed_index_with_vectors() {
let index = SearchIndex::new("documents", "crdt:documents:")
.text_sortable("title")
.text("content")
.tag("category")
.numeric("created_at")
.vector_hnsw("embedding", 1536, DistanceMetric::Cosine);
let args = index.to_ft_create_args();
let cmd = format!("FT.CREATE {}", args.join(" "));
assert!(cmd.contains("AS title TEXT SORTABLE"));
assert!(cmd.contains("AS content TEXT"));
assert!(cmd.contains("AS category TAG"));
assert!(cmd.contains("AS created_at NUMERIC"));
assert!(cmd.contains("AS embedding VECTOR HNSW"));
}
#[test]
fn test_distance_metrics_display() {
assert_eq!(format!("{}", DistanceMetric::L2), "L2");
assert_eq!(format!("{}", DistanceMetric::InnerProduct), "IP");
assert_eq!(format!("{}", DistanceMetric::Cosine), "COSINE");
}
#[test]
fn test_vector_algorithm_display() {
assert_eq!(format!("{}", VectorAlgorithm::Hnsw), "HNSW");
assert_eq!(format!("{}", VectorAlgorithm::Flat), "FLAT");
}
}