Added registerTypes function

ankane · ankane · commit 86af42dedfe3 · 2024-06-25T20:58:28.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## 0.1.5 (unreleased)
 
 - Added support for `halfvec`, `bit`, and `sparsevec` types
+- Added `registerTypes` function
 
 ## 0.1.4 (2023-12-08)
 
diff --git a/src/main/java/com/pgvector/PGhalfvec.java b/src/main/java/com/pgvector/PGhalfvec.java
@@ -100,14 +100,4 @@ public String getValue() {
     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/main/java/com/pgvector/PGsparsevec.java b/src/main/java/com/pgvector/PGsparsevec.java
@@ -278,14 +278,4 @@ public int[] getIndices() {
     public float[] getValues() {
         return values;
     }
-
-    /**
-     * Registers the sparsevec type
-     *
-     * @param conn connection
-     * @throws SQLException exception
-     */
-    public static void addSparsevecType(Connection conn) throws SQLException {
-        conn.unwrap(PGConnection.class).addDataType("sparsevec", PGsparsevec.class);
-    }
 }
diff --git a/src/main/java/com/pgvector/PGvector.java b/src/main/java/com/pgvector/PGvector.java
@@ -149,4 +149,17 @@ public float[] toArray() {
     public static void addVectorType(Connection conn) throws SQLException {
         conn.unwrap(PGConnection.class).addDataType("vector", PGvector.class);
     }
+
+    /**
+     * Registers the vector, halfvec, and sparsevec types
+     *
+     * @param conn connection
+     * @throws SQLException exception
+     */
+    public static void registerTypes(Connection conn) throws SQLException {
+        // bit type should be registered separately
+        addVectorType(conn);
+        conn.unwrap(PGConnection.class).addDataType("halfvec", PGhalfvec.class);
+        conn.unwrap(PGConnection.class).addDataType("sparsevec", PGsparsevec.class);
+    }
 }
diff --git a/src/test/java/com/pgvector/JDBCJavaTest.java b/src/test/java/com/pgvector/JDBCJavaTest.java
@@ -90,7 +90,7 @@ void halfvecExample(boolean readBinary) throws SQLException {
         setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
         setupStmt.executeUpdate("DROP TABLE IF EXISTS jdbc_items");
 
-        PGhalfvec.addHalfvecType(conn);
+        PGvector.registerTypes(conn);
 
         Statement createStmt = conn.createStatement();
         createStmt.executeUpdate("CREATE TABLE jdbc_items (id bigserial PRIMARY KEY, embedding halfvec(3))");
@@ -191,7 +191,7 @@ void sparsevecExample(boolean readBinary) throws SQLException {
         setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
         setupStmt.executeUpdate("DROP TABLE IF EXISTS jdbc_items");
 
-        PGsparsevec.addSparsevecType(conn);
+        PGvector.registerTypes(conn);
 
         Statement createStmt = conn.createStatement();
         createStmt.executeUpdate("CREATE TABLE jdbc_items (id bigserial PRIMARY KEY, embedding sparsevec(3))");
