///|
/// Selects an index strategy based on corpus size and caller requirements.
pub(all) enum IndexStrategy {
ExactFlat
ClusteredIvf
SpatialKdTree
ApproximateLsh
}
///|
pub impl Show for IndexStrategy with fn output(self, logger) {
let text = match self {
ExactFlat => "ExactFlat"
ClusteredIvf => "ClusteredIvf"
SpatialKdTree => "SpatialKdTree"
ApproximateLsh => "ApproximateLsh"
}
logger.write_string(text)
}
///|
/// Query options shared by application-facing search adapters.
pub(all) struct QueryPlan {
query : Array[Double]
top_k : Int
metric : DistanceMetric
filters : Array[(String, String)]
strategy : IndexStrategy
nprobe : Int
}
///|
/// Create a conservative exact query plan.
pub fn QueryPlan::exact(
query : Array[Double],
top_k : Int,
metric : DistanceMetric,
filters : Array[(String, String)],
) -> QueryPlan {
{
query,
top_k: if top_k < 0 {
0
} else {
top_k
},
metric,
filters,
strategy: ExactFlat,
nprobe: 1,
}
}
///|
/// Create an IVF plan with a bounded probe count.
pub fn QueryPlan::ivf(
query : Array[Double],
top_k : Int,
metric : DistanceMetric,
filters : Array[(String, String)],
nprobe : Int,
) -> QueryPlan {
{
query,
top_k: if top_k < 0 {
0
} else {
top_k
},
metric,
filters,
strategy: ClusteredIvf,
nprobe: if nprobe < 1 {
1
} else {
nprobe
},
}
}
///|
/// Recommend a strategy using simple, explainable defaults.
pub fn recommend_strategy(
corpus_size : Int,
dimension : Int,
exact_required : Bool,
) -> IndexStrategy {
if exact_required || corpus_size <= 256 {
return ExactFlat
}
if dimension <= 16 {
return SpatialKdTree
}
if corpus_size >= 4096 {
return ApproximateLsh
}
ClusteredIvf
}
///|
/// Explain the recommendation in user-facing language.
pub fn explain_strategy(strategy : IndexStrategy) -> String {
match strategy {
ExactFlat =>
"Flat search scans every vector and is the correctness baseline."
ClusteredIvf =>
"IVF-Flat narrows search to nearby K-Means clusters; tune nprobe for recall."
SpatialKdTree =>
"KD-Tree partitions low-dimensional Euclidean or Manhattan space exactly."
ApproximateLsh =>
"LSH uses projection buckets for fast approximate cosine retrieval."
}
}
///|
/// Execute a plan with the exact Flat index.
pub fn execute_exact_plan(
index : FlatIndex,
plan : QueryPlan,
) -> Array[SearchResult] raise VectorError {
index.search(plan.query, plan.top_k, plan.metric, plan.filters)
}
///|
/// Execute a plan with an already-built IVF index.
pub fn execute_ivf_plan(
index : IvfIndex,
plan : QueryPlan,
) -> Array[SearchResult] raise VectorError {
index.search(plan.query, plan.top_k, plan.nprobe, plan.filters)
}
///|
/// Execute a plan with an already-built KD-Tree index.
pub fn execute_kd_plan(
index : KdTreeIndex,
plan : QueryPlan,
) -> Array[SearchResult] raise VectorError {
index.search(plan.query, plan.top_k, plan.filters)
}
///|
/// Execute a plan with an already-built LSH index.
pub fn execute_lsh_plan(
index : LshIndex,
plan : QueryPlan,
) -> Array[SearchResult] raise VectorError {
index.search(plan.query, plan.top_k, plan.filters)
}
///|
/// Build an exact index and execute a plan against a corpus.
pub fn search_documents(
docs : Array[Document],
plan : QueryPlan,
) -> Array[SearchResult] raise VectorError {
validate_documents(docs)
let index = FlatIndex::new()
for doc in docs {
index.add(doc)
}
execute_exact_plan(index, plan)
}
///|
/// Return a plan with a stricter top-k bound, useful for public API limits.
pub fn cap_query_plan(plan : QueryPlan, maximum_k : Int) -> QueryPlan {
let max_k = if maximum_k < 0 { 0 } else { maximum_k }
let top_k = if plan.top_k < max_k { plan.top_k } else { max_k }
{
query: plan.query,
top_k,
metric: plan.metric,
filters: plan.filters,
strategy: plan.strategy,
nprobe: plan.nprobe,
}
}
///|
/// Return whether all filters are syntactically non-empty.
pub fn filters_are_valid(filters : Array[(String, String)]) -> Bool {
for pair in filters {
if pair.0.trim().length() == 0 || pair.1.trim().length() == 0 {
return false
}
}
true
}
///|
/// Return a plan with filters that match the documented empty-value policy.
pub fn sanitize_query_plan(plan : QueryPlan) -> QueryPlan {
let filters = []
for pair in plan.filters {
if pair.0.trim().length() > 0 && pair.1.trim().length() > 0 {
filters.push((pair.0.trim().to_owned(), pair.1.trim().to_owned()))
}
}
{
query: plan.query,
top_k: plan.top_k,
metric: plan.metric,
filters,
strategy: plan.strategy,
nprobe: plan.nprobe,
}
}