Added support for binary representation - closes #5

ankane · davecramer · ankane · commit 0b21f088bc72 · 2023-07-15T11:56:23.000-07:00
Co-authored-by: Dave Cramer &lt;davecramer@gmail.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.1.3 (unreleased)
+
+- Added support for binary representation
+
 ## 0.1.2 (2023-05-08)
 
 - Added support for Java 11 and 17
diff --git a/src/main/java/com/pgvector/PGvector.java b/src/main/java/com/pgvector/PGvector.java
@@ -5,12 +5,14 @@
 import java.sql.SQLException;
 import java.util.Arrays;
 import org.postgresql.PGConnection;
+import org.postgresql.util.ByteConverter;
+import org.postgresql.util.PGBinaryObject;
 import org.postgresql.util.PGobject;
 
 /**
  * PGvector class
  */
-public class PGvector extends PGobject implements Serializable, Cloneable {
+public class PGvector extends PGobject implements PGBinaryObject, Serializable, Cloneable {
     private float[] vec;
 
     /**
@@ -62,6 +64,45 @@ public String getValue() {
         }
     }
 
+    /**
+     * Returns the number of bytes for the binary representation
+     */
+    public int lengthInBytes() {
+        return vec == null ? 0 : 4 + vec.length * 4;
+    }
+
+    /**
+     * Sets the value from a binary representation of a vector
+     */
+    public void setByteValue(byte[] value, int offset) throws SQLException {
+        int dim = ByteConverter.int2(value, offset);
+
+        int unused = ByteConverter.int2(value, offset + 2);
+        if (unused != 0) {
+            throw new SQLException("expected unused to be 0");
+        }
+
+        vec = new float[dim];
+        for (int i = 0; i < dim; i++) {
+            vec[i] = ByteConverter.float4(value, offset + 4 + i * 4);
+        }
+    }
+
+    /**
+     * Writes the binary representation of a vector
+     */
+    public void toBytes(byte[] bytes, int offset) {
+        if (vec == null) {
+            return;
+        }
+
+        ByteConverter.int2(bytes, offset, vec.length);
+        ByteConverter.int2(bytes, offset + 2, 0);
+        for (int i = 0; i < vec.length; i++) {
+            ByteConverter.float4(bytes, offset + 4 + i * 4, vec[i]);
+        }
+    }
+
     /**
      * Returns an array
      */
diff --git a/src/test/java/com/pgvector/JDBCJava.java b/src/test/java/com/pgvector/JDBCJava.java
@@ -2,10 +2,14 @@
 
 import java.sql.*;
 import com.pgvector.PGvector;
+import org.postgresql.PGConnection;
 
 public class JDBCJava {
-    public static void example() throws SQLException {
+    public static void example(boolean read_binary) throws SQLException {
         Connection conn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_java_test");
+        if (read_binary) {
+            conn.unwrap(PGConnection.class).setPrepareThreshold(-1);
+        }
 
         Statement setupStmt = conn.createStatement();
         setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
diff --git a/src/test/scala/com/pgvector/MainSpec.scala b/src/test/scala/com/pgvector/MainSpec.scala
@@ -2,7 +2,11 @@ package com.pgvector
 
 class HelloSpec extends munit.FunSuite {
   test("JDBCJava") {
-    JDBCJava.example()
+    JDBCJava.example(false)
+  }
+
+  test("JDBCJava read binary") {
+    JDBCJava.example(true)
   }
 
   test("SpringJDBC") {
