pgvector-java/examples/disco/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
66
67
68
69
packagecom.example;

importjava.sql.Connection;
importjava.sql.DriverManager;
importjava.sql.PreparedStatement;
importjava.sql.ResultSet;
importjava.sql.SQLException;
importjava.sql.Statement;
importjava.util.ArrayList;
importcom.pgvector.PGvector;
importorg.ankane.disco.Data;
importorg.ankane.disco.Dataset;
importorg.ankane.disco.Recommender;

publicclassExample {
publicstaticvoidmain(String[] args) throwsException {
Connectionconn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_example");
StatementsetupStmt = conn.createStatement();
setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
PGvector.addVectorType(conn);

StatementcreateStmt = conn.createStatement();
createStmt.executeUpdate("DROP TABLE IF EXISTS users");
createStmt.executeUpdate("DROP TABLE IF EXISTS movies");
createStmt.executeUpdate("CREATE TABLE users (id integer PRIMARY KEY, factors vector(20))");
createStmt.executeUpdate("CREATE TABLE movies (name text PRIMARY KEY, factors vector(20))");

Dataset<Integer, String> data = Data.loadMovieLens();
Recommender<Integer, String> recommender = Recommender
 .builder()
 .factors(20)
 .fitExplicit(data);

for (IntegeruserId : recommender.userIds()) {
PreparedStatementinsertStmt = conn.prepareStatement("INSERT INTO users (id, factors) VALUES (?, ?)");
insertStmt.setInt(1, userId);
insertStmt.setObject(2, newPGvector(recommender.userFactors(userId).get()));
insertStmt.executeUpdate();
 }

for (StringitemId : recommender.itemIds()) {
PreparedStatementinsertStmt = conn.prepareStatement("INSERT INTO movies (name, factors) VALUES (?, ?)");
insertStmt.setString(1, itemId);
insertStmt.setObject(2, newPGvector(recommender.itemFactors(itemId).get()));
insertStmt.executeUpdate();
 }

Stringmovie = "Star Wars (1977)";
System.out.printf("Item-based recommendations for %s\n", movie);
PreparedStatementneighborStmt = conn.prepareStatement("SELECT name FROM movies WHERE name != ? ORDER BY factors <=> (SELECT factors FROM movies WHERE name = ?) LIMIT 5");
neighborStmt.setString(1, movie);
neighborStmt.setString(2, movie);
ResultSetrs = neighborStmt.executeQuery();
while (rs.next()) {
System.out.println("- " + rs.getString("name"));
 }

intuserId = 123;
System.out.printf("\nUser-based recommendations for user %d\n", userId);
neighborStmt = conn.prepareStatement("SELECT name FROM movies ORDER BY factors <#> (SELECT factors FROM users WHERE id = ?) LIMIT 5");
neighborStmt.setInt(1, userId);
rs = neighborStmt.executeQuery();
while (rs.next()) {
System.out.println("- " + rs.getString("name"));
 }

conn.close();
 }
}