.github/workflows/build.yml

@@ -28,6 +28,7 @@ jobs:
       # Hibernate 6.4 and HttpClient require Java 11+
       - if: ${{ matrix.java == 8 }}
         run: |
+          rm src/test/java/com/pgvector/CohereTest.java
           rm src/test/java/com/pgvector/HibernateTest.java
           rm src/test/java/com/pgvector/OpenAITest.java
       - run: mvn -B -ntp test

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.1.6 (2024-07-17)
+
+- Added `byte[]` constructor to `PGbit`
+
 ## 0.1.5 (2024-06-25)
 
 - Added support for `halfvec`, `bit`, and `sparsevec` types

README.md

@@ -14,14 +14,14 @@ For Maven, add to `pom.xml` under `<dependencies>`:
 <dependency>
     <groupId>com.pgvector</groupId>
     <artifactId>pgvector</artifactId>
-    <version>0.1.5</version>
+    <version>0.1.6</version>
 </dependency>
 ```
 
 For sbt, add to `build.sbt`:
 
 ```sbt
-libraryDependencies += "com.pgvector" % "pgvector" % "0.1.5"
+libraryDependencies += "com.pgvector" % "pgvector" % "0.1.6"
 ```
 
 For other build tools, see [this page](https://central.sonatype.com/artifact/com.pgvector/pgvector).
@@ -36,6 +36,8 @@ And follow the instructions for your database library:
 Or check out an example:
 
 - [Embeddings](src/test/java/com/pgvector/OpenAITest.java) with OpenAI
+- [Binary embeddings](src/test/java/com/pgvector/CohereTest.java) with Cohere
+- [Bulk loading](src/test/java/com/pgvector/LoadingTest.java) with `COPY`
 
 ## JDBC (Java)
 

pom.xml

@@ -6,7 +6,7 @@
     <packaging>jar</packaging>
     <description>pgvector support for Java, Kotlin, Groovy, and Scala</description>
     <url>https://github.com/pgvector/pgvector-java</url>
-    <version>0.1.5</version>
+    <version>0.1.6</version>
     <licenses>
         <license>
             <name>MIT</name>

src/main/java/com/pgvector/PGbit.java

@@ -39,6 +39,17 @@ public PGbit(boolean[] v) {
         }
     }
 
+    /**
+     * Constructor
+     *
+     * @param v byte array
+     */
+    public PGbit(byte[] v) {
+        this();
+        length = Math.multiplyExact(v.length, 8);
+        data = v;
+    }
+
     /**
      * Constructor
      *

src/test/java/com/pgvector/CohereTest.java

@@ -0,0 +1,101 @@
+package com.pgvector;
+
+import java.io.IOException;
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpRequest.BodyPublishers;
+import java.net.http.HttpResponse;
+import java.net.http.HttpResponse.BodyHandlers;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.List;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.pgvector.PGvector;
+import org.postgresql.PGConnection;
+import org.junit.jupiter.api.Test;
+
+public class CohereTest {
+    @Test
+    void example() throws IOException, InterruptedException, SQLException {
+        String apiKey = System.getenv("CO_API_KEY");
+        if (apiKey == null) {
+            return;
+        }
+
+        Connection conn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_example");
+
+        Statement setupStmt = conn.createStatement();
+        setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
+        setupStmt.executeUpdate("DROP TABLE IF EXISTS documents");
+
+        PGvector.addVectorType(conn);
+
+        Statement createStmt = conn.createStatement();
+        createStmt.executeUpdate("CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))");
+
+        String[] input = {
+            "The dog is barking",
+            "The cat is purring",
+            "The bear is growling"
+        };
+        List<byte[]> embeddings = fetchEmbeddings(input, "search_document", apiKey);
+        for (int i = 0; i < input.length; i++) {
+            PreparedStatement insertStmt = conn.prepareStatement("INSERT INTO documents (content, embedding) VALUES (?, ?)");
+            insertStmt.setString(1, input[i]);
+            insertStmt.setObject(2, new PGbit(embeddings.get(i)));
+            insertStmt.executeUpdate();
+        }
+
+        String query = "forest";
+        byte[] queryEmbedding = fetchEmbeddings(new String[] {query}, "search_query", apiKey).get(0);
+        PreparedStatement neighborStmt = conn.prepareStatement("SELECT * FROM documents ORDER BY embedding <~> ? LIMIT 5");
+        neighborStmt.setObject(1, new PGbit(queryEmbedding));
+        ResultSet rs = neighborStmt.executeQuery();
+        while (rs.next()) {
+            System.out.println(rs.getString("content"));
+        }
+
+        conn.close();
+    }
+
+    // https://docs.cohere.com/reference/embed
+    private List<byte[]> fetchEmbeddings(String[] texts, String inputType, String apiKey) throws IOException, InterruptedException {
+        ObjectMapper mapper = new ObjectMapper();
+        ObjectNode root = mapper.createObjectNode();
+        for (String v : texts) {
+            root.withArray("texts").add(v);
+        }
+        root.put("model", "embed-english-v3.0");
+        root.put("input_type", inputType);
+        root.withArray("embedding_types").add("binary");
+        String json = mapper.writeValueAsString(root);
+
+        HttpClient client = HttpClient.newHttpClient();
+        HttpRequest request = HttpRequest.newBuilder()
+            .uri(URI.create("https://api.cohere.com/v1/embed"))
+            .header("Authorization", "Bearer " + apiKey)
+            .header("Content-Type", "application/json")
+            .POST(BodyPublishers.ofString(json))
+            .build();
+        HttpResponse<String> response = client.send(request, BodyHandlers.ofString());
+
+        List<byte[]> embeddings = new ArrayList<>();
+        for (JsonNode n : mapper.readTree(response.body()).get("embeddings").get("binary")) {
+            byte[] embedding = new byte[n.size()];
+            int i = 0;
+            for (JsonNode v : n) {
+                embedding[i++] = (byte) v.asInt();
+            }
+            embeddings.add(embedding);
+        }
+        return embeddings;
+    }
+}

src/test/java/com/pgvector/LoadingTest.java

@@ -0,0 +1,82 @@
+package com.pgvector;
+
+import java.io.UnsupportedEncodingException;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.ArrayList;
+import com.pgvector.PGvector;
+import org.postgresql.PGConnection;
+import org.postgresql.copy.CopyIn;
+import org.postgresql.copy.CopyManager;
+import org.postgresql.core.BaseConnection;
+import org.junit.jupiter.api.Test;
+
+public class LoadingTest {
+    @Test
+    void example() throws SQLException {
+        if (System.getenv("TEST_LOADING") == null) {
+            return;
+        }
+
+        // generate random data
+        int rows = 1000000;
+        int dimensions = 128;
+        ArrayList<float[]> embeddings = new ArrayList<>(rows);
+        for (int i = 0; i < rows; i++) {
+            float[] embedding = new float[dimensions];
+            for (int j = 0; j < dimensions; j++) {
+                embedding[j] = (float) Math.random();
+            }
+            embeddings.add(embedding);
+        }
+
+        // enable extension
+        Connection conn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_example");
+        Statement setupStmt = conn.createStatement();
+        setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
+        PGvector.addVectorType(conn);
+
+        // create table
+        setupStmt.executeUpdate("DROP TABLE IF EXISTS items");
+        setupStmt.executeUpdate("CREATE TABLE items (id bigserial, embedding vector(128))");
+
+        // load data
+        System.out.println("Loading 1000000 rows");
+
+        CopyManager copyManager = new CopyManager((BaseConnection) conn);
+        // TODO use binary format
+        CopyIn copyIn = copyManager.copyIn("COPY items (embedding) FROM STDIN");
+        for (int i = 0; i < rows; i++) {
+            if (i % 10000 == 0) {
+                System.out.print(".");
+            }
+
+            PGvector embedding = new PGvector(embeddings.get(i));
+            byte[] bytes = (embedding.getValue() + "\n").getBytes();
+            copyIn.writeToCopy(bytes, 0, bytes.length);
+        }
+        copyIn.endCopy();
+
+        System.out.println("\nSuccess!");
+
+        // create any indexes *after* loading initial data (skipping for this example)
+        boolean createIndex = false;
+        if (createIndex) {
+            System.out.println("Creating index");
+            Statement createIndexStmt = conn.createStatement();
+            createIndexStmt.executeUpdate("SET maintenance_work_mem = '8GB'");
+            createIndexStmt.executeUpdate("SET max_parallel_maintenance_workers = 7");
+            createIndexStmt.executeUpdate("CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)");
+        }
+
+        // update planner statistics for good measure
+        Statement analyzeStmt = conn.createStatement();
+        analyzeStmt.executeUpdate("ANALYZE items");
+
+        conn.close();
+    }
+}

src/test/java/com/pgvector/OpenAITest.java

@@ -73,7 +73,7 @@ private List<float[]> fetchEmbeddings(String[] input, String apiKey) throws IOEx
         for (String v : input) {
             root.withArray("input").add(v);
         }
-        root.put("model", "text-embedding-ada-002");
+        root.put("model", "text-embedding-3-small");
         String json = mapper.writeValueAsString(root);
 
         HttpClient client = HttpClient.newHttpClient();

src/test/java/com/pgvector/PGbitTest.java

@@ -17,6 +17,13 @@ void testArrayConstructor() {
         assertArrayEquals(new boolean[] {false, true, false, true, false, false, false, false, true}, vec.toArray());
     }
 
+    void testEmptyArrayConstructor() {
+        PGbit vec = new PGbit(new boolean[] {});
+        assertEquals(0, vec.length());
+        assertArrayEquals(new byte[] {}, vec.toByteArray());
+        assertArrayEquals(new boolean[] {}, vec.toArray());
+    }
+
     @Test
     void testStringConstructor() throws SQLException {
         PGbit vec = new PGbit("010100001");