在PostgreSQL数据库中，我应该使用哪种JPA + Hibernate数据类型来支持向量扩展？

Question

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

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

pgvector

- user352290

Answer 1

你可以使用vladmihalcea的Hibernate类型将向量类型转换为List，这样就可以在JpaRepository中保存或查询。

1. 在 pom.xml 文件中添加依赖项：

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

2. 创建 Item 类：

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> {

    // 根据向量查找最近的邻居，例如 value = "[1,2,3]"
    // 这也可以工作，cast 相当于 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);

    // 根据同一表中的记录查找最近的邻居
    @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. 测试创建、查询和查找最近的邻居：

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

- zhouhsc_cn

Answer 2

在Hibernate 6和Hibernate 5中，处理这个问题的方式略有不同。 这是在Hibernate 6中应该处理的方式，因为Hibernate 5已经有了答案。

你需要使用来自io.hypersistence的依赖项。检查你使用的Hibernate版本，因为每个版本都有不同的构件。对于Hibernate ORM 6.2：

<!-- https://mvnrepository.com/artifact/io.hypersistence/hypersistence-utils-hibernate-62 -->
<dependency>
    <groupId>io.hypersistence</groupId>
    <artifactId>hypersistence-utils-hibernate-62</artifactId>
    <version>3.5.3</version>
</dependency>

对于实体映射（在v.6中已删除TypeDef）：

import jakarta.persistence.*;

import org.hibernate.annotations.Type;
import io.hypersistence.utils.hibernate.type.json.JsonType;
import java.util.List;
import java.util.UUID;


@Entity
@Table(name = "items")
public class Item {

    @Id
    @Column(name = "id")
    private UUID id;

    @Basic
    @Type(JsonType.class)
    @Column(name = "embedding", columnDefinition = "vector")
    private List<Double> embedding;
    
    ...
}

JPA Repository 查询：

@Repository
public interface ItemRepository extends JpaRepository<Item, UUID> {

    @Query(nativeQuery = true, 
           value = "SELECT * FROM items ORDER BY embedding <-> cast(? as vector) LIMIT 3")
    List<Item> findNearestNeighbors(String embedding);

}

- Chris Sekas