nabled 0.0.11

A linear algebra library written in Rust with an ndarray-first architecture.
Documentation 
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171

#![cfg(all(feature = "arrow", feature = "embeddings"))]

use approx::assert_abs_diff_eq;
use arrow_array::types::{Float32Type, Float64Type};
use arrow_array::{Array, FixedSizeListArray, Float32Array, Float64Array, Int64Array, StructArray};
use nabled::arrow::embeddings;
use nabled::embeddings::{Metric, brute_force_knn, normalize_rows, query_corpus_scores, rerank};
use nabled::ndarrow::{IntoArrow, fixed_size_list_as_array2};
use ndarray::{Array1, array};

fn struct_index(structs: &StructArray) -> Vec<i64> {
    structs
        .column_by_name("index")
        .unwrap()
        .as_any()
        .downcast_ref::<Int64Array>()
        .unwrap()
        .values()
        .to_vec()
}

fn struct_scores_f64(structs: &StructArray) -> Vec<f64> {
    structs
        .column_by_name("score")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap()
        .values()
        .to_vec()
}

#[test]
fn arrow_query_corpus_scores_matches_ndarray_path() {
    let queries_nd = array![[1.0_f64, 0.0, 0.0], [0.0, 1.0, 0.0]];
    let corpus_nd = array![[1.0_f64, 0.0, 0.0], [0.0, 1.0, 0.0], [1.0, 1.0, 0.0]];

    let queries = queries_nd.clone().into_arrow().unwrap();
    let corpus = corpus_nd.clone().into_arrow().unwrap();

    for metric in [Metric::Cosine, Metric::Dot, Metric::L2] {
        let expected = query_corpus_scores(&queries_nd, &corpus_nd, metric).unwrap();
        let arrow_scores =
            embeddings::arrow_query_corpus_scores::<Float64Type>(&queries, &corpus, metric)
                .unwrap();
        let view = fixed_size_list_as_array2::<Float64Type>(&arrow_scores).unwrap();
        assert_eq!(view.dim(), expected.dim());
        for (lhs, rhs) in view.iter().zip(expected.iter()) {
            assert_abs_diff_eq!(lhs, rhs, epsilon = 1.0e-12);
        }
    }
}

#[test]
fn arrow_query_corpus_scores_f32_parity() {
    let queries_nd = array![[1.0_f32, 0.0, 0.0]];
    let corpus_nd = array![[1.0_f32, 0.0, 0.0], [0.0, 1.0, 0.0]];

    let queries = queries_nd.clone().into_arrow().unwrap();
    let corpus = corpus_nd.clone().into_arrow().unwrap();

    let expected = query_corpus_scores(&queries_nd, &corpus_nd, Metric::Cosine).unwrap();
    let arrow_scores =
        embeddings::arrow_query_corpus_scores::<Float32Type>(&queries, &corpus, Metric::Cosine)
            .unwrap();
    let view = fixed_size_list_as_array2::<Float32Type>(&arrow_scores).unwrap();
    for (lhs, rhs) in view.iter().zip(expected.iter()) {
        assert_abs_diff_eq!(lhs, rhs, epsilon = 1.0e-6);
    }
}

#[test]
fn arrow_normalize_rows_matches_ndarray_path() {
    let rows_nd = array![[3.0_f64, 4.0], [0.0, 2.0], [1.0, 1.0]];
    let rows = rows_nd.clone().into_arrow().unwrap();

    let expected = normalize_rows(&rows_nd).unwrap();
    let arrow_normalized = embeddings::arrow_normalize_rows::<Float64Type>(&rows).unwrap();
    let view = fixed_size_list_as_array2::<Float64Type>(&arrow_normalized).unwrap();
    assert_eq!(view.dim(), expected.dim());
    for (lhs, rhs) in view.iter().zip(expected.iter()) {
        assert_abs_diff_eq!(lhs, rhs, epsilon = 1.0e-12);
    }
}

