Какой тип данных JPA + Hibernate следует использовать для поддержки векторного расширения в базе данных PostgreSQL?

Question

Какой тип данных JPA + Hibernate мне следует использовать для поддержки векторного расширения в базе данных PostgreSQL, чтобы оно позволяло мне создавать внедрения с использованием объекта JPA?

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

pgvector

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. Создайте класс 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, поддерживающий сохранение и поиск. Вы можете написать собственные методы findNearestNeighbors с помощью собственного SQL.

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

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>

Для сопоставления Entity (TypeDef был удален в версии 6):

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

}