{{ message }}
-
Notifications
You must be signed in to change notification settings - Fork 6
Expand file tree
/
Copy pathJDBCJavaTest.java
More file actions
221 lines (183 loc) · 9.48 KB
/
Copy pathJDBCJavaTest.java
File metadata and controls
221 lines (183 loc) · 9.48 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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));
}
}
You can’t perform that action at this time.