Skip to main content

oxirs_vec/
tree_indices_types.rs

1//! Shared configuration and helper types for tree-based indices.
2//!
3//! Defines [`TreeIndexConfig`], [`TreeType`], [`DistanceMetric`], and the
4//! crate-internal `SearchResult` heap entry used by every tree implementation.
5
6use oxirs_core::simd::SimdOps;
7use std::cmp::Ordering;
8
9/// Configuration for tree-based indices
10#[derive(Debug, Clone)]
11pub struct TreeIndexConfig {
12    /// Type of tree to use
13    pub tree_type: TreeType,
14    /// Maximum leaf size before splitting
15    pub max_leaf_size: usize,
16    /// Random seed for reproducibility
17    pub random_seed: Option<u64>,
18    /// Enable parallel construction
19    pub parallel_construction: bool,
20    /// Distance metric
21    pub distance_metric: DistanceMetric,
22}
23
24impl Default for TreeIndexConfig {
25    fn default() -> Self {
26        Self {
27            tree_type: TreeType::BallTree,
28            max_leaf_size: 16, // Larger leaf size to prevent deep recursion and stack overflow
29            random_seed: None,
30            parallel_construction: true,
31            distance_metric: DistanceMetric::Euclidean,
32        }
33    }
34}
35
36/// Available tree types
37#[derive(Debug, Clone, Copy)]
38pub enum TreeType {
39    BallTree,
40    KdTree,
41    VpTree,
42    CoverTree,
43    RandomProjectionTree,
44}
45
46/// Distance metrics
47#[derive(Debug, Clone, Copy)]
48pub enum DistanceMetric {
49    Euclidean,
50    Manhattan,
51    Cosine,
52    Minkowski(f32),
53}
54
55impl DistanceMetric {
56    pub(crate) fn distance(&self, a: &[f32], b: &[f32]) -> f32 {
57        match self {
58            DistanceMetric::Euclidean => f32::euclidean_distance(a, b),
59            DistanceMetric::Manhattan => f32::manhattan_distance(a, b),
60            DistanceMetric::Cosine => f32::cosine_distance(a, b),
61            DistanceMetric::Minkowski(p) => a
62                .iter()
63                .zip(b.iter())
64                .map(|(x, y)| (x - y).abs().powf(*p))
65                .sum::<f32>()
66                .powf(1.0 / p),
67        }
68    }
69}
70
71/// Search result with distance
72#[derive(Debug, Clone)]
73pub(crate) struct SearchResult {
74    pub(crate) index: usize,
75    pub(crate) distance: f32,
76}
77
78impl PartialEq for SearchResult {
79    fn eq(&self, other: &Self) -> bool {
80        self.distance == other.distance
81    }
82}
83
84impl Eq for SearchResult {}
85
86impl PartialOrd for SearchResult {
87    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
88        Some(self.cmp(other))
89    }
90}
91
92impl Ord for SearchResult {
93    fn cmp(&self, other: &Self) -> Ordering {
94        self.partial_cmp(other).unwrap_or(Ordering::Equal)
95    }
96}