GeezerVison is showing github.com — pgvector-java/src/test/java/com/pgvector/JDBCJavaTest.java at master · pgvector/pgvector-java

Expand file tree

JDBCJavaTest.java

Copy path

More file actions

Latest commit

History

221 lines (183 loc) · 9.48 KB

JDBCJavaTest.java

Copy path

File metadata and controls

221 lines (183 loc) · 9.48 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

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));

}