Add approximate k-NN search support to OpenSearchVectorStore

This commit adds HNSW-based approximate k-NN search as an alternative to
the existing exact k-NN search, providing better performance and
scalability for large vector datasets.

Key changes:
- Add useApproximateKnn, dimensions, and similarity configuration options
- Implement buildApproximateQuery() using OpenSearch native knn query
- Maintain buildExactQuery() for backwards compatibility (default)
- Add comprehensive test coverage for both exact and approximate modes
- Test multiple similarity functions: cosinesimil, l1, l2, linf, innerproduct
- Fix test isolation by using dedicated beans and indexes per similarity function
- Update documentation with new properties and usage examples

The implementation maintains full backward compatibility with approximate
k-NN disabled by default.

Signed-off-by: Jemin Huh <hjm1980@gmail.com>

Co-authored-by: Soby Chacko <soby.chacko@broadcom.com>

Author: JM-Lab

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/main/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/OpenSearchVectorStoreAutoConfiguration.java

@@ -88,14 +88,19 @@ OpenSearchVectorStore vectorStore(OpenSearchVectorStoreProperties properties, Op
 		var mappingJson = Optional.ofNullable(properties.getMappingJson())
 			.orElse(OpenSearchVectorStore.DEFAULT_MAPPING_EMBEDDING_TYPE_KNN_VECTOR_DIMENSION);
 
-		return OpenSearchVectorStore.builder(openSearchClient, embeddingModel)
+		var builder = OpenSearchVectorStore.builder(openSearchClient, embeddingModel)
 			.index(indexName)
 			.mappingJson(mappingJson)
 			.initializeSchema(properties.isInitializeSchema())
 			.observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
 			.customObservationConvention(customObservationConvention.getIfAvailable(() -> null))
-			.batchingStrategy(batchingStrategy)
-			.build();
+			.batchingStrategy(batchingStrategy);
+
+		Optional.ofNullable(properties.getUseApproximateKnn()).ifPresent(builder::useApproximateKnn);
+		Optional.ofNullable(properties.getDimensions()).ifPresent(builder::dimensions);
+		Optional.ofNullable(properties.getSimilarity()).ifPresent(builder::similarityFunction);
+
+		return builder.build();
 	}
 
 	@Configuration(proxyBeanMethods = false)

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/main/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/OpenSearchVectorStoreProperties.java

@@ -38,6 +38,12 @@ public class OpenSearchVectorStoreProperties extends CommonVectorStoreProperties
 
 	private String password;
 
