dynamic_learned_index 0.1.0

Dynamic Learned Index for efficient search in unstructured data.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

use log::debug;
use measure_time_macro::log_time;
use ndarray::{Array2, ArrayView};

use crate::structs::LabelMethod;

#[log_time]
pub(crate) fn compute_labels(
    data: &[f32],
    label_method: &LabelMethod,
    k: usize,
    input_shape: usize,
    max_iters: usize,
) -> Vec<i64> {
    debug_assert!(!data.is_empty());
    let data_len = data.len() / input_shape;
    assert!(data_len * input_shape == data.len());
    debug!(data_len = data_len, k = k; "clustering:compute_labels");
    match label_method {
        LabelMethod::KMeans => k_means_clustering(data, input_shape, k, max_iters),
        LabelMethod::SphericalKMeans => {
            k_means_clustering_spherical(data, input_shape, k, max_iters)
        }
    }
}

fn k_means_clustering(data: &[f32], input_shape: usize, k: usize, max_iters: usize) -> Vec<i64> {
    let count = data.len() / input_shape;
    assert!(count * input_shape == data.len());
    assert!(k > 0);
    assert!(k <= count);
    assert!(max_iters > 0);
    let mut kmeans = kentro::KMeans::new(k)
        .with_iterations(max_iters)
        .with_euclidean(true);

    // Convert the Vec<f32> to Array2 with shape (count, input_shape)
    let data = ArrayView::from_shape((count, input_shape), data)
        .expect("Failed to reshape data into Array2");
    let clusters = kmeans.train(data.view(), None).unwrap();
    let mut result = vec![0; count];
    clusters
        .into_iter()
        .enumerate()
        .for_each(|(cluster, query_ids)| {
            query_ids.into_iter().for_each(|query_id| {
                assert!(query_id < count);
                result[query_id] = cluster as i64;
            });
        });
    // debug!(error = result.distsum ;"kmeans:metrics");
    assert!(result.len() == count);
    result
}

fn k_means_clustering_spherical(
    data: &[f32],
    input_shape: usize,
    k: usize,
    max_iters: usize,
) -> Vec<i64> {
    let count = data.len() / input_shape;
    assert!(count * input_shape == data.len());
    assert!(k > 0);
    assert!(k <= count);
    assert!(max_iters > 0);
    let mut kmeans = kentro::KMeans::new(k).with_iterations(max_iters);

    // Convert the Vec<f32> to Array2 with shape (count, input_shape)
    let data = Array2::from_shape_vec((count, input_shape), data.to_vec())
        .expect("Failed to reshape data into Array2");
    let clusters = kmeans.train(data.view(), None).unwrap();
    let mut result = vec![0; count];
    clusters
        .into_iter()
        .enumerate()
        .for_each(|(cluster, query_ids)| {
            query_ids.into_iter().for_each(|query_id| {
                assert!(query_id < count);
                result[query_id] = cluster as i64;
            });
        });
    // debug!(error = result.distsum ;"kmeans:metrics");
    assert!(result.len() == count);
    result
}