Add support for vector search with BSON Vector (#1549)

JAVA-5650
---------

Co-authored-by: Valentin Kovalenko <[email protected]>

Author: vbabanin

.evergreen/run-atlas-search-tests.sh

@@ -16,4 +16,4 @@ echo "Running Atlas Search tests"
 ./gradlew --stacktrace --info \
     -Dorg.mongodb.test.atlas.search=true \
     -Dorg.mongodb.test.uri=${MONGODB_URI} \
-    driver-core:test --tests AggregatesSearchIntegrationTest
+    driver-core:test --tests AggregatesSearchIntegrationTest --tests AggregatesVectorSearchIntegrationTest

driver-core/src/main/com/mongodb/client/model/Aggregates.java

@@ -37,6 +37,7 @@
 import org.bson.BsonType;
 import org.bson.BsonValue;
 import org.bson.Document;
+import org.bson.Vector;
 import org.bson.codecs.configuration.CodecRegistry;
 import org.bson.conversions.Bson;
 
@@ -963,28 +964,37 @@ public static Bson vectorSearch(
         notNull("queryVector", queryVector);
         notNull("index", index);
         notNull("options", options);
-        return new Bson() {
-            @Override
-            public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> documentClass, final CodecRegistry codecRegistry) {
-                Document specificationDoc = new Document("path", path.toValue())
-                        .append("queryVector", queryVector)
-                        .append("index", index)
-                        .append("limit", limit);
-                specificationDoc.putAll(options.toBsonDocument(documentClass, codecRegistry));
-                return new Document("$vectorSearch", specificationDoc).toBsonDocument(documentClass, codecRegistry);
-            }
+        return new VectorSearchBson(path, queryVector, index, limit, options);
+    }
 
-            @Override
-            public String toString() {
-                return "Stage{name=$vectorSearch"
-                        + ", path=" + path
-                        + ", queryVector=" + queryVector
-                        + ", index=" + index
-                        + ", limit=" + limit
-                        + ", options=" + options
-                        + '}';
-            }
-        };
+    /**
+     * Creates a {@code $vectorSearch} pipeline stage supported by MongoDB Atlas.
+     * You may use the {@code $meta: "vectorSearchScore"} expression, e.g., via {@link Projections#metaVectorSearchScore(String)},
+     * to extract the relevance score assigned to each found document.
+     *
+     * @param queryVector The {@linkplain Vector query vector}. The number of dimensions must match that of the {@code index}.
+     * @param path        The field to be searched.
+     * @param index       The name of the index to use.
+     * @param limit       The limit on the number of documents produced by the pipeline stage.
+     * @param options     Optional {@code $vectorSearch} pipeline stage fields.
+     * @return The {@code $vectorSearch} pipeline stage.
+     * @mongodb.atlas.manual atlas-vector-search/vector-search-stage/ $vectorSearch
+     * @mongodb.atlas.manual atlas-search/scoring/ Scoring
+     * @mongodb.server.release 6.0
+     * @see Vector
+     * @since 5.3
+     */
+    public static Bson vectorSearch(
+            final FieldSearchPath path,
+            final Vector queryVector,
+            final String index,
+            final long limit,
+            final VectorSearchOptions options) {
+        notNull("path", path);
+        notNull("queryVector", queryVector);
+        notNull("index", index);
+        notNull("options", options);
+        return new VectorSearchBson(path, queryVector, index, limit, options);
     }
 
     /**
@@ -2145,6 +2155,45 @@ public String toString() {
         }
     }
 
+    private static class VectorSearchBson implements Bson {
+        private final FieldSearchPath path;
+        private final Object queryVector;
+        private final String index;
+        private final long limit;
+        private final VectorSearchOptions options;
+
+        VectorSearchBson(final FieldSearchPath path, final Object queryVector,
+                                final String index, final long limit,
+                                final VectorSearchOptions options) {
+            this.path = path;
+            this.queryVector = queryVector;
+            this.index = index;
+            this.limit = limit;
+            this.options = options;
+        }
+
+        @Override
+        public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> documentClass, final CodecRegistry codecRegistry) {
+            Document specificationDoc = new Document("path", path.toValue())
+                    .append("queryVector", queryVector)
+                    .append("index", index)
+                    .append("limit", limit);
+            specificationDoc.putAll(options.toBsonDocument(documentClass, codecRegistry));
+            return new Document("$vectorSearch", specificationDoc).toBsonDocument(documentClass, codecRegistry);
+        }
+
+        @Override
+        public String toString() {
+            return "Stage{name=$vectorSearch"
+                    + ", path=" + path
+                    + ", queryVector=" + queryVector
+                    + ", index=" + index
+                    + ", limit=" + limit
+                    + ", options=" + options
+                    + '}';
+        }
+    }
+
     private Aggregates() {
     }
 }

driver-core/src/main/com/mongodb/internal/operation/CreateSearchIndexesOperation.java

@@ -32,11 +32,11 @@
  *
  * <p>This class is not part of the public API and may be removed or changed at any time</p>
  */
