JDBCJavaTest.java

package com.pgvector;

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.pgvector.PGvector;
import org.postgresql.PGConnection;
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 {
    @Test
    void testVectorReadText() throws SQLException {
        vectorExample(false);
    }

    @Test
    void testVectorReadBinary() throws SQLException {
        vectorExample(true);
    }

    void vectorExample(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");

        PGvector.addVectorType(conn);

        Statement createStmt = conn.createStatement();
        createStmt.executeUpdate("CREATE TABLE jdbc_items (id bigserial PRIMARY KEY, embedding vector(3))");

        PreparedStatement insertStmt = conn.prepareStatement("INSERT INTO jdbc_items (embedding) VALUES (?), (?), (?), (?)");
        insertStmt.setObject(1, new PGvector(new float[] {1, 1, 1}));
        insertStmt.setObject(2, new PGvector(new float[] {2, 2, 2}));
        insertStmt.setObject(3, new PGvector(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 PGvector(new float[] {1, 1, 1}));
        ResultSet rs = neighborStmt.executeQuery();
        List<Long> ids = new ArrayList<>();
        List<PGvector> embeddings = new ArrayList<>();
        while (rs.next()) {
            ids.add(rs.getLong("id"));
            embeddings.add((PGvector) 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));

        Statement indexStmt = conn.createStatement();
        indexStmt.executeUpdate("CREATE INDEX ON jdbc_items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)");

        conn.close();
    }

    @Test
    void testHalfvecReadText() throws SQLException {
        halfvecExample(false);
    }

    @Test
    void testHalfvecReadBinary() throws SQLException {
        halfvecExample(true);
    }

    void halfvecExample(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");

        PGvector.registerTypes(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));
    }

    @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.registerType(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);
    }

    @Test
    void testSparsevecReadBinary() throws SQLException {
        sparsevecExample(true);
    }

    void sparsevecExample(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");

        PGvector.registerTypes(conn);

        Statement createStmt = conn.createStatement();
        createStmt.executeUpdate("CREATE TABLE jdbc_items (id bigserial PRIMARY KEY, embedding sparsevec(3))");

        PreparedStatement insertStmt = conn.prepareStatement("INSERT INTO jdbc_items (embedding) VALUES (?), (?), (?), (?)");
        insertStmt.setObject(1, new PGsparsevec(new float[] {1, 1, 1}));
        insertStmt.setObject(2, new PGsparsevec(new float[] {2, 2, 2}));
        insertStmt.setObject(3, new PGsparsevec(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 PGsparsevec(new float[] {1, 1, 1}));
        ResultSet rs = neighborStmt.executeQuery();
        List<Long> ids = new ArrayList<>();
        List<PGsparsevec> embeddings = new ArrayList<>();
        while (rs.next()) {
            ids.add(rs.getLong("id"));
            embeddings.add((PGsparsevec) 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));
    }
}