///|
/// IVF (Inverted File) index implementation.
/// Uses k-means clustering for approximate nearest neighbor search.
pub struct IVFState {
metric : @types.Metric
nlist : Int // number of clusters
nprobe : Int // number of clusters to search
mut centroid_count : Int
centroids : Array[Double] // flattened: centroids[c*dim..(c+1)*dim]
lists : Array[Array[@types.VectorId]] // posting lists
id_to_list : Map[@types.VectorId, Int] // id -> list index
}
///|
/// Create IVF state
pub fn IVFState::new(
params : @types.IVFParams,
metric : @types.Metric,
dim : Int,
) -> IVFState {
let nlist = if params.nlist < 1 { 1 } else { params.nlist }
let nprobe_default = nlist.to_double().sqrt().to_int()
let nprobe = if params.nprobe < 1 {
if nprobe_default < 1 {
1
} else {
nprobe_default
}
} else if params.nprobe > nlist {
nlist
} else {
params.nprobe
}
let lists : Array[Array[@types.VectorId]] = []
for _ in 0.. Int {
let dim = store.dim
let score_fn = @vecmath.get_score_fn(state.metric)
let mut best = -1
let mut best_score = -1.0e308
let count = if state.centroid_count < 1 { 1 } else { state.centroid_count }
for c in 0.. best_score {
best_score = sc
best = c
}
}
if best < 0 {
0
} else {
best
}
}
///|
/// Add a vector to IVF
pub fn ivf_add(
state : IVFState,
store : @store.CoreStore,
id : @types.VectorId,
) -> Unit {
with_store_vector(store, id, fn(_, vec) {
// Initialize centroids with first nlist vectors
if state.centroid_count < state.nlist {
let c = state.centroid_count
for j in 0.. Unit {
with_store_vector(store, id, fn(_, vec) {
// Always assign to nearest existing centroid (no bootstrap)
let c = nearest_centroid(state, store, vec)
state.lists[c].push(id)
state.id_to_list.set(id, c)
})
}
///|
/// Remove a vector from IVF
pub fn ivf_remove(state : IVFState, id : @types.VectorId) -> Unit {
match state.id_to_list.get(id) {
None => ()
Some(li) => {
let arr = state.lists[li]
// Find and remove id from list
let mut found_idx = -1
for i in 0..= 0 {
// Swap with last and pop
let last = arr.length() - 1
if found_idx != last {
arr[found_idx] = arr[last]
}
let _ = arr.pop()
}
state.id_to_list.remove(id)
}
}
}
///|
/// IVF search - find k nearest neighbors
pub fn ivf_search(
state : IVFState,
store : @store.CoreStore,
q : Array[Double],
k : Int,
filter : ((@types.VectorId, @types.Attrs) -> Bool)?,
) -> Array[@types.SearchHit] {
let dim = store.dim
guard q.length() == dim else {
abort(
"dim mismatch: got " +
q.length().to_string() +
", want " +
dim.to_string(),
)
}
// Normalize query for cosine
let query = store.normalize_query(q)
let score_fn = @vecmath.get_score_fn(state.metric)
// Score all centroids
let scores : Array[(Int, Double)] = []
let count = if state.centroid_count < 1 { 1 } else { state.centroid_count }
for c in 0.. continue
Some(at) => {
let attrs = store.attrs[at]
// Apply filter
match filter {
Some(f) => if !f(id, attrs) { continue }
None => ()
}
let base = at * dim
let s = score_fn(store.data, base, query, dim)
push_search_hit_top_k(
out,
@types.SearchHit::{ id, score: s, attrs },
k,
)
}
}
}
}
out
}
///|
/// IVF find one - returns the best match
pub fn ivf_find(
state : IVFState,
store : @store.CoreStore,
q : Array[Double],
filter : ((@types.VectorId, @types.Attrs) -> Bool)?,
) -> @types.SearchHit? {
let results = ivf_search(state, store, q, 1, filter)
first_search_hit(results)
}
///|
/// Serialize IVF state to bytes
pub fn ivf_serialize(state : IVFState, dim : Int) -> Bytes {
let w = @binary.BinaryWriter::new()
// Version marker: 0x49564632 = "IVF2" in ASCII — cannot collide with centroid_count
w.push_u32(0x49564632U)
// Write centroid_count
w.push_i32(state.centroid_count)
// Write centroids (flattened array)
let centroid_len = state.nlist * dim
for i in 0.. Unit {
let r = @binary.BinaryReader::new(data)
// Detect format version: v2 starts with magic 0x49564632 ("IVF2")
let first_u32 = r.read_u32()
let version = if first_u32 == 0x49564632U { 2 } else { 1 }
// Read centroid_count
state.centroid_count = if version >= 2 {
r.read_i32()
} else {
first_u32.reinterpret_as_int()
}
// Read centroids
let centroid_len = state.nlist * dim
for i in 0.. 0 {
let _ = list.pop()
}
}
state.id_to_list.clear()
// Read new lists
for li in 0..= 2 {
let wire_len = r.read_u32().reinterpret_as_int()
let wire = r.read_bytes(wire_len)
match @types.VectorId::from_wire_bytes(wire) {
Some(id) => {
state.lists[li].push(id)
state.id_to_list.set(id, li)
}
None => ()
}
} else {
// v1: fixed 8-byte Int64
let id_val = r.read_u64().reinterpret_as_int64()
let id = @types.Int64Id(id_val)
state.lists[li].push(id)
state.id_to_list.set(id, li)
}
}
}
}
///|
/// Train IVF centroids using k-means++ initialization
pub fn ivf_train(
state : IVFState,
store : @store.CoreStore,
iterations? : Int = 10,
) -> Unit {
if store.count == 0 {
return
}
let dim = store.dim
let k = state.nlist
// K-means++ initialization: select first centroid randomly
// For simplicity, use first k vectors as initial centroids
let centroid_count = if store.count < k { store.count } else { k }
state.centroid_count = centroid_count
// Copy first centroid_count vectors as initial centroids
for c in 0.. 0 {
let _ = list.pop()
}
}
state.id_to_list.clear()
// Assign each vector to nearest centroid
for idx in 0.. continue
Some(idx) => {
let base = idx * dim
for i in 0..