What JPA + Hibernate data type should I use to support the vector extension in a PostgreSQL database?

Question

What JPA + Hibernate data type should I use to support the vector extension in a PostgreSQL database, so that it allows me to create embeddings using a JPA Entity?

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

pgvector

Answer 1

You can use vladmihalcea Hibernate types to convert a vector type to List<Double>, so it is possible to save or query with JpaRepository.

1. Add a dependency to the pom.xml file:

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

2. Create the Item class:

   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. Create a JpaRepository interface that supports save and find. You can write custom findNearestNeighbors methods with native 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. Test create, query and findNearestNeighbors:

   @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

Since 6.4 I think, Hibernate has released a vector module : https://docs.jboss.org/hibernate/orm/current/userguide/html_single/Hibernate_User_Guide.html#vector-module

You import the hibernate-vector lib (you won't need converter anymore) and then this should work :

@Column( name = "the_vector" )
@JdbcTypeCode(SqlTypes.VECTOR)
@Array(length = 3)
private float[] theVector;

(maybe https://stackoverflow.com/a/78212559/4499463 should be closed as duplicated?)

Answer 3

There are slightly different ways to handle this in Hibernate 6 vs Hibernate 5. This is how you should handle this in Hibernate 6, since there is already an answer for Hibernate 5.

You need to use a dependency from io.hypersistence. Check what version of Hibernate you use, since each version has different artifacts. For 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>

For the Entity mapping (TypeDef has been removed in v.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 Query:

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

}