与存储标量数据(如整数、字符串等)的传统数据库不同,矢量数据库旨在有效地存储和检索矢量数据——表示多维空间中的点的数值集合。
本文将探讨如何在 Rust 中实现基本的向量数据库。
让我们开始吧!🦀
什么是矢量数据库?
矢量数据库是一种针对存储和查询矢量(表示高维空间中的点的数字数组)进行了优化的数据库。这些数据库对于在大型数据集中进行相似性搜索是关键操作的应用至关重要,例如推荐系统、图像检索和自然语言处理。
矢量数据库中的关键概念包括:
- 向量表示:这些数据库中的向量表示数据点。例如,在图像识别中,图像可能表示为高维向量,其中每个维度对应于图像的一个特征。
- 距离度量:为了检索相似的向量,数据库需要一种方法来量化两个向量的“接近”或“相似”程度。常用度量包括欧几里得距离、曼哈顿距离和余弦相似度。
- 索引和搜索算法:在高维空间中进行高效搜索是一个具有挑战性的问题。矢量数据库通常采用专门的索引策略来加快查询时间,例如 KD 树、R 树或基于哈希的方法。
使用 Rust 实现矢量数据库
步骤 1:设置 Rust 环境
在开始编码之前,请确保您已安装 Rust。Rust 的包管理器 Cargo 可让您轻松设置新项目:
cargo new vector_db
cd vector_db
步骤 2:定义向量类型
在 Rust 中,我们可以将向量定义为固定大小的数组,也可以使用标准库中的动态数组。为简单起见,我们使用固定大小的数组 ( [f32; N]
),其中N
是向量的维度:
type Vector = [ f32 ; 3 ]; // 3D 向量示例
步骤3:创建数据库结构
我们将创建一个VectorDB
作为数据库的结构:
struct VectorDB {
vectors: Vec<Vector>,
}
步骤4:实现基本操作
现在,让我们添加添加和检索向量的方法:
impl VectorDB {
fn new() -> Self {
VectorDB { vectors: Vec::new() }
}
fn add_vector(&mut self, vector: Vector) {
self.vectors.push(vector);
}
fn get_vector(&self, index: usize) -> Option<&Vector> {
self.vectors.get(index)
}
}
步骤 5:添加搜索功能
为了找到最接近给定查询向量的向量,我们将实现基于欧几里得距离的简单线性搜索:
impl VectorDB {
// Existing methods...
fn find_closest(&self, query: Vector) -> Option<&Vector> {
self.vectors.iter().min_by(|&a, &b| {
let distance_a = VectorDB::euclidean_distance(&query, a);
let distance_b = VectorDB::euclidean_distance(&query, b);
distance_a.partial_cmp(&distance_b).unwrap()
})
}
fn euclidean_distance(a: &Vector, b: &Vector) -> f32 {
a.iter().zip(b.iter()).map(|(x, y)| (x - y).powi(2)).sum::<f32>().sqrt()
}
}
步骤 6:测试我们的矢量数据库
最后,让我们在main
函数中测试我们的数据库:
fn main() {
let mut db = VectorDB::new();
db.add_vector([1.0, 2.0, 3.0]);
db.add_vector([4.0, 5.0, 6.0]);
// Retrieving and printing a vector
if let Some(vector) = db.get_vector(0) {
println!("Vector at index 0: {:?}", vector);
}
// Finding and printing the closest vector
if let Some(closest) = db.find_closest([2.0, 3.0, 4.0]) {
println!("Closest vector: {:?}", closest);
}
}
实施更高效的索引
我们将重点介绍一种基本但有效的索引策略,即 KD 树(K 维树),这是一种用于组织 K 维空间中的点的空间分区数据结构。KD 树对于多维键中的高效最近邻搜索特别有用。
步骤 1:了解 KD 树
KD 树是一种二叉树,其中每个节点都是一个 K 维点。每个非叶节点都会生成一个分裂超平面,将空间分成两半。此超平面左侧的点由左子树表示,而右侧的点由右子树表示。
关键概念:
- 分割维度:在树的每一层,选择不同的维度来分割数据。维度的选择通常会循环遍历所有维度。
- 中位数查找:要分割每个节点上的点,请选择所选维度上的中位数。这种方法有助于平衡树。
- 递归分区:该过程以递归方式继续,最终形成一棵树,其中每个叶节点代表空间中的一个点。
步骤 2:在 Rust 中定义 KD 树结构
首先,定义 KD-Tree 的结构。您需要表示内部节点(包含拆分信息)和叶节点(包含实际矢量数据):
enum KDTreeNode {
Leaf(Vector),
Internal {
left: Box<KDTreeNode>,
right: Box<KDTreeNode>,
split_value: f32,
split_dimension: usize,
},
}
struct KDTree {
root: KDTreeNode,
}
步骤 3:构建 KD 树
该过程包括根据分裂维度对点进行排序并递归构建树:
impl KDTree {
fn build(points: Vec<Vector>, depth: usize) -> KDTreeNode {
if points.len() == 1 {
return KDTreeNode::Leaf(points[0]);
}
let dim = depth % K; // K is the number of dimensions
let mut sorted_points = points;
sorted_points.sort_by(|a, b| a[dim].partial_cmp(&b[dim]).unwrap());
let median_idx = sorted_points.len() / 2;
let median_value = sorted_points[median_idx][dim];
KDTreeNode::Internal {
left: Box::new(KDTree::build(sorted_points[..median_idx].to_vec(), depth + 1)),
right: Box::new(KDTree::build(sorted_points[median_idx..].to_vec(), depth + 1)),
split_value: median_value,
split_dimension: dim,
}
}
}
步骤 4:在 KD 树中实现搜索
寻找最近的邻居涉及遍历树、检查距离,并可能进行回溯:
impl KDTree {
fn nearest_neighbor(&self, query: &Vector) -> Option<&Vector> {
self.nearest(query, &self.root, None, f32::MAX)
}
fn nearest<'a>(&'a self, query: &Vector, node: &'a KDTreeNode, best: Option<&'a Vector>, best_dist: f32) -> Option<&'a Vector> {
match node {
KDTreeNode::Leaf(point) => {
let dist = VectorDB::euclidean_distance(query, point);
if dist < best_dist {
Some(point)
} else {
best
}
},
KDTreeNode::Internal { left, right, split_value, split_dimension } => {
let next_node = if query[*split_dimension] < *split_value { left } else { right };
let other_node = if query[*split_dimension] < *split_value { right } else { left };
let updated_best = self.nearest(query, next_node, best, best_dist);
let updated_best_dist = updated_best.map_or(best_dist, |v| VectorDB::euclidean_distance(query, v));
if (query[*split_dimension] - split_value).abs() < updated_best_dist {
self.nearest(query, other_node, updated_best, updated_best_dist)
} else {
updated_best
}
}
}
}
}
步骤 5:测试 KD 树
最后,将KD-Tree集成到您的矢量数据库中并进行测试:
fn main ( ) {
let points = vec![[ 1.0 , 2.0 , 3.0 ], [ 4.0 , 5.0 , 6.0 ], [ 2.0 , 3.0 , 4.0 ], ...];
let kd_tree = KDTree :: build (points, 0 );
If let Some (nearest) = kd_tree.nearest_neighbor ( &[ 3.0 , 3.0 , 3.0 ] ) { println!( "最近邻居: {:?}" , nearest); } }
就这些了
RA/SD 衍生者AI训练营。发布者:chris,转载请注明出处:https://www.shxcj.com/archives/4047