+	private Boolean useApproximateKnn;
+
+	private Integer dimensions;
+
+	private String similarity;
+
 	private String mappingJson;
 
 	/**
@@ -100,6 +106,30 @@ public String getMappingJson() {
 		return this.mappingJson;
 	}
 
+	public Boolean getUseApproximateKnn() {
+		return this.useApproximateKnn;
+	}
+
+	public void setUseApproximateKnn(Boolean useApproximateKnn) {
+		this.useApproximateKnn = useApproximateKnn;
+	}
+
+	public Integer getDimensions() {
+		return this.dimensions;
+	}
+
+	public void setDimensions(Integer dimensions) {
+		this.dimensions = dimensions;
+	}
+
+	public String getSimilarity() {
+		return this.similarity;
+	}
+
+	public void setSimilarity(String similarity) {
+		this.similarity = similarity;
+	}
+
 	public void setMappingJson(String mappingJson) {
 		this.mappingJson = mappingJson;
 	}

spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/opensearch.adoc

@@ -108,7 +108,10 @@ Properties starting with `spring.ai.vectorstore.opensearch.*` are used to config
 |`spring.ai.vectorstore.opensearch.password`| Password for the specified username | -
 |`spring.ai.vectorstore.opensearch.index-name`| Name of the index to store vectors | `spring-ai-document-index`
 |`spring.ai.vectorstore.opensearch.initialize-schema`| Whether to initialize the required schema | `false`
-|`spring.ai.vectorstore.opensearch.similarity-function`| The similarity function to use | `cosinesimil`
+|`spring.ai.vectorstore.opensearch.similarity-function`| The similarity function to use (cosinesimil, l1, l2, linf, innerproduct) | `cosinesimil`
+|`spring.ai.vectorstore.opensearch.use-approximate-knn`| Whether to use approximate k-NN for faster searches. If true, uses HNSW-based approximate search. If false, uses exact brute-force k-NN. See link:https://opensearch.org/docs/latest/search-plugins/knn/approximate-knn/[Approximate k-NN] and link:https://opensearch.org/docs/latest/search-plugins/knn/knn-score-script/[Exact k-NN] | `false`
+|`spring.ai.vectorstore.opensearch.dimensions`| Number of dimensions for vector embeddings. Used when creating index mapping for approximate k-NN. If not set, uses the embedding model's dimensions. | `1536`
+|`spring.ai.vectorstore.opensearch.mapping-json`| Custom JSON mapping for the index. Overrides default mapping generation. | -
 |`spring.ai.vectorstore.opensearch.read-timeout`| Time to wait for response from the opposite endpoint. 0 - infinity. | -
 |`spring.ai.vectorstore.opensearch.connect-timeout`| Time to wait until connection established. 0 - infinity. | -
 |`spring.ai.vectorstore.opensearch.path-prefix`| Path prefix for OpenSearch API endpoints. Useful when OpenSearch is behind a reverse proxy with a non-root path. | -
@@ -191,6 +194,8 @@ public VectorStore vectorStore(OpenSearchClient openSearchClient, EmbeddingModel
     return OpenSearchVectorStore.builder(openSearchClient, embeddingModel)
         .index("custom-index")                // Optional: defaults to "spring-ai-document-index"
         .similarityFunction("l2")             // Optional: defaults to "cosinesimil"
+        .useApproximateKnn(true)              // Optional: defaults to false (exact k-NN)
+        .dimensions(1536)                     // Optional: defaults to 1536 or embedding model's dimensions
         .initializeSchema(true)               // Optional: defaults to false
         .batchingStrategy(new TokenCountBatchingStrategy()) // Optional: defaults to TokenCountBatchingStrategy
         .build();

vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStore.java

@@ -170,6 +170,10 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem
 
 	private String similarityFunction;
 
+	private final boolean useApproximateKnn;
+
+	private final int dimensions;
+
 	/**
 	 * Creates a new OpenSearchVectorStore using the builder pattern.
 	 * @param builder The configured builder instance
@@ -187,6 +191,8 @@ protected OpenSearchVectorStore(Builder builder) {
 		// https://opensearch.org/docs/latest/search-plugins/knn/approximate-knn/#spaces
 		this.similarityFunction = builder.similarityFunction;
 		this.initializeSchema = builder.initializeSchema;
+		this.useApproximateKnn = builder.useApproximateKnn;
+		this.dimensions = builder.dimensions;
 	}
 
 	/**
@@ -270,17 +276,38 @@ public List<Document> doSimilaritySearch(SearchRequest searchRequest) {
 
 	public List<Document> similaritySearch(float[] embedding, int topK, double similarityThreshold,
 			Filter.Expression filterExpression) {
-		return similaritySearch(new org.opensearch.client.opensearch.core.SearchRequest.Builder()
-			.query(getOpenSearchSimilarityQuery(embedding, filterExpression))
+		return similaritySearch(
+				this.useApproximateKnn ? buildApproximateQuery(embedding, topK, similarityThreshold, filterExpression)
+						: buildExactQuery(embedding, topK, similarityThreshold, filterExpression));
+	}
+
+	private org.opensearch.client.opensearch.core.SearchRequest buildApproximateQuery(float[] embedding, int topK,
+			double similarityThreshold, Filter.Expression filterExpression) {
+		return new org.opensearch.client.opensearch.core.SearchRequest.Builder().index(this.index)
+			.query(Query.of(builder -> builder.knn(knnQueryBuilder -> knnQueryBuilder
+				.filter(Query
+					.of(queryBuilder -> queryBuilder.queryString(queryStringQuerybuilder -> queryStringQuerybuilder
+						.query(getOpenSearchQueryString(filterExpression)))))
+				.field("embedding")
+				.k(topK)
+				.vector(embedding))))
+			.minScore(similarityThreshold)
+			.build();
+	}
+
+	private org.opensearch.client.opensearch.core.SearchRequest buildExactQuery(float[] embedding, int topK,
+			double similarityThreshold, Filter.Expression filterExpression) {
+		return new org.opensearch.client.opensearch.core.SearchRequest.Builder()
+			.query(buildExactQuery(embedding, filterExpression))
 			.index(this.index)
 			.sort(sortOptionsBuilder -> sortOptionsBuilder
 				.score(scoreSortBuilder -> scoreSortBuilder.order(SortOrder.Desc)))
 			.size(topK)
 			.minScore(similarityThreshold)
-			.build());
+			.build();
 	}
 
-	private Query getOpenSearchSimilarityQuery(float[] embedding, Filter.Expression filterExpression) {
+	private Query buildExactQuery(float[] embedding, Filter.Expression filterExpression) {
 		return Query.of(queryBuilder -> queryBuilder.scriptScore(scriptScoreQueryBuilder -> {
 			scriptScoreQueryBuilder
 				.query(queryBuilder2 -> queryBuilder2.queryString(queryStringQuerybuilder -> queryStringQuerybuilder
@@ -358,8 +385,41 @@ private CreateIndexResponse createIndexMapping(String index, String mappingJson)
 
 	@Override
 	public void afterPropertiesSet() {
+		/**
+		 * Generates a JSON string for the k-NN vector mapping configuration. The
+		 * knn_vector field allows k-NN vectors ingestion into OpenSearch and supports
+		 * various k-NN searches.
+		 * @see <a href=
+		 * "https://opensearch.org/docs/latest/search-plugins/knn/knn-index#method-definitions">OpenSearch
+		 * k-NN Method Definitions</a>
+		 */
 		if (this.initializeSchema && !exists(this.index)) {
-			createIndexMapping(this.index, String.format(this.mappingJson, this.embeddingModel.dimensions()));
+			String finalMappingJson;
+			if (this.useApproximateKnn
+					&& this.mappingJson.equals(DEFAULT_MAPPING_EMBEDDING_TYPE_KNN_VECTOR_DIMENSION)) {
+				// Generate approximate k-NN mapping with HNSW method
+				finalMappingJson = """
+						{
+							"properties": {
+								"embedding": {
+									"type": "knn_vector",
+									"dimension": %d,
+									"method": {
+										"name": "hnsw",
+										"engine": "lucene",
+										"space_type": "%s"
+									}
+								}
+							}
+						}
+						""".formatted(this.dimensions > 0 ? this.dimensions : this.embeddingModel.dimensions(),
+						this.similarityFunction);
+			}
+			else {
+				// Use provided mapping or default exact k-NN mapping
+				finalMappingJson = String.format(this.mappingJson, this.embeddingModel.dimensions());
+			}
+			createIndexMapping(this.index, finalMappingJson);
 		}
 	}
 