-final class CreateSearchIndexesOperation extends AbstractWriteSearchIndexOperation {
+public final class CreateSearchIndexesOperation extends AbstractWriteSearchIndexOperation {
     private static final String COMMAND_NAME = "createSearchIndexes";
     private final List<SearchIndexRequest> indexRequests;
 
-    CreateSearchIndexesOperation(final MongoNamespace namespace, final List<SearchIndexRequest> indexRequests) {
+    public CreateSearchIndexesOperation(final MongoNamespace namespace, final List<SearchIndexRequest> indexRequests) {
         super(namespace);
         this.indexRequests = assertNotNull(indexRequests);
     }

driver-core/src/main/com/mongodb/internal/operation/ListSearchIndexesOperation.java

@@ -42,7 +42,7 @@
  *
  * <p>This class is not part of the public API and may be removed or changed at any time</p>
  */
-final class ListSearchIndexesOperation<T>
+public final class ListSearchIndexesOperation<T>
         implements AsyncExplainableReadOperation<AsyncBatchCursor<T>>, ExplainableReadOperation<BatchCursor<T>> {
     private static final String STAGE_LIST_SEARCH_INDEXES = "$listSearchIndexes";
     private final MongoNamespace namespace;
@@ -59,9 +59,10 @@ final class ListSearchIndexesOperation<T>
     private final String indexName;
     private final boolean retryReads;
 
-    ListSearchIndexesOperation(final MongoNamespace namespace, final Decoder<T> decoder, @Nullable final String indexName,
-            @Nullable final Integer batchSize, @Nullable final Collation collation, @Nullable final BsonValue comment,
-            @Nullable final Boolean allowDiskUse, final boolean retryReads) {
+    public ListSearchIndexesOperation(final MongoNamespace namespace, final Decoder<T> decoder, @Nullable final String indexName,
+                                      @Nullable final Integer batchSize, @Nullable final Collation collation,
+                                      @Nullable final BsonValue comment,
+                                      @Nullable final Boolean allowDiskUse, final boolean retryReads) {
         this.namespace = namespace;
         this.decoder = decoder;
         this.allowDiskUse = allowDiskUse;

driver-core/src/main/com/mongodb/internal/operation/SearchIndexRequest.java

@@ -31,14 +31,15 @@
  *
  * <p>This class is not part of the public API and may be removed or changed at any time</p>
  */
-final class SearchIndexRequest {
+public final class SearchIndexRequest {
     private final BsonDocument definition;
     @Nullable
     private final String indexName;
     @Nullable
     private final SearchIndexType searchIndexType;
 
-    SearchIndexRequest(final BsonDocument definition, @Nullable final String indexName, @Nullable final SearchIndexType searchIndexType) {
+    public SearchIndexRequest(final BsonDocument definition, @Nullable final String indexName,
+                              @Nullable final SearchIndexType searchIndexType) {
         assertNotNull(definition);
         this.definition = definition;
         this.indexName = indexName;