pgvector-java/examples/cdk/src/main/java/com/example/Example.java at master · pgvector/pgvector-java · GitHub

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
packagecom.example;

importjava.sql.Connection;
importjava.sql.DriverManager;
importjava.sql.PreparedStatement;
importjava.sql.ResultSet;
importjava.sql.SQLException;
importjava.sql.Statement;
importcom.pgvector.PGbit;
importorg.openscience.cdk.exception.CDKException;
importorg.openscience.cdk.exception.InvalidSmilesException;
importorg.openscience.cdk.fingerprint.CircularFingerprinter;
importorg.openscience.cdk.fingerprint.IBitFingerprint;
importorg.openscience.cdk.interfaces.IAtomContainer;
importorg.openscience.cdk.silent.SilentChemObjectBuilder;
importorg.openscience.cdk.smiles.SmilesParser;

publicclassExample {
publicstaticvoidmain(String[] args) throwsCDKException, InvalidSmilesException, SQLException {
Connectionconn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_example");

StatementsetupStmt = conn.createStatement();
setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
setupStmt.executeUpdate("DROP TABLE IF EXISTS molecules");

StatementcreateStmt = conn.createStatement();
createStmt.executeUpdate("CREATE TABLE molecules (id text PRIMARY KEY, fingerprint bit(2048))");

String[] molecules = {
"Cc1ccccc1",
"Cc1ncccc1",
"c1ccccn1"
 };
for (Stringmolecule : molecules) {
PreparedStatementinsertStmt = conn.prepareStatement("INSERT INTO molecules (id, fingerprint) VALUES (?, ?)");
insertStmt.setString(1, molecule);
insertStmt.setObject(2, newPGbit(generateFingerprint(molecule)));
insertStmt.executeUpdate();
 }

StringqueryMolecule = "c1ccco1";
PreparedStatementqueryStmt = conn.prepareStatement("SELECT id, fingerprint <%> ? AS distance FROM molecules ORDER BY distance LIMIT 5");
queryStmt.setObject(1, newPGbit(generateFingerprint(queryMolecule)));
ResultSetrs = queryStmt.executeQuery();
while (rs.next()) {
System.out.println(String.format("%s: %f", rs.getString("id"), rs.getDouble("distance")));
 }

conn.close();
 }

privatestaticboolean[] generateFingerprint(Stringmolecule) throwsCDKException, InvalidSmilesException {
SmilesParsersp = newSmilesParser(SilentChemObjectBuilder.getInstance());
IAtomContainerm = sp.parseSmiles(molecule);

CircularFingerprinterfingerprinter = newCircularFingerprinter(CircularFingerprinter.CLASS_ECFP6, 2048);
IBitFingerprintfp = fingerprinter.getBitFingerprint(m);

boolean[] ba = newboolean[(int) fp.size()];
for (inti : fp.getSetbits()) {
ba[i] = true;
 }
returnba;
 }
}