Added test for binary copy

pgvector · Mar 6, 2024 · 435d4f4 · 435d4f4
1 parent b183250
commit 435d4f4
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/src/postgres_ext.rs b/src/postgres_ext.rs
@@ -38,6 +38,7 @@ impl ToSql for Vector {
 #[cfg(test)]
 mod tests {
     use crate::Vector;
+    use postgres::binary_copy::BinaryCopyInWriter;
     use postgres::{Client, NoTls};
 
     #[test]
@@ -98,6 +99,15 @@ mod tests {
         let text_res: String = text_row.get(0);
         assert_eq!("[1,2,3]", text_res);
 
+        // copy
+        let vector_type = row.columns()[0].type_().clone();
+        let writer =
+            client.copy_in("COPY postgres_items (embedding) FROM STDIN WITH (FORMAT BINARY)")?;
+        let mut writer = BinaryCopyInWriter::new(writer, &[vector_type]);
+        writer.write(&[&Vector::from(vec![1.0, 2.0, 3.0])]).unwrap();
+        writer.write(&[&Vector::from(vec![4.0, 5.0, 6.0])]).unwrap();
+        writer.finish()?;
+
         Ok(())
     }
 }