Added from_f32_slice function to HalfVector

ankane · ankane · commit 0012b629f925 · 2025-02-13T23:03:33.000-08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.4.1 (unreleased)
+
+- Added `from_f32_slice` function to `HalfVector`
+
 ## 0.4.0 (2024-07-28)
 
 - Added support for SQLx 0.8
diff --git a/README.md b/README.md
@@ -287,7 +287,7 @@ let slice = vec.as_slice();
 
 Note: Use the `halfvec` feature to enable half vectors
 
-Create a half vector
+Create a half vector from a `Vec<f16>`
 
 ```rust
 use half::f16;
@@ -296,6 +296,12 @@ use pgvector::HalfVector;
 let vec = HalfVector::from(vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]);
 ```
 
+or a `f32` slice [unreleased]
+
+```rust
+let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
+```
+
 Convert to a `Vec<f16>`
 
 ```rust
diff --git a/src/diesel_ext/halfvec.rs b/src/diesel_ext/halfvec.rs
@@ -36,7 +36,6 @@ impl FromSql<HalfVectorType, Pg> for HalfVector {
 mod tests {
     use crate::{HalfVector, VectorExpressionMethods};
     use diesel::prelude::*;
-    use half::f16;
 
     table! {
         use diesel::sql_types::*;
@@ -74,25 +73,13 @@ mod tests {
 
         let new_items = vec![
             NewItem {
-                embedding: Some(HalfVector::from(vec![
-                    f16::from_f32(1.0),
-                    f16::from_f32(1.0),
-                    f16::from_f32(1.0),
-                ])),
+                embedding: Some(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])),
             },
             NewItem {
-                embedding: Some(HalfVector::from(vec![
-                    f16::from_f32(2.0),
-                    f16::from_f32(2.0),
-                    f16::from_f32(2.0),
-                ])),
+                embedding: Some(HalfVector::from_f32_slice(&[2.0, 2.0, 2.0])),
             },
             NewItem {
-                embedding: Some(HalfVector::from(vec![
-                    f16::from_f32(1.0),
-                    f16::from_f32(1.0),
-                    f16::from_f32(2.0),
-                ])),
+                embedding: Some(HalfVector::from_f32_slice(&[1.0, 1.0, 2.0])),
             },
             NewItem { embedding: None },
         ];
@@ -105,32 +92,20 @@ mod tests {
         assert_eq!(4, all.len());
 
         let neighbors = items::table
-            .order(items::embedding.l2_distance(HalfVector::from(vec![
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-            ])))
+            .order(items::embedding.l2_distance(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])))
             .limit(5)
             .load::<Item>(&mut conn)?;
         assert_eq!(
             vec![1, 3, 2, 4],
             neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
         );
         assert_eq!(
-            Some(HalfVector::from(vec![
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-                f16::from_f32(1.0)
-            ])),
+            Some(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])),
             neighbors.first().unwrap().embedding
         );
 
         let neighbors = items::table
-            .order(items::embedding.max_inner_product(HalfVector::from(vec![
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-            ])))
+            .order(items::embedding.max_inner_product(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])))
             .limit(5)
             .load::<Item>(&mut conn)?;
         assert_eq!(
@@ -139,11 +114,7 @@ mod tests {
         );
 
         let neighbors = items::table
-            .order(items::embedding.cosine_distance(HalfVector::from(vec![
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-            ])))
+            .order(items::embedding.cosine_distance(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])))
             .limit(5)
             .load::<Item>(&mut conn)?;
         assert_eq!(
@@ -152,11 +123,7 @@ mod tests {
         );
 
         let neighbors = items::table
-            .order(items::embedding.l1_distance(HalfVector::from(vec![
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-            ])))
+            .order(items::embedding.l1_distance(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])))
             .limit(5)
             .load::<Item>(&mut conn)?;
         assert_eq!(
@@ -165,11 +132,9 @@ mod tests {
         );
 
         let distances = items::table
-            .select(items::embedding.max_inner_product(HalfVector::from(vec![
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-                f16::from_f32(1.0),
-            ])))
+            .select(
+                items::embedding.max_inner_product(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])),
+            )
             .order(items::id)
             .load::<Option<f64>>(&mut conn)?;
         assert_eq!(vec![Some(-3.0), Some(-6.0), Some(-4.0), None], distances);
diff --git a/src/halfvec.rs b/src/halfvec.rs
@@ -25,6 +25,11 @@ impl From<HalfVector> for Vec<f16> {
 }
 
 impl HalfVector {
+    /// Creates a half vector from a `f32` slice.
+    pub fn from_f32_slice(slice: &[f32]) -> HalfVector {
+        HalfVector(slice.iter().map(|v| f16::from_f32(*v)).collect())
+    }
+
     /// Returns a copy of the half vector as a `Vec<f16>`.
     pub fn to_vec(&self) -> Vec<f16> {
         self.0.clone()
@@ -76,11 +81,7 @@ mod tests {
 
     #[test]
     fn test_to_vec() {
-        let vec = HalfVector::from(vec![
-            f16::from_f32(1.0),
-            f16::from_f32(2.0),
-            f16::from_f32(3.0),
-        ]);
+        let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
         assert_eq!(
             vec.to_vec(),
             vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]
@@ -89,11 +90,7 @@ mod tests {
 
     #[test]
     fn test_as_slice() {
-        let vec = HalfVector::from(vec![
-            f16::from_f32(1.0),
-            f16::from_f32(2.0),
-            f16::from_f32(3.0),
-        ]);
+        let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
         assert_eq!(
             vec.as_slice(),
             &[f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]
diff --git a/src/postgres_ext/halfvec.rs b/src/postgres_ext/halfvec.rs
@@ -59,26 +59,14 @@ mod tests {
             &[],
         )?;
 
-        let vec = HalfVector::from(vec![
-            f16::from_f32(1.0),
-            f16::from_f32(2.0),
-            f16::from_f32(3.0),
-        ]);
-        let vec2 = HalfVector::from(vec![
-            f16::from_f32(4.0),
-            f16::from_f32(5.0),
-            f16::from_f32(6.0),
-        ]);
+        let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
+        let vec2 = HalfVector::from_f32_slice(&[4.0, 5.0, 6.0]);
         client.execute(
             "INSERT INTO postgres_half_items (embedding) VALUES ($1), ($2), (NULL)",
             &[&vec, &vec2],
         )?;
 
-        let query_vec = HalfVector::from(vec![
-            f16::from_f32(3.0),
-            f16::from_f32(1.0),
-            f16::from_f32(2.0),
-        ]);
+        let query_vec = HalfVector::from_f32_slice(&[3.0, 1.0, 2.0]);
         let row = client.query_one(
             "SELECT embedding FROM postgres_half_items ORDER BY embedding <-> $1 LIMIT 1",
             &[&query_vec],
diff --git a/src/sqlx_ext/halfvec.rs b/src/sqlx_ext/halfvec.rs
@@ -65,27 +65,15 @@ mod tests {
         .execute(&pool)
         .await?;
 
-        let vec = HalfVector::from(vec![
-            f16::from_f32(1.0),
-            f16::from_f32(2.0),
-            f16::from_f32(3.0),
-        ]);
-        let vec2 = HalfVector::from(vec![
-            f16::from_f32(4.0),
-            f16::from_f32(5.0),
-            f16::from_f32(6.0),
-        ]);
+        let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
+        let vec2 = HalfVector::from_f32_slice(&[4.0, 5.0, 6.0]);
         sqlx::query("INSERT INTO sqlx_half_items (embedding) VALUES ($1), ($2), (NULL)")
             .bind(&vec)
             .bind(&vec2)
             .execute(&pool)
             .await?;
 
-        let query_vec = HalfVector::from(vec![
-            f16::from_f32(3.0),
-            f16::from_f32(1.0),
-            f16::from_f32(2.0),
-        ]);
+        let query_vec = HalfVector::from_f32_slice(&[3.0, 1.0, 2.0]);
         let row =
             sqlx::query("SELECT embedding FROM sqlx_half_items ORDER BY embedding <-> $1 LIMIT 1")
                 .bind(query_vec)
