Skip to content

Commit 3db3994

Browse files
ankaneankane
committed
Added from_slice function to HalfVector
1 parent 84103ca commit 3db3994

File tree

6 files changed

+55
-87
lines changed

6 files changed

+55
-87
lines changed

‎CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
## 0.4.1 (unreleased)
2+
3+
- Added `from_slice` function to `HalfVector`
4+
15
## 0.4.0 (2024-07-28)
26

37
- Added support for SQLx 0.8

‎README.md‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -287,7 +287,7 @@ let slice = vec.as_slice();
287287

288288
Note: Use the `halfvec` feature to enable half vectors
289289

290-
Create a half vector
290+
Create a half vector from a `Vec<f16>`
291291

292292
```rust
293293
use half::f16;
@@ -296,6 +296,12 @@ use pgvector::HalfVector;
296296
let vec = HalfVector::from(vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]);
297297
```
298298

299+
or a slice [unreleased]
300+
301+
```rust
302+
let vec = HalfVector::from_slice(&[1.0, 2.0, 3.0]);
303+
```
304+
299305
Convert to a `Vec<f16>`
300306

301307
```rust

‎src/diesel_ext/halfvec.rs‎

Lines changed: 9 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@ impl FromSql<HalfVectorType, Pg> for HalfVector {
3636
mod tests {
3737
use crate::{HalfVector, VectorExpressionMethods};
3838
use diesel::prelude::*;
39-
use half::f16;
4039

4140
table! {
4241
use diesel::sql_types::*;
@@ -74,25 +73,13 @@ mod tests {
7473

7574
let new_items = vec![
7675
NewItem {
77-
embedding: Some(HalfVector::from(vec![
78-
f16::from_f32(1.0),
79-
f16::from_f32(1.0),
80-
f16::from_f32(1.0),
81-
])),
76+
embedding: Some(HalfVector::from_slice(&[1.0, 1.0, 1.0])),
8277
},
8378
NewItem {
84-
embedding: Some(HalfVector::from(vec![
85-
f16::from_f32(2.0),
86-
f16::from_f32(2.0),
87-
f16::from_f32(2.0),
88-
])),
79+
embedding: Some(HalfVector::from_slice(&[2.0, 2.0, 2.0])),
8980
},
9081
NewItem {
91-
embedding: Some(HalfVector::from(vec![
92-
f16::from_f32(1.0),
93-
f16::from_f32(1.0),
94-
f16::from_f32(2.0),
95-
])),
82+
embedding: Some(HalfVector::from_slice(&[1.0, 1.0, 2.0])),
9683
},
9784
NewItem { embedding: None },
9885
];
@@ -105,32 +92,20 @@ mod tests {
10592
assert_eq!(4, all.len());
10693

10794
let neighbors = items::table
108-
.order(items::embedding.l2_distance(HalfVector::from(vec![
109-
f16::from_f32(1.0),
110-
f16::from_f32(1.0),
111-
f16::from_f32(1.0),
112-
])))
95+
.order(items::embedding.l2_distance(HalfVector::from_slice(&[1.0, 1.0, 1.0])))
11396
.limit(5)
11497
.load::<Item>(&mut conn)?;
11598
assert_eq!(
11699
vec![1, 3, 2, 4],
117100
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
118101
);
119102
assert_eq!(
120-
Some(HalfVector::from(vec![
121-
f16::from_f32(1.0),
122-
f16::from_f32(1.0),
123-
f16::from_f32(1.0)
124-
])),
103+
Some(HalfVector::from_slice(&[1.0, 1.0, 1.0])),
125104
neighbors.first().unwrap().embedding
126105
);
127106

128107
let neighbors = items::table
129-
.order(items::embedding.max_inner_product(HalfVector::from(vec![
130-
f16::from_f32(1.0),
131-
f16::from_f32(1.0),
132-
f16::from_f32(1.0),
133-
])))
108+
.order(items::embedding.max_inner_product(HalfVector::from_slice(&[1.0, 1.0, 1.0])))
134109
.limit(5)
135110
.load::<Item>(&mut conn)?;
136111
assert_eq!(
@@ -139,11 +114,7 @@ mod tests {
139114
);
140115

141116
let neighbors = items::table
142-
.order(items::embedding.cosine_distance(HalfVector::from(vec![
143-
f16::from_f32(1.0),
144-
f16::from_f32(1.0),
145-
f16::from_f32(1.0),
146-
])))
117+
.order(items::embedding.cosine_distance(HalfVector::from_slice(&[1.0, 1.0, 1.0])))
147118
.limit(5)
148119
.load::<Item>(&mut conn)?;
149120
assert_eq!(
@@ -152,11 +123,7 @@ mod tests {
152123
);
153124

154125
let neighbors = items::table
155-
.order(items::embedding.l1_distance(HalfVector::from(vec![
156-
f16::from_f32(1.0),
157-
f16::from_f32(1.0),
158-
f16::from_f32(1.0),
159-
])))
126+
.order(items::embedding.l1_distance(HalfVector::from_slice(&[1.0, 1.0, 1.0])))
160127
.limit(5)
161128
.load::<Item>(&mut conn)?;
162129
assert_eq!(
@@ -165,11 +132,7 @@ mod tests {
165132
);
166133

167134
let distances = items::table
168-
.select(items::embedding.max_inner_product(HalfVector::from(vec![
169-
f16::from_f32(1.0),
170-
f16::from_f32(1.0),
171-
f16::from_f32(1.0),
172-
])))
135+
.select(items::embedding.max_inner_product(HalfVector::from_slice(&[1.0, 1.0, 1.0])))
173136
.order(items::id)
174137
.load::<Option<f64>>(&mut conn)?;
175138
assert_eq!(vec![Some(-3.0), Some(-6.0), Some(-4.0), None], distances);

‎src/halfvec.rs‎

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,34 @@ impl From<HalfVector> for Vec<f16> {
2424
}
2525
}
2626

27+
pub trait AsPrimitive {
28+
fn as_(&self) -> f16;
29+
}
30+
31+
impl AsPrimitive for f16 {
32+
fn as_(&self) -> f16 {
33+
*self
34+
}
35+
}
36+
37+
impl AsPrimitive for f32 {
38+
fn as_(&self) -> f16 {
39+
f16::from_f32(*self)
40+
}
41+
}
42+
43+
impl AsPrimitive for f64 {
44+
fn as_(&self) -> f16 {
45+
f16::from_f64(*self)
46+
}
47+
}
48+
2749
impl HalfVector {
50+
/// Creates a half vector from a slice.
51+
pub fn from_slice<T: AsPrimitive>(slice: &[T]) -> HalfVector {
52+
HalfVector(slice.iter().map(|v| (*v).as_()).collect())
53+
}
54+
2855
/// Returns a copy of the half vector as a `Vec<f16>`.
2956
pub fn to_vec(&self) -> Vec<f16> {
3057
self.0.clone()
@@ -76,11 +103,7 @@ mod tests {
76103

77104
#[test]
78105
fn test_to_vec() {
79-
let vec = HalfVector::from(vec![
80-
f16::from_f32(1.0),
81-
f16::from_f32(2.0),
82-
f16::from_f32(3.0),
83-
]);
106+
let vec = HalfVector::from_slice(&[1.0, 2.0, 3.0]);
84107
assert_eq!(
85108
vec.to_vec(),
86109
vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]
@@ -89,11 +112,7 @@ mod tests {
89112

90113
#[test]
91114
fn test_as_slice() {
92-
let vec = HalfVector::from(vec![
93-
f16::from_f32(1.0),
94-
f16::from_f32(2.0),
95-
f16::from_f32(3.0),
96-
]);
115+
let vec = HalfVector::from_slice(&[1.0_f32, 2.0, 3.0]);
97116
assert_eq!(
98117
vec.as_slice(),
99118
&[f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]

‎src/postgres_ext/halfvec.rs‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -59,26 +59,14 @@ mod tests {
5959
&[],
6060
)?;
6161

62-
let vec = HalfVector::from(vec![
63-
f16::from_f32(1.0),
64-
f16::from_f32(2.0),
65-
f16::from_f32(3.0),
66-
]);
67-
let vec2 = HalfVector::from(vec![
68-
f16::from_f32(4.0),
69-
f16::from_f32(5.0),
70-
f16::from_f32(6.0),
71-
]);
62+
let vec = HalfVector::from_slice(&[1.0, 2.0, 3.0]);
63+
let vec2 = HalfVector::from_slice(&[4.0, 5.0, 6.0]);
7264
client.execute(
7365
"INSERT INTO postgres_half_items (embedding) VALUES ($1), ($2), (NULL)",
7466
&[&vec, &vec2],
7567
)?;
7668

77-
let query_vec = HalfVector::from(vec![
78-
f16::from_f32(3.0),
79-
f16::from_f32(1.0),
80-
f16::from_f32(2.0),
81-
]);
69+
let query_vec = HalfVector::from_slice(&[3.0, 1.0, 2.0]);
8270
let row = client.query_one(
8371
"SELECT embedding FROM postgres_half_items ORDER BY embedding <-> $1 LIMIT 1",
8472
&[&query_vec],

‎src/sqlx_ext/halfvec.rs‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -65,27 +65,15 @@ mod tests {
6565
.execute(&pool)
6666
.await?;
6767

68-
let vec = HalfVector::from(vec![
69-
f16::from_f32(1.0),
70-
f16::from_f32(2.0),
71-
f16::from_f32(3.0),
72-
]);
73-
let vec2 = HalfVector::from(vec![
74-
f16::from_f32(4.0),
75-
f16::from_f32(5.0),
76-
f16::from_f32(6.0),
77-
]);
68+
let vec = HalfVector::from_slice(&[1.0, 2.0, 3.0]);
69+
let vec2 = HalfVector::from_slice(&[4.0, 5.0, 6.0]);
7870
sqlx::query("INSERT INTO sqlx_half_items (embedding) VALUES ($1), ($2), (NULL)")
7971
.bind(&vec)
8072
.bind(&vec2)
8173
.execute(&pool)
8274
.await?;
8375

84-
let query_vec = HalfVector::from(vec![
85-
f16::from_f32(3.0),
86-
f16::from_f32(1.0),
87-
f16::from_f32(2.0),
88-
]);
76+
let query_vec = HalfVector::from_slice(&[3.0, 1.0, 2.0]);
8977
let row =
9078
sqlx::query("SELECT embedding FROM sqlx_half_items ORDER BY embedding <-> $1 LIMIT 1")
9179
.bind(query_vec)

0 commit comments

Comments
 (0)