#[test]
fn arrow_normalize_rows_f32_parity() {
    let rows_nd = array![[3.0_f32, 4.0]];
    let rows = rows_nd.clone().into_arrow().unwrap();

    let expected = normalize_rows(&rows_nd).unwrap();
    let arrow_normalized = embeddings::arrow_normalize_rows::<Float32Type>(&rows).unwrap();
    let view = fixed_size_list_as_array2::<Float32Type>(&arrow_normalized).unwrap();
    for (lhs, rhs) in view.iter().zip(expected.iter()) {
        assert_abs_diff_eq!(lhs, rhs, epsilon = 1.0e-6);
    }
}

#[test]
fn arrow_rerank_matches_ndarray_path() {
    let candidates_nd = array![[1.0_f64, 0.0], [0.0, 1.0], [0.9, 0.1]];
    let candidates = candidates_nd.clone().into_arrow().unwrap();
    let query_nd = Array1::from(vec![1.0_f64, 0.0]);
    let query = Float64Array::from(vec![1.0_f64, 0.0]);

    let expected = rerank(&query_nd.view(), &candidates_nd.view(), 2, Metric::Cosine).unwrap();
    let arrow_neighbors =
        embeddings::arrow_rerank::<Float64Type>(&query, &candidates, 2, Metric::Cosine).unwrap();

    assert_eq!(arrow_neighbors.len(), expected.len());
    let indices = struct_index(&arrow_neighbors);
    let scores = struct_scores_f64(&arrow_neighbors);
    for (i, neighbor) in expected.iter().enumerate() {
        assert_eq!(indices[i], i64::try_from(neighbor.index).unwrap());
        assert_abs_diff_eq!(scores[i], neighbor.score, epsilon = 1.0e-12);
    }
}

#[test]
fn arrow_rerank_f32_parity() {
    let candidates_nd = array![[1.0_f32, 0.0], [0.0, 1.0]];
    let candidates = candidates_nd.clone().into_arrow().unwrap();
    let query_nd = Array1::from(vec![1.0_f32, 0.0]);
    let query = Float32Array::from(vec![1.0_f32, 0.0]);

    let expected = rerank(&query_nd.view(), &candidates_nd.view(), 2, Metric::Cosine).unwrap();
    let arrow_neighbors =
        embeddings::arrow_rerank::<Float32Type>(&query, &candidates, 2, Metric::Cosine).unwrap();

    let indices = struct_index(&arrow_neighbors);
    assert_eq!(indices[0], i64::try_from(expected[0].index).unwrap());
}

#[test]
fn arrow_brute_force_knn_matches_ndarray_path() {
    let queries_nd = array![[1.0_f64, 0.0], [0.0, 1.0]];
    let corpus_nd = array![[1.0_f64, 0.0], [0.0, 1.0], [0.9, 0.1]];
    let queries = queries_nd.clone().into_arrow().unwrap();
    let corpus = corpus_nd.clone().into_arrow().unwrap();

    let expected =
        brute_force_knn(&queries_nd.view(), &corpus_nd.view(), 2, Metric::Cosine).unwrap();
    let arrow_lists =
        embeddings::arrow_brute_force_knn::<Float64Type>(&queries, &corpus, 2, Metric::Cosine)
            .unwrap();

    assert_eq!(arrow_lists.len(), expected.len());
    for (q, expected_neighbors) in expected.iter().enumerate() {
        let row = arrow_lists.value(q);
        let structs = row.as_any().downcast_ref::<StructArray>().unwrap();
        let indices = struct_index(structs);
        let scores = struct_scores_f64(structs);
        assert_eq!(indices.len(), expected_neighbors.len());
        for (i, neighbor) in expected_neighbors.iter().enumerate() {
            assert_eq!(indices[i], i64::try_from(neighbor.index).unwrap());
            assert_abs_diff_eq!(scores[i], neighbor.score, epsilon = 1.0e-12);
        }
    }
}

#[test]
fn arrow_embedding_inputs_bridge_zero_copy() {
    let rows_nd = array![[3.0_f64, 4.0], [1.0, 2.0]];
    let rows: FixedSizeListArray = rows_nd.into_arrow().unwrap();

    let view = fixed_size_list_as_array2::<Float64Type>(&rows).unwrap();
    let values = rows.values().as_any().downcast_ref::<Float64Array>().unwrap();

    // A zero-copy bridge shares the same backing buffer pointer as the Arrow values array.
    assert_eq!(view.as_ptr(), values.values().as_ptr());
}