Add VectorIndex and SearchIndexDefinition abstraction.

Author: mp911de

spring-data-mongodb/pom.xml

@@ -131,6 +131,13 @@
 			<optional>true</optional>
 		</dependency>
 
+		<dependency>
+			<groupId>org.awaitility</groupId>
+			<artifactId>awaitility</artifactId>
+			<version>4.2.2</version>
+			<scope>test</scope>
+		</dependency>
+
 		<dependency>
 			<groupId>io.reactivex.rxjava3</groupId>
 			<artifactId>rxjava</artifactId>

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java

@@ -185,8 +185,8 @@
  * @author Michael Krog
  * @author Jakub Zurawa
  */
-public class MongoTemplate
-		implements MongoOperations, ApplicationContextAware, IndexOperationsProvider, SearchIndexOperationsProvider, ReadPreferenceAware {
+public class MongoTemplate implements MongoOperations, ApplicationContextAware, IndexOperationsProvider,
+		SearchIndexOperationsProvider, ReadPreferenceAware {
 
 	private static final Log LOGGER = LogFactory.getLog(MongoTemplate.class);
 	private static final WriteResultChecking DEFAULT_WRITE_RESULT_CHECKING = WriteResultChecking.NONE;
@@ -771,6 +771,21 @@ public IndexOperations indexOps(Class<?> entityClass) {
 		return indexOps(getCollectionName(entityClass), entityClass);
 	}
 
+	@Override
+	public SearchIndexOperations searchIndexOps(String collectionName) {
+		return searchIndexOps(null, collectionName);
+	}
+
+	@Override
+	public SearchIndexOperations searchIndexOps(Class<?> type) {
+		return new DefaultSearchIndexOperations(this, type);
+	}
+
+	@Override
+	public SearchIndexOperations searchIndexOps(@Nullable Class<?> type, String collectionName) {
+		return new DefaultSearchIndexOperations(this, collectionName, type);
+	}
+
 	@Override
 	public BulkOperations bulkOps(BulkMode mode, String collectionName) {
 		return bulkOps(mode, null, collectionName);
@@ -1316,7 +1331,7 @@ private WriteConcern potentiallyForceAcknowledgedWrite(@Nullable WriteConcern wc
 
 		if (ObjectUtils.nullSafeEquals(WriteResultChecking.EXCEPTION, writeResultChecking)) {
 			if (wc == null || wc.getWObject() == null
-					|| (wc.getWObject()instanceof Number concern && concern.intValue() < 1)) {
+					|| (wc.getWObject() instanceof Number concern && concern.intValue() < 1)) {
 				return WriteConcern.ACKNOWLEDGED;
 			}
 		}
@@ -1968,7 +1983,8 @@ public <T> List<T> mapReduce(Query query, Class<?> domainType, String inputColle
 			}
 
 			if (mapReduceOptions.getOutputSharded().isPresent()) {
-				MongoCompatibilityAdapter.mapReduceIterableAdapter(mapReduce).sharded(mapReduceOptions.getOutputSharded().get());
+				MongoCompatibilityAdapter.mapReduceIterableAdapter(mapReduce)
+						.sharded(mapReduceOptions.getOutputSharded().get());
 			}
 
 			if (StringUtils.hasText(mapReduceOptions.getOutputCollection()) && !mapReduceOptions.usesInlineOutput()) {
@@ -2067,7 +2083,7 @@ public <T> List<T> findAllAndRemove(Query query, Class<T> entityClass, String co
 	}
 
 	@Override
-	public <T> UpdateResult replace(Query query, T replacement, ReplaceOptions options, String collectionName){
+	public <T> UpdateResult replace(Query query, T replacement, ReplaceOptions options, String collectionName) {
 
 		Assert.notNull(replacement, "Replacement must not be null");
 		return replace(query, (Class<T>) ClassUtils.getUserClass(replacement), replacement, options, collectionName);
@@ -2743,8 +2759,7 @@ protected <T> T doFindAndModify(CollectionPreparer collectionPreparer, String co
 			LOGGER.debug(String.format(
 					"findAndModify using query: %s fields: %s sort: %s for class: %s and update: %s in collection: %s",
 					serializeToJsonSafely(mappedQuery), fields, serializeToJsonSafely(sort), entityClass,
-					serializeToJsonSafely(mappedUpdate),
-					collectionName));
+					serializeToJsonSafely(mappedUpdate), collectionName));
 		}
 
 		return executeFindOneInternal(
@@ -3013,21 +3028,6 @@ static RuntimeException potentiallyConvertRuntimeException(RuntimeException ex,
 		return resolved == null ? ex : resolved;
 	}
 
-	@Override
-	public SearchIndexOperations searchIndexOps(String collectionName) {
-		return searchIndexOps(null, collectionName);
-	}
-
-	@Override
-	public SearchIndexOperations searchIndexOps(Class<?> type) {
-		return new DefaultSearchIndexOperations(this, type);
-	}
-
-	@Override
-	public SearchIndexOperations searchIndexOps(Class<?> type, String collectionName) {
-		return new DefaultSearchIndexOperations(this, collectionName, type);
-	}
-
 	// Callback implementations
 
 	/**

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/QueryMapper.java

@@ -1130,7 +1130,7 @@ public Class<?> getFieldType() {
 	 * @author Oliver Gierke
 	 * @author Thomas Darimont
 	 */
-	protected static class MetadataBackedField extends Field {
+	public static class MetadataBackedField extends Field {
 
 		private static final Pattern POSITIONAL_PARAMETER_PATTERN = Pattern.compile("\\.\\$(\\[.*?\\])?");
 		private static final Pattern NUMERIC_SEGMENT = Pattern.compile("\\d+");

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/ABetterNameForVector.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2014-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core.index;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * {@link ABetterNameForVector} marks a property as Vector Index Field that contains the actual vector.
+ *
+ * @author Mark Paluch
+ * @since 3.5
+ */
+@Documented
+@Target({ ElementType.FIELD, ElementType.METHOD, ElementType.ANNOTATION_TYPE })
+@Retention(RetentionPolicy.RUNTIME)
+public @interface ABetterNameForVector {
+
+	/**
+	 * @return number of dimensions used for the vector index.
+	 */
+	int dimensions();
+
+	/**
+	 * @return similarity function used for distance computation.
+	 */
+	VectorIndex.SimilarityFunction similarity();
+
+	/**
+	 * @return quantization strategy used for the vector index. Defaults to none expecting {@code float} vectors.
+	 */
+	VectorIndex.Quantization quantization() default VectorIndex.Quantization.NONE;
+
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/ABetterNameForVectorIndexed.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2014-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core.index;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * {@link ABetterNameForVectorIndexed} marks a property for inclusion in the Vector Index for filtering.
+ *
+ * @author Mark Paluch
+ * @since 3.5
+ */
+@Documented
+@Target({ ElementType.FIELD, ElementType.METHOD, ElementType.ANNOTATION_TYPE })
+@Retention(RetentionPolicy.RUNTIME)
+public @interface ABetterNameForVectorIndexed {
+
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/DefaultSearchIndexOperations.java

@@ -21,33 +21,46 @@
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.bson.Document;
-import org.springframework.data.mongodb.core.DefaultIndexOperations;
+
+import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.aggregation.Aggregation;
 import org.springframework.data.mongodb.core.aggregation.AggregationResults;
-import org.springframework.data.mongodb.core.convert.QueryMapper;
-import org.springframework.data.mongodb.core.index.SearchIndex.Filter;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
-import org.springframework.lang.NonNull;
+import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
+import org.springframework.data.util.TypeInformation;
 import org.springframework.lang.Nullable;
 
 /**
  * @author Christoph Strobl
+ * @author Mark Paluch
+ * @since 3.5
  */
-public class DefaultSearchIndexOperations extends DefaultIndexOperations implements SearchIndexOperations {
+public class DefaultSearchIndexOperations implements SearchIndexOperations {
+
+	private static final Log LOGGER = LogFactory.getLog(DefaultSearchIndexOperations.class);
 
-	private static final Log LOGGER = LogFactory.getLog(SearchIndexOperations.class);
+	private final MongoOperations mongoOperations;
+	private final String collectionName;
+	private final TypeInformation<?> entityTypeInformation;
 
 	public DefaultSearchIndexOperations(MongoOperations mongoOperations, Class<?> type) {
 		this(mongoOperations, mongoOperations.getCollectionName(type), type);
 	}
 
 	public DefaultSearchIndexOperations(MongoOperations mongoOperations, String collectionName, @Nullable Class<?> type) {
-		super(mongoOperations, collectionName, type);
-	}
+		this.collectionName = collectionName;
+
+		if (type != null) {
+
+			MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext = mongoOperations
+					.getConverter().getMappingContext();
+			entityTypeInformation = mappingContext.getRequiredPersistentEntity(type).getTypeInformation();
+		} else {
+			entityTypeInformation = null;
+		}
 
-	private static String getMappedPath(String path, MongoPersistentEntity<?> entity, QueryMapper mapper) {
-		return mapper.getMappedFields(new Document(path, 1), entity).entrySet().iterator().next().getKey();
+		this.mongoOperations = mongoOperations;
 	}
 
 	@Override
@@ -62,11 +75,10 @@ public boolean exists(String indexName) {
 	}
 
 	@Override
-	public void updateIndex(SearchIndex index) {
-
-		MongoPersistentEntity<?> entity = lookupPersistentEntity(type, collectionName);
+	public void updateIndex(SearchIndexDefinition index) {
 
-		Document indexDocument = createIndexDocument(index, entity);
+		Document indexDocument = index.getIndexDocument(entityTypeInformation,
+				mongoOperations.getConverter().getMappingContext());
 
 		Document cmdResult = mongoOperations.execute(db -> {
 
@@ -108,57 +120,40 @@ public List<IndexInfo> getIndexInfo() {
 	@Override
 	public String ensureIndex(SearchIndexDefinition indexDefinition) {
 
-		if (!(indexDefinition instanceof SearchIndex vsi)) {
+		if (!(indexDefinition instanceof VectorIndex vsi)) {
 			throw new IllegalStateException("Index definitions must be of type VectorIndex");
 		}
 
-		MongoPersistentEntity<?> entity = lookupPersistentEntity(type, collectionName);
-
-		Document index = createIndexDocument(vsi, entity);
+		Document index = indexDefinition.getIndexDocument(entityTypeInformation,
+				mongoOperations.getConverter().getMappingContext());
 
 		Document cmdResult = mongoOperations.execute(db -> {
 
 			Document command = new Document().append("createSearchIndexes", collectionName).append("indexes", List.of(index));
+
 			if (LOGGER.isDebugEnabled()) {
-				LOGGER.debug("Creating VectorIndex: db.runCommand(%s)".formatted(command.toJson()));
+				LOGGER.debug("Creating SearchIndex: db.runCommand(%s)".formatted(command.toJson()));
 			}
+
 			return db.runCommand(command);
 		});
 
 		return cmdResult.get("ok").toString().equalsIgnoreCase("1.0") ? vsi.getName() : cmdResult.toJson();
 	}
 
-	@NonNull
-	private Document createIndexDocument(SearchIndex vsi, MongoPersistentEntity<?> entity) {
-
-		Document index = new Document(vsi.getIndexOptions());
-		Document definition = new Document();
-
-		List<Document> fields = new ArrayList<>(vsi.getFilters().size() + 1);
-
-		Document vectorField = new Document("type", "vector");
-		vectorField.append("path", getMappedPath(vsi.getPath(), entity, mapper));
-		vectorField.append("numDimensions", vsi.getDimensions());
-		vectorField.append("similarity", vsi.getSimilarity());
-
-		fields.add(vectorField);
-
-		for (Filter filter : vsi.getFilters()) {
-			fields.add(new Document("type", "filter").append("path", getMappedPath(filter.path(), entity, mapper)));
-		}
-
-		definition.append("fields", fields);
-		index.append("definition", definition);
-		return index;
+	@Override
+	public void dropAllIndexes() {
+		getIndexInfo().forEach(indexInfo -> dropIndex(indexInfo.getName()));
 	}
 
 	@Override
 	public void dropIndex(String name) {
 
 		Document command = new Document().append("dropSearchIndex", collectionName).append("name", name);
 		if (LOGGER.isDebugEnabled()) {
-			LOGGER.debug("Dropping VectorIndex: db.runCommand(%s)".formatted(command.toJson()));
+			LOGGER.debug("Dropping SearchIndex: db.runCommand(%s)".formatted(command.toJson()));
 		}
 		mongoOperations.execute(db -> db.runCommand(command));
 	}
+
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/IndexResolver.java

@@ -78,4 +78,14 @@ default Iterable<? extends IndexDefinition> resolveIndexFor(Class<?> entityType)
 		return resolveIndexFor(TypeInformation.of(entityType));
 	}
 
+	/**
+	 * Find and create {@link SearchIndexDefinition}s for properties of given {@link TypeInformation}.
+	 * {@link SearchIndexDefinition}s are created for properties and types with {@link VectorIndexed}.
+	 *
+	 * @param typeInformation must not be {@literal null}.
+	 * @return Empty {@link Iterable} in case no {@link IndexDefinition} could be resolved for type.
+	 * @since 3.5
+	 */
+	Iterable<? extends SearchIndexDefinition> resolveSearchIndexFor(TypeInformation<?> typeInformation);
+
 }