Added support for halfvec type

ankane · ankane · commit 54b8f98c0a95 · 2024-05-18T22:05:57.000-04:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.1.5 (unreleased)
+
+- Added support for `halfvec` type
+
 ## 0.1.4 (2023-12-08)
 
 - Added `List` constructor
diff --git a/src/main/java/com/pgvector/PGhalfvec.java b/src/main/java/com/pgvector/PGhalfvec.java
@@ -0,0 +1,109 @@
+package com.pgvector;
+
+import java.io.Serializable;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Objects;
+import org.postgresql.PGConnection;
+import org.postgresql.util.PGobject;
+
+/**
+ * PGhalfvec class
+ */
+public class PGhalfvec extends PGobject implements Serializable, Cloneable {
+    private float[] vec;
+
+    /**
+     * Constructor
+     */
+    public PGhalfvec() {
+        type = "halfvec";
+    }
+
+    /**
+     * Constructor
+     *
+     * @param v float array
+     */
+    public PGhalfvec(float[] v) {
+        this();
+        vec = v;
+    }
+
+    /**
+     * Constructor
+     *
+     * @param <T> number
+     * @param v list of numbers
+     */
+    public <T extends Number> PGhalfvec(List<T> v) {
+        this();
+        if (Objects.isNull(v)) {
+            vec = null;
+        } else {
+            vec = new float[v.size()];
+            int i = 0;
+            for (T f : v) {
+                vec[i++] = f.floatValue();
+            }
+        }
+    }
+
+    /**
+     * Constructor
+     *
+     * @param s text representation of a half vector
+     * @throws SQLException exception
+     */
+    public PGhalfvec(String s) throws SQLException {
+        this();
+        setValue(s);
+    }
+
+    /**
+     * Sets the value from a text representation of a half vector
+     */
+    public void setValue(String s) throws SQLException {
+        if (s == null) {
+            vec = null;
+        } else {
+            String[] sp = s.substring(1, s.length() - 1).split(",");
+            vec = new float[sp.length];
+            for (int i = 0; i < sp.length; i++) {
+                vec[i] = Float.parseFloat(sp[i]);
+            }
+        }
+    }
+
+    /**
+     * Returns the text representation of a half vector
+     */
+    public String getValue() {
+        if (vec == null) {
+            return null;
+        } else {
+            return Arrays.toString(vec).replace(" ", "");
+        }
+    }
+
+    /**
+     * Returns an array
+     *
+     * @return an array
+     */
+    public float[] toArray() {
+        return vec;
+    }
+
+    /**
+     * Registers the halfvec type
+     *
+     * @param conn connection
+     * @throws SQLException exception
+     */
+    public static void addHalfvecType(Connection conn) throws SQLException {
+        conn.unwrap(PGConnection.class).addDataType("halfvec", PGhalfvec.class);
+    }
+}
diff --git a/src/test/java/com/pgvector/JDBCJavaTest.java b/src/test/java/com/pgvector/JDBCJavaTest.java
@@ -68,4 +68,40 @@ void example(boolean readBinary) throws SQLException {
 
         conn.close();
     }
+
+    @Test
+    void testHalfvec() throws SQLException {
+        Connection conn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_java_test");
+
+        Statement setupStmt = conn.createStatement();
+        setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
+        setupStmt.executeUpdate("DROP TABLE IF EXISTS jdbc_items");
+
+        PGhalfvec.addHalfvecType(conn);
+
+        Statement createStmt = conn.createStatement();
+        createStmt.executeUpdate("CREATE TABLE jdbc_items (id bigserial PRIMARY KEY, embedding halfvec(3))");
+
+        PreparedStatement insertStmt = conn.prepareStatement("INSERT INTO jdbc_items (embedding) VALUES (?), (?), (?), (?)");
+        insertStmt.setObject(1, new PGhalfvec(new float[] {1, 1, 1}));
+        insertStmt.setObject(2, new PGhalfvec(new float[] {2, 2, 2}));
+        insertStmt.setObject(3, new PGhalfvec(new float[] {1, 1, 2}));
+        insertStmt.setObject(4, null);
+        insertStmt.executeUpdate();
+
+        PreparedStatement neighborStmt = conn.prepareStatement("SELECT * FROM jdbc_items ORDER BY embedding <-> ? LIMIT 5");
+        neighborStmt.setObject(1, new PGhalfvec(new float[] {1, 1, 1}));
+        ResultSet rs = neighborStmt.executeQuery();
+        List<Long> ids = new ArrayList<>();
+        List<PGhalfvec> embeddings = new ArrayList<>();
+        while (rs.next()) {
+            ids.add(rs.getLong("id"));
+            embeddings.add((PGhalfvec) rs.getObject("embedding"));
+        }
+        assertArrayEquals(new Long[] {1L, 3L, 2L, 4L}, ids.toArray());
+        assertArrayEquals(new float[] {1, 1, 1}, embeddings.get(0).toArray());
+        assertArrayEquals(new float[] {1, 1, 2}, embeddings.get(1).toArray());
+        assertArrayEquals(new float[] {2, 2, 2}, embeddings.get(2).toArray());
+        assertNull(embeddings.get(3));
+    }
 }
diff --git a/src/test/java/com/pgvector/PGhalfvecTest.java b/src/test/java/com/pgvector/PGhalfvecTest.java
@@ -0,0 +1,43 @@
+package com.pgvector;
+
+import java.sql.SQLException;
+import java.util.Arrays;
+import com.pgvector.PGhalfvec;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class PGhalfvecTest {
+    @Test
+    void testArrayConstructor() {
+        PGhalfvec vec = new PGhalfvec(new float[] {1, 2, 3});
+        assertArrayEquals(new float[] {1, 2, 3}, vec.toArray());
+    }
+
+    @Test
+    void testStringConstructor() throws SQLException {
+        PGhalfvec vec = new PGhalfvec("[1,2,3]");
+        assertArrayEquals(new float[] {1, 2, 3}, vec.toArray());
+    }
+
+    @Test
+    void testFloatListConstructor() {
+        Float[] a = new Float[] {Float.valueOf(1), Float.valueOf(2), Float.valueOf(3)};
+        PGhalfvec vec = new PGhalfvec(Arrays.asList(a));
+        assertArrayEquals(new float[] {1, 2, 3}, vec.toArray());
+    }
+
+    @Test
+    void testDoubleListConstructor() {
+        Double[] a = new Double[] {Double.valueOf(1), Double.valueOf(2), Double.valueOf(3)};
+        PGhalfvec vec = new PGhalfvec(Arrays.asList(a));
+        assertArrayEquals(new float[] {1, 2, 3}, vec.toArray());
+    }
+
+    @Test
+    void testGetValue() {
+        PGhalfvec vec = new PGhalfvec(new float[] {1, 2, 3});
+        assertEquals("[1.0,2.0,3.0]", vec.getValue());
+    }
+}
