Working with Repository Vector Search Methods in Spring Data

1. Overview

In this tutorial, we’ll learn how Spring Data can help integrate with vector databases. Basically, we’ll learn how to create Spring repository classes with methods for searching the underlying vector database, similar to traditional databases. However, we must know that these are preview features available in the Spring Data version 4.0.0-M6.

In the following sections, we’ll explore how Spring Data integrates with both PGvector and MongoDB. While PGvector extends a relational database with vector support, MongoDB provides vector capabilities within a document-oriented model. Together, these examples will illustrate how Spring Data adapts to different types of databases.

2. Key Concepts

Let’s begin by understanding a few key concepts to appreciate the Spring Data framework for performing vector search.

2.1. Vector Databases

In GenAI applications such as AI chatbots, vector databases have a crucial role in implementing the RAG technique. Vector databases help store document chunks in the form of multi-dimensional arrays called embeddings or vectors.

Applications convert user queries into vectors and search the underlying vector database to fetch semantically similar records to the query. Later, they use the search results to build a contextualized prompt and send it to the underlying LLM. Finally, the LLM responds within the boundary of the context provided in the prompt.

2.2. Spring Data Repositories

For searching in traditional relational databases, the Spring repository classes use methods starting with the keyword findBy. To define further query boundaries, the repository methods can end with keywords such as In, Like, NotLike, StartingWith, and others. At runtime, Spring Data JPA translates them to the database queries.

To support vector search,  Spring Data supports keywords that can prefix the method names, such as searchBy and searchTop. Likewise, to support query boundaries based on similarity between vectors, keywords such as Within or Near can be appended to the method names. Furthermore, it also supports the @Query annotations on the repository methods, offering greater control and flexibility when querying the underlying database.

These concepts would become clearer when we explain them in detail through code examples in the later sections.

2.3. Common Spring Data Classes

Now, we’ll discuss a few important common Spring Data classes that are necessary to support vector search across different vector databases:

Moving on, let’s apply all these concepts and see how Spring Data’s vector search feature works.

3. Prerequisites for Spring Data With PGvector

To enable vector search using Spring Data and PGvector in PostgreSQL, the following dependencies and database setup steps are necessary.

3.1. Maven Dependencies

We’ll include the Postgres driver for integrating with PGvector and the Spring Boot JPA starter dependencies:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-jpa</artifactId>
    <version>4.0.0-M2</version>	
</dependency>

<dependency>
    <groupId>org.postgresql</groupId>
    <artifactId>postgresql</artifactId>
    <version>42.7.7</version>
</dependency>

Spring Boot JPA transitively adds the Spring JDBC dependencies. However, to enable vector-based ORM, we must add the Hibernate vector library:

<dependency>
    <groupId>org.hibernate.orm</groupId>
    <artifactId>hibernate-vector</artifactId>
    <version>7.1.0.Final</version>
</dependency>

3.2. Data Setup

PGvector is a PostgreSQL database with the pgvector extension enabled, for supporting vector columns. Hence, in the SQL script, we’ll create the extension before creating tables and inserting data:

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE Book (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding VECTOR(5) NOT NULL,
    year_published VARCHAR(10) NOT NULL
);

INSERT INTO book (content, embedding, year_published) VALUES
('Spring Boot Basics',  '[-0.49966827034950256, -0.025236541405320168, 0.736327588558197, -0.20225830376148224, 0.4081762731075287]'::vector, '2022'),
('Spring Boot Advanced',  '[-0.20951677858829498, 0.17629066109657288, 0.7875414490699768, -0.13002122938632965, 0.5365606546401978]'::vector, '2022'),
-- ..and so on

The script creates a Book table and then inserts data into it, which includes the 5-dimensional Vector type in its embedding column.

Moving on, let’s examine some sample code that demonstrates the Spring Data framework integrating with databases that support vector search capabilities.

4. Integration With PGvector

We’ll start with defining the repository class and the entity class, and finally, we’ll showcase the integration code.

4.1. Repository Class

Let’s begin by taking a look at the PGvectorBookRepository class:

@Repository("pgvectorBookRepository")
public interface PGvectorBookRepository extends JpaRepository<Book, String> {
    SearchResults<Book> searchByYearPublishedAndEmbeddingNear(
      String yearPublished, Vector vector, Score scoreThreshold
    );
    SearchResults<Book> searchByYearPublishedAndEmbeddingWithin(
      String yearPublished, Vector vector, Range<Similarity> range, Limit topK
    );
}

