Added support for bit type

Author: ankane

CHANGELOG.md

@@ -1,6 +1,6 @@
 ## 0.1.5 (unreleased)
 
-- Added support for `halfvec` and `sparsevec` types
+- Added support for `halfvec`, `bit`, and `sparsevec` types
 
 ## 0.1.4 (2023-12-08)
 

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

@@ -0,0 +1,124 @@
+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.ByteConverter;
+import org.postgresql.util.PGBinaryObject;
+import org.postgresql.util.PGobject;
+
+/**
+ * PGbit class
+ */
+public class PGbit extends PGobject implements PGBinaryObject, Serializable, Cloneable {
+    private int length;
+    private byte[] data;
+
+    /**
+     * Constructor
+     */
+    public PGbit() {
+        type = "bit";
+    }
+
+    /**
+     * Constructor
+     *
+     * @param v boolean array
+     */
+    public PGbit(boolean[] v) {
+        this();
+        length = v.length;
+        data = new byte[(length + 7) / 8];
+        for (int i = 0; i < length; i++) {
+            data[i / 8] |= (v[i] ? 1 : 0) << (7 - (i % 8));
+        }
+    }
+
+    /**
+     * Constructor
+     *
+     * @param s text representation of a bit string
+     * @throws SQLException exception
+     */
+    public PGbit(String s) throws SQLException {
+        this();
+        setValue(s);
+    }
+
+    /**
+     * Sets the value from a text representation of a bit string
+     */
+    public void setValue(String s) throws SQLException {
+        if (s == null) {
+            data = null;
+        } else {
+            length = s.length();
+            data = new byte[(length + 7) / 8];
+            for (int i = 0; i < length; i++) {
+                data[i / 8] |= (s.charAt(i) != '0' ? 1 : 0) << (7 - (i % 8));
+            }
+        }
+    }
+
+    /**
+     * Returns the text representation of a bit string
+     */
+    public String getValue() {
+        if (data == null) {
+            return null;
+        } else {
+            StringBuilder sb = new StringBuilder(length);
+            for (int i = 0; i < length; i++) {
+                sb.append(((data[i / 8] >> (7 - (i % 8))) & 1) == 1 ? '1' : '0');
+            }
+            return sb.toString();
+        }
+    }
+
+    /**
+     * Returns the number of bytes for the binary representation
+     */
+    public int lengthInBytes() {
+        return data == null ? 0 : 4 + data.length;
+    }
+
+    /**
+     * Sets the value from a binary representation of a bit string
+     */
+    public void setByteValue(byte[] value, int offset) throws SQLException {
+        length = ByteConverter.int4(value, offset);
+        data = new byte[(length + 7) / 8];
+        for (int i = 0; i < data.length; i++) {
+            data[i] = value[offset + 4 + i];
+        }
+    }
+
+    /**
+     * Writes the binary representation of a bit string
+     */
+    public void toBytes(byte[] bytes, int offset) {
+        if (data == null) {
+            return;
+        }
+
+        ByteConverter.int4(bytes, offset, length);
+        for (int i = 0; i < data.length; i++) {
+            bytes[offset + 4 + i] = data[i];
+        }
+    }
+
+    /**
+     * Registers the bit type
+     *
+     * @param conn connection
+     * @throws SQLException exception
+     */
+    public static void addBitType(Connection conn) throws SQLException {
+        conn.unwrap(PGConnection.class).addDataType("bit", PGbit.class);
+    }
+}

src/test/java/com/pgvector/JDBCJavaTest.java

@@ -13,6 +13,7 @@
 import org.junit.jupiter.api.Test;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
 
 public class JDBCJavaTest {
@@ -117,6 +118,59 @@ void halfvecExample(boolean readBinary) throws SQLException {
         assertNull(embeddings.get(3));
     }
 
+    @Test
+    void testBitReadText() throws SQLException {
+        bitExample(false);
+    }
+
+    @Test
+    void testBitReadBinary() throws SQLException {
+        bitExample(true);
+    }
+
+    void bitExample(boolean readBinary) throws SQLException {
+        Connection conn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_java_test");
+        if (readBinary) {
+            conn.unwrap(PGConnection.class).setPrepareThreshold(-1);
+        }
+
+        Statement setupStmt = conn.createStatement();
+        setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
+        setupStmt.executeUpdate("DROP TABLE IF EXISTS jdbc_items");
+
+        PGbit.addBitType(conn);
+
+        Statement createStmt = conn.createStatement();
+        createStmt.executeUpdate("CREATE TABLE jdbc_items (id bigserial PRIMARY KEY, embedding bit(9))");
+
+        PreparedStatement insertStmt = conn.prepareStatement("INSERT INTO jdbc_items (embedding) VALUES (?), (?), (?), (?)");
+        insertStmt.setObject(1, new PGbit(new boolean[] {false, false, false, false, false, false, false, false, false}));
+        insertStmt.setObject(2, new PGbit(new boolean[] {false, true, false, true, false, false, false, false, true}));
+        insertStmt.setObject(3, new PGbit(new boolean[] {false, true, true, true, false, false, false, false, true}));
+        insertStmt.setObject(4, null);
+        insertStmt.executeUpdate();
+
+        PreparedStatement neighborStmt = conn.prepareStatement("SELECT * FROM jdbc_items ORDER BY embedding <~> ? LIMIT 5");
+        neighborStmt.setObject(1, new PGbit(new boolean[] {false, true, false, true, false, false, false, false, true}));
+        ResultSet rs = neighborStmt.executeQuery();
+        List<Long> ids = new ArrayList<>();
+        List<PGbit> embeddings = new ArrayList<>();
+        while (rs.next()) {
+            ids.add(rs.getLong("id"));
+            embeddings.add((PGbit) rs.getObject("embedding"));
+        }
+        assertArrayEquals(new Long[] {2L, 3L, 1L, 4L}, ids.toArray());
+        assertEquals("010100001", embeddings.get(0).getValue());
+        assertEquals("011100001", embeddings.get(1).getValue());
+        assertEquals("000000000", embeddings.get(2).getValue());
+        assertNull(embeddings.get(3));
+
+        Statement indexStmt = conn.createStatement();
+        indexStmt.executeUpdate("CREATE INDEX ON jdbc_items USING ivfflat (embedding bit_hamming_ops) WITH (lists = 100)");
+
+        conn.close();
+    }
+
     @Test
     void testSparsevecReadText() throws SQLException {
         sparsevecExample(false);

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

@@ -0,0 +1,23 @@
+package com.pgvector;
+
+import java.sql.SQLException;
+import java.util.Arrays;
+import com.pgvector.PGbit;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class PGbitTest {
+    @Test
+    void testArrayConstructor() {
+        PGbit vec = new PGbit(new boolean[] {true, false, true});
+        assertEquals("101", vec.getValue());
+    }
+
+    @Test
+    void testStringConstructor() throws SQLException {
+        PGbit vec = new PGbit("101");
+        assertEquals("101", vec.getValue());
+    }
+}