我应该使用什么 JPA + Hibernate 数据类型来支持 PostgreSQL 数据库中的向量扩展？

Question

我应该使用什么 JPA + Hibernate 数据类型来支持 PostgreSQL 数据库中的向量扩展，以便它允许我使用 JPA 实体创建嵌入？

CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));

PG向量

Answer 1

您可以使用 vladmihalcea Hibernate 类型将向量类型转换为 List<Double>，这样就可以使用 JpaRepository 进行保存或查询。

1. 添加依赖到pom.xml文件：

   <dependency>
     <groupId>io.hypersistence</groupId>
     <artifactId>hypersistence-utils-hibernate-55</artifactId>
     <version>3.5.0</version>
   </dependency>
   

2. 创建项目类：

   import com.fasterxml.jackson.annotation.JsonInclude;
   import io.hypersistence.utils.hibernate.type.json.JsonType;
   import lombok.Data;
   import lombok.NoArgsConstructor;
   import org.hibernate.annotations.Type;
   import org.hibernate.annotations.TypeDef;
   
   import javax.persistence.*;
   import java.util.List;
   
   @Data
   @NoArgsConstructor
   @Entity
   @Table(name = "items")
   @JsonInclude(JsonInclude.Include.NON_NULL)
   @TypeDef(name = "json", typeClass = JsonType.class)
   public class Item {
       @Id
       @GeneratedValue(strategy = GenerationType.IDENTITY)
       private Long id;
   
       @Type(type = "json")
       @Column(columnDefinition = "vector")
       private List<Double> embedding;
   }
   

3. 创建一个支持保存和查找的 JpaRepository 接口。您可以使用本机 SQL 编写自定义 findNearestNeighbors 方法

   import org.springframework.data.jpa.repository.JpaRepository;
   
   public interface ItemRepository extends JpaRepository<Item, Long> {
   
       // Find nearest neighbors by a vector, for example value = "[1,2,3]"
       // This also works, cast is equals to the :: operator in postgresql
       //@Query(nativeQuery = true, value = "SELECT * FROM items ORDER BY embedding <-> cast(? as vector) LIMIT 5")
       @Query(nativeQuery = true, value = "SELECT * FROM items ORDER BY embedding <-> ? \\:\\:vector LIMIT 5")
       List<Item> findNearestNeighbors(String value);
   
       // Find nearest neighbors by a record in the same table
       @Query(nativeQuery = true, value = "SELECT * FROM items WHERE id != :id ORDER BY embedding <-> (SELECT embedding FROM items WHERE id = :id) LIMIT 5")
       List<Item> findNearestNeighbors(Long id);
   }
   

4. 测试创建、查询和查找NearestNeighbors：

   @Autowired
   private ItemRepository itemRepository;
   
   @Test
   @Rollback(false)
   @Transactional
   public void createItem() {
       Item item = new Item();
       Random rand = new Random();
       List<Double> embedding = new ArrayList<>();
       for (int i = 0; i < 3; i++)
           embedding.add(rand.nextDouble());
       item.setEmbedding(embedding);
       itemRepository.save(item);
   }
   
   @Test
   public void loadItems() {
       final List<Item> items = itemRepository.findAll();
       System.out.println(items);
   }
   
   @Test
   public void findNearestNeighbors() {
       final String value = "[0.1, 0.2, 0.3]";
       final List<Item> items = itemRepository.findNearestNeighbors(value);
       System.out.println(items);
   }