In the repository class, all the Near and Within methods return the results wrapped in the SearchResults<Book> class. We’ll later explore how these methods behave when executed.

The class Book represents the underlying PostgreSQL table book:

@Entity(name = "book")
public class Book {
    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private String id;

    private String content;

    @Column(name = "year_published")
    private String yearPublished;

    @JdbcTypeCode(SqlTypes.VECTOR)
    @Array(length = 5)
    private float[] embedding;

//..standard getters and setters
}

In the Book class, apart from the usual fields, the embedding field of type float[] represents the embedding column in the Book table. Furthermore, the Hibernate core library’s @JdbcTypeCode and @Array annotations help map to the 5-dimensional embedding column of type Vector in the Book table.

4.2. Execute Repository Methods

Moving on, let’s run a few repository methods and see how Spring Data fulfills the requirement related to vector search. In each of these methods, the embedding used corresponds to the sentence Which document has the details about Django?. Furthermore, before invoking the methods, we call getEmbedding() to fetch the query’s LLM vector representation. Typically, this method calls an embedding model to convert the text into embeddings.

First, let’s run the method searchByYearPublishedAndEmbeddingNear():

void whenSearchByYearPublishedAndEmbeddingNear_thenResult() {
    Vector embedding = getEmbedding("Which document has the details about Django?");
    
    SearchResults<Book> results = pgvectorBookRepository.searchByYearPublishedAndEmbeddingNear(
      "2022", embedding,
      Score.of(0.9, ScoringFunction.euclidean())
    );
    assertThat(results).isNotNull();

    List<SearchResult<Book>> resultList = results.getContent();

    assertThat(resultList.size()).isGreaterThan(0);

    resultList.forEach(book -> assertThat(book.getContent()
      .getYearPublished()).isEqualTo("2022")
    );
}

The search method filters by the publication year 2022 and then fetches the book records that are near the embedding. The Score argument, with a value of 0.9, is passed to the repository method. Furthermore, the method uses the Euclidean function to return embeddings that are closer to or near the query vector.

Next, let’s execute the searchByYearPublishedAndEmbeddingWithin() method and study its behavior:

void whenSearchByYearPublishedAndEmbeddingWithin_thenResult() {
    Vector embedding = getEmbedding("Which document has the details about Django?");

    Range<Similarity> range = Range.closed(
        Similarity.of(0.7, ScoringFunction.cosine()),
        Similarity.of(0.9, ScoringFunction.cosine())
    );
    SearchResults<Book> results = pgvectorBookRepository.searchByYearPublishedAndEmbeddingWithin(
      "2022", embedding, range, Limit.of(5)
    );
    assertThat(results).isNotNull();

    List<SearchResult<Book>> resultList = results.getContent();

    assertThat(resultList.size()).isGreaterThan(0).isLessThanOrEqualTo(5);

    resultList.forEach(book -> {
        assertThat(book.getContent().getYearPublished()).isEqualTo("2022");
        assertThat(book.getScore().getValue()).isBetween(0.7, 0.9);
    });
}

In this example, we set a similarity range between 0.7 and 0.9, and use assertions to ensure that all results fall within this range. Since we also set a limit of 5, we verify that the number of records returned doesn’t exceed 5 and that each book’s publication year matches 2022.

Up to this point, we have combined vector similarity with a column filter. However, queries can also be run directly on the vector field alone using methods such as searchByEmbeddingWithin() and searchByEmbeddingNear(), even though we aren’t demonstrating them here.

5. Prerequisites for Spring Data With MongoDB

To get started with integrating Spring Data and MongoDB for vector search, it’s essential to configure the required dependencies and provide the appropriate dataset.

5.1. Maven Dependencies

Let’s include the Spring Boot MongoDB dependencies for demonstrating the vector search capabilities of the Spring Data framework through its integration with MongoDB:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-mongodb</artifactId>
    <version>4.0.0-M2</version>   
</dependency>

5.2. Data Setup

Unlike PostgreSQL, because MongoDB is a non-relational database, the program will read the data from a CSV file:

