use super::vector_types::{VectorDistanceMetricType, VectorQuantType, VectorSetFlags};
pub const INDEX_SIZE: usize = 56;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Index {
pub context: u64,
pub index_ptr: u64,
pub dimensions: u32,
pub reduce_dims: u32,
pub num_links: u32,
pub build_exploration_factor: u32,
pub quant_type: VectorQuantType,
pub distance_metric: VectorDistanceMetricType,
pub flags: VectorSetFlags,
}
impl Default for Index {
fn default() -> Self {
Self {
context: 0,
index_ptr: 0,
dimensions: 0,
reduce_dims: 0,
num_links: 0,
build_exploration_factor: 0,
quant_type: VectorQuantType::Invalid,
distance_metric: VectorDistanceMetricType::Cosine,
flags: VectorSetFlags::NONE,
}
}
}
impl Index {
pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
if bytes.len() != INDEX_SIZE {
return None;
}
let rd_u64 = |off: usize| u64::from_le_bytes(bytes[off..off + 8].try_into().unwrap());
let rd_u32 = |off: usize| u32::from_le_bytes(bytes[off..off + 4].try_into().unwrap());
let rd_i32 = |off: usize| i32::from_le_bytes(bytes[off..off + 4].try_into().unwrap());
Some(Self {
context: rd_u64(0),
index_ptr: rd_u64(8),
dimensions: rd_u32(16),
reduce_dims: rd_u32(20),
num_links: rd_u32(24),
build_exploration_factor: rd_u32(28),
quant_type: quant_from_i32(rd_i32(32)),
distance_metric: metric_from_i32(rd_i32(36)),
flags: VectorSetFlags::from_bits(bytes[40]),
})
}
pub fn to_bytes(&self) -> [u8; INDEX_SIZE] {
let mut out = [0u8; INDEX_SIZE];
out[0..8].copy_from_slice(&self.context.to_le_bytes());
out[8..16].copy_from_slice(&self.index_ptr.to_le_bytes());
out[16..20].copy_from_slice(&self.dimensions.to_le_bytes());
out[20..24].copy_from_slice(&self.reduce_dims.to_le_bytes());
out[24..28].copy_from_slice(&self.num_links.to_le_bytes());
out[28..32].copy_from_slice(&self.build_exploration_factor.to_le_bytes());
out[32..36].copy_from_slice(&(self.quant_type as i32).to_le_bytes());
out[36..40].copy_from_slice(&(self.distance_metric as i32).to_le_bytes());
out[40] = self.flags.bits();
out
}
}
use super::vector_manager::VectorManager;
impl VectorManager {
pub fn create_index_in(
&self,
key: &[u8],
params: &super::vector_manager__locking::CreateIndexParams,
) -> Result<Index, super::vector_manager::VectorManagerResult> {
let context = self
.next_vector_set_context(0)
.ok_or(super::vector_manager::VectorManagerResult::Invalid)?;
let index = Index {
context,
index_ptr: 1,
dimensions: params.dims,
reduce_dims: params.reduce_dims,
num_links: params.num_links,
build_exploration_factor: params.build_exploration_factor,
quant_type: params.quant,
distance_metric: params.distance_metric,
flags: super::vector_types::VectorSetFlags::NONE,
};
self.service.create_index(
context,
index.dimensions,
index.reduce_dims,
index.quant_type,
index.build_exploration_factor,
index.num_links,
index.distance_metric,
);
self.write_stored_index(key, &index.to_bytes());
Ok(index)
}
pub fn drop_index_record(&self, index_value: &[u8]) {
self.drop_in_memory_index(index_value);
}
pub fn set_context_for_migration(index_value: &mut [u8], new_context: u64) {
debug_assert_ne!(new_context, 0, "0 为特殊上下文,不可指派");
debug_assert_eq!(index_value.len(), INDEX_SIZE, "索引记录尺寸不符");
index_value[0..8].copy_from_slice(&new_context.to_le_bytes());
index_value[8..16].fill(0);
}
pub fn mark_suppress_cleanup(&self, key: &[u8]) {
let mut flags = super::vector_types::VectorSetFlags::NONE;
if let Some(index) = self
.read_stored_index(key)
.and_then(|bytes| Index::from_bytes(&bytes))
{
flags = index.flags;
}
flags = flags.union(super::vector_types::VectorSetFlags::SUPPRESS_CLEANUP);
self.set_flags(key, flags);
}
pub fn clear_suppress_cleanup(&self, key: &[u8]) {
self.set_flags(key, super::vector_types::VectorSetFlags::NONE);
}
pub fn set_flags(&self, key: &[u8], flags: super::vector_types::VectorSetFlags) {
let mut bytes = match self.read_stored_index(key) {
Some(bytes) => bytes,
None => return,
};
bytes[40] = flags.bits();
self.write_stored_index(key, &bytes);
}
pub fn set_index_flags(index_value: &mut [u8], flags: super::vector_types::VectorSetFlags) {
debug_assert_eq!(index_value.len(), INDEX_SIZE, "索引记录尺寸不符");
index_value[40] = flags.bits();
}
}
pub fn quant_from_i32(v: i32) -> VectorQuantType {
match v {
1 => VectorQuantType::NoQuant,
2 => VectorQuantType::Bin,
3 => VectorQuantType::Q8,
4 => VectorQuantType::XNoQuant_U8,
5 => VectorQuantType::XNoQuant_I8,
6 => VectorQuantType::XBin_I8,
7 => VectorQuantType::XBin_U8,
_ => VectorQuantType::Invalid,
}
}
pub fn metric_from_i32(v: i32) -> VectorDistanceMetricType {
match v {
0 => VectorDistanceMetricType::Cosine,
1 => VectorDistanceMetricType::InnerProduct,
2 => VectorDistanceMetricType::L2,
3 => VectorDistanceMetricType::XCosineNormalized,
_ => VectorDistanceMetricType::Cosine,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn sample() -> Index {
Index {
context: 8,
index_ptr: 0xDEAD_BEEF,
dimensions: 32,
reduce_dims: 0,
num_links: 8,
build_exploration_factor: 200,
quant_type: VectorQuantType::Q8,
distance_metric: VectorDistanceMetricType::InnerProduct,
flags: VectorSetFlags::SUPPRESS_CLEANUP,
}
}
#[test]
fn index_layout_roundtrip() {
let idx = sample();
let bytes = idx.to_bytes();
assert_eq!(bytes.len(), INDEX_SIZE);
assert_eq!(u64::from_le_bytes(bytes[0..8].try_into().unwrap()), 8);
assert_eq!(u32::from_le_bytes(bytes[16..20].try_into().unwrap()), 32);
assert_eq!(bytes[40], 1);
let back = Index::from_bytes(&bytes).unwrap();
assert_eq!(back, idx);
}
#[test]
fn read_index_rejects_bad_size() {
assert!(Index::from_bytes(&[0u8; 55]).is_none());
assert!(Index::from_bytes(&[0u8; 57]).is_none());
}
#[test]
fn enum_decode_tables() {
assert_eq!(quant_from_i32(3), VectorQuantType::Q8);
assert_eq!(quant_from_i32(0), VectorQuantType::Invalid);
assert_eq!(quant_from_i32(99), VectorQuantType::Invalid);
assert_eq!(metric_from_i32(1), VectorDistanceMetricType::InnerProduct);
assert_eq!(metric_from_i32(-1), VectorDistanceMetricType::Cosine);
}
}