@@ -417,6 +477,10 @@ public static class Builder extends AbstractVectorStoreBuilder<Builder> {
 
 		private String similarityFunction = COSINE_SIMILARITY_FUNCTION;
 
+		private boolean useApproximateKnn = false;
+
+		private int dimensions = 1536;
+
 		/**
 		 * Sets the OpenSearch client.
 		 * @param openSearchClient The OpenSearch client to use
@@ -488,6 +552,39 @@ public Builder similarityFunction(String similarityFunction) {
 			return this;
 		}
 
+		/**
+		 * Sets whether to use approximate k-NN search. If true, the approximate k-NN
+		 * method is used for faster searches and maintains good performance even at large
+		 * scales. If false, the exact brute-force k-NN method is used for precise and
+		 * highly accurate searches.
+		 * @param useApproximateKnn true to use approximate k-NN, false for exact k-NN
+		 * @return The builder instance
+		 * @see <a href=
+		 * "https://opensearch.org/docs/latest/search-plugins/knn/approximate-knn/">Approximate
+		 * k-NN</a>
+		 * @see <a href=
+		 * "https://opensearch.org/docs/latest/search-plugins/knn/knn-score-script/">Exact
+		 * k-NN with scoring script</a>
+		 */
+		public Builder useApproximateKnn(boolean useApproximateKnn) {
+			this.useApproximateKnn = useApproximateKnn;
+			return this;
+		}
+
+		/**
+		 * Sets the number of dimensions for the vector embeddings. This is used when
+		 * creating the index mapping for approximate k-NN. If not set, defaults to 1536
+		 * or uses the embedding model's dimensions.
+		 * @param dimensions The number of dimensions
+		 * @return The builder instance
+		 * @throws IllegalArgumentException if dimensions is less than or equal to 0
+		 */
+		public Builder dimensions(int dimensions) {
+			Assert.isTrue(dimensions > 0, "dimensions must be greater than 0");
+			this.dimensions = dimensions;
+			return this;
+		}
+
 		/**
 		 * Builds a new OpenSearchVectorStore instance with the configured properties.
 		 * @return A new OpenSearchVectorStore instance