content,yearPublished,embedding
Spring Boot Basics,2022,"[-0.49966827034950256, -0.025236541405320168, 
0.736327588558197, -0.20225830376148224, 0.4081762731075287]"
Spring Boot Advanced,2022,"[-0.20951677858829498, 0.17629066109657288, 
0.7875414490699768, -0.13002122938632965, 0.5365606546401978]"
..and so on

The CSV file contains the records of a collection of books similar to the PostgreSQL table defined earlier. Furthermore, the sample program will create the Book collection and the vector index book-vector-index on the field embedding.

6. Integration With MongoDB

Similar to PGvector, we’ll query the MongoDB collection Book to meet the multiple requirements of a vector search.

6.1. Repository Class

As usual, let’s first start with the MongoDbBookRepository class:

@Repository("mongoDbBookRepository")
public interface MongoDbBookRepository extends MongoRepository<Book, String> {
    @VectorSearch(indexName = "book-vector-index", limit = "10", numCandidates="200")
    SearchResults<Book> searchByYearPublishedAndEmbeddingNear(String yearPublished, Vector vector, Score score);

    @VectorSearch(indexName = "book-vector-index", limit = "10", numCandidates="200")
    SearchResults<Book> searchByYearPublishedAndEmbeddingWithin(String yearPublished, Vector vector, Range<Similarity> range);
}

The methods in MongoDbBookRepository must use the @VectorSearch annotation to define the vector index and control search constraints such as result limits and candidate size. By combining logical keywords like Near and Within, these method names are automatically translated into corresponding vector search queries at runtime.

All the Repository methods return the records wrapped in the class SearchResults<Book>. Further, the underlying MongoDB Books collection is represented by the class Book:

@Document(collection = "books")
public class Book {
    @Id
    private String id;

    private String name;

    private String yearPublished;

    private Vector embedding;

    // Standard constructor, getters, setters
}

The class is annotated with the @Document annotation from the Spring MongoDB library to represent the books collection in MongoDB. Additionally, it has a constructor and other standard getters and setters for all the fields.

6.2. Execute Repository Methods

Moving on, let’s run the repository methods and verify the results fetched from MongoDB.

First, let’s execute the method searchByYearPublishedAndEmbeddingNear():

void whenSearchByYearPublishedAndEmbeddingNear_thenReturnResult() {
    Vector embedding = getEmbedding("Which document has the details about Django?");

    SearchResults<Book> results = mongoDbBookRepository.searchByYearPublishedAndEmbeddingNear(
      "2022", embedding, Score.of(0.9)
    );
    List<SearchResult<Book>> resultLst = results.getContent();

    assertThat(resultLst.size()).isGreaterThan(0);
    
    resultLst.forEach(content -> {
        Book book = content.getContent();
        assertThat(book.getYearPublished()).isEqualTo("2022");
    });
}

The repository method searchByYearPublishedAndEmbeddingNear() fetches books published in 2022 whose embeddings are close to the given query vector with a similarity score of 0.9.

Next, let’s execute the method searchByYearPublishedAndEmbeddingWithin():

void whenSearchByYearPublishedAndEmbeddingWithin_thenReturnResult() {
    Vector embedding = getEmbedding("Which document has the details about Django?");

    Range<Similarity> range = Range.closed(Similarity.of(0.7), Similarity.of(0.9));
    SearchResults<Book> results = mongoDbBookRepository.searchByYearPublishedAndEmbeddingWithin(
      "2022", embedding, range
    );
    
    assertThat(results).isNotNull();

    List<SearchResult<Book>> resultList = results.getContent();

    assertThat(resultList.size()).isGreaterThan(0).isLessThanOrEqualTo(10);

    resultList.forEach(book -> {
      assertThat(book.getContent().getYearPublished()).isEqualTo("2022");
      assertThat(book.getScore().getValue()).isBetween(0.7, 0.9);
    });
}

The repository method fetches books published in 2022 whose embeddings fall within a similarity score range of 0.7 to 0.9 relative to the query vector.

7. Conclusion

In this article, we learned about the emerging preview vector search feature in the Spring Data framework.

With a standard and consistent approach integrating with vector databases, the development and maintenance of the code becomes easier. However, each vector database has its own constraints or boundaries. Hence, we must thoroughly learn the capabilities of the underlying vector database to maximize the output from the framework.