Skip to content

Commit 38fc561

Browse files
vbabaninstIncMale
vbabanin
and
stIncMale
authored
Add support for vector search with BSON Vector (#1549)
JAVA-5650 --------- Co-authored-by: Valentin Kovalenko <valentin.male.kovalenko@gmail.com>
1 parent 3d61254 commit 38fc561

File tree

8 files changed

+476
-34
lines changed

8 files changed

+476
-34
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,4 +16,4 @@ echo "Running Atlas Search tests"
1616
./gradlew --stacktrace --info \
1717
-Dorg.mongodb.test.atlas.search=true \
1818
-Dorg.mongodb.test.uri=${MONGODB_URI} \
19-
driver-core:test --tests AggregatesSearchIntegrationTest
19+
driver-core:test --tests AggregatesSearchIntegrationTest --tests AggregatesVectorSearchIntegrationTest

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

Lines changed: 70 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import org.bson.BsonType;
3838
import org.bson.BsonValue;
3939
import org.bson.Document;
40+
import org.bson.Vector;
4041
import org.bson.codecs.configuration.CodecRegistry;
4142
import org.bson.conversions.Bson;
4243

@@ -963,28 +964,37 @@ public static Bson vectorSearch(
963964
notNull("queryVector", queryVector);
964965
notNull("index", index);
965966
notNull("options", options);
966-
return new Bson() {
967-
@Override
968-
public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> documentClass, final CodecRegistry codecRegistry) {
969-
Document specificationDoc = new Document("path", path.toValue())
970-
.append("queryVector", queryVector)
971-
.append("index", index)
972-
.append("limit", limit);
973-
specificationDoc.putAll(options.toBsonDocument(documentClass, codecRegistry));
974-
return new Document("$vectorSearch", specificationDoc).toBsonDocument(documentClass, codecRegistry);
975-
}
967+
return new VectorSearchBson(path, queryVector, index, limit, options);
968+
}
976969

977-
@Override
978-
public String toString() {
979-
return "Stage{name=$vectorSearch"
980-
+ ", path=" + path
981-
+ ", queryVector=" + queryVector
982-
+ ", index=" + index
983-
+ ", limit=" + limit
984-
+ ", options=" + options
985-
+ '}';
986-
}
987-
};
970+
/**
971+
* Creates a {@code $vectorSearch} pipeline stage supported by MongoDB Atlas.
972+
* You may use the {@code $meta: "vectorSearchScore"} expression, e.g., via {@link Projections#metaVectorSearchScore(String)},
973+
* to extract the relevance score assigned to each found document.
974+
*
975+
* @param queryVector The {@linkplain Vector query vector}. The number of dimensions must match that of the {@code index}.
976+
* @param path The field to be searched.
977+
* @param index The name of the index to use.
978+
* @param limit The limit on the number of documents produced by the pipeline stage.
979+
* @param options Optional {@code $vectorSearch} pipeline stage fields.
980+
* @return The {@code $vectorSearch} pipeline stage.
981+
* @mongodb.atlas.manual atlas-vector-search/vector-search-stage/ $vectorSearch
982+
* @mongodb.atlas.manual atlas-search/scoring/ Scoring
983+
* @mongodb.server.release 6.0
984+
* @see Vector
985+
* @since 5.3
986+
*/
987+
public static Bson vectorSearch(
988+
final FieldSearchPath path,
989+
final Vector queryVector,
990+
final String index,
991+
final long limit,
992+
final VectorSearchOptions options) {
993+
notNull("path", path);
994+
notNull("queryVector", queryVector);
995+
notNull("index", index);
996+
notNull("options", options);
997+
return new VectorSearchBson(path, queryVector, index, limit, options);
988998
}
989999

9901000
/**
@@ -2145,6 +2155,45 @@ public String toString() {
21452155
}
21462156
}
21472157

2158+
private static class VectorSearchBson implements Bson {
2159+
private final FieldSearchPath path;
2160+
private final Object queryVector;
2161+
private final String index;
2162+
private final long limit;
2163+
private final VectorSearchOptions options;
2164+
2165+
VectorSearchBson(final FieldSearchPath path, final Object queryVector,
2166+
final String index, final long limit,
2167+
final VectorSearchOptions options) {
2168+
this.path = path;
2169+
this.queryVector = queryVector;
2170+
this.index = index;
2171+
this.limit = limit;
2172+
this.options = options;
2173+
}
2174+
2175+
@Override
2176+
public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> documentClass, final CodecRegistry codecRegistry) {
2177+
Document specificationDoc = new Document("path", path.toValue())
2178+
.append("queryVector", queryVector)
2179+
.append("index", index)
2180+
.append("limit", limit);
2181+
specificationDoc.putAll(options.toBsonDocument(documentClass, codecRegistry));
2182+
return new Document("$vectorSearch", specificationDoc).toBsonDocument(documentClass, codecRegistry);
2183+
}
2184+
2185+
@Override
2186+
public String toString() {
2187+
return "Stage{name=$vectorSearch"
2188+
+ ", path=" + path
2189+
+ ", queryVector=" + queryVector
2190+
+ ", index=" + index
2191+
+ ", limit=" + limit
2192+
+ ", options=" + options
2193+
+ '}';
2194+
}
2195+
}
2196+
21482197
private Aggregates() {
21492198
}
21502199
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,11 @@
3232
*
3333
* <p>This class is not part of the public API and may be removed or changed at any time</p>
3434
*/
35-
final class CreateSearchIndexesOperation extends AbstractWriteSearchIndexOperation {
35+
public final class CreateSearchIndexesOperation extends AbstractWriteSearchIndexOperation {
3636
private static final String COMMAND_NAME = "createSearchIndexes";
3737
private final List<SearchIndexRequest> indexRequests;
3838

39-
CreateSearchIndexesOperation(final MongoNamespace namespace, final List<SearchIndexRequest> indexRequests) {
39+
public CreateSearchIndexesOperation(final MongoNamespace namespace, final List<SearchIndexRequest> indexRequests) {
4040
super(namespace);
4141
this.indexRequests = assertNotNull(indexRequests);
4242
}

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@
4242
*
4343
* <p>This class is not part of the public API and may be removed or changed at any time</p>
4444
*/
45-
final class ListSearchIndexesOperation<T>
45+
public final class ListSearchIndexesOperation<T>
4646
implements AsyncExplainableReadOperation<AsyncBatchCursor<T>>, ExplainableReadOperation<BatchCursor<T>> {
4747
private static final String STAGE_LIST_SEARCH_INDEXES = "$listSearchIndexes";
4848
private final MongoNamespace namespace;
@@ -59,9 +59,10 @@ final class ListSearchIndexesOperation<T>
5959
private final String indexName;
6060
private final boolean retryReads;
6161

62-
ListSearchIndexesOperation(final MongoNamespace namespace, final Decoder<T> decoder, @Nullable final String indexName,
63-
@Nullable final Integer batchSize, @Nullable final Collation collation, @Nullable final BsonValue comment,
64-
@Nullable final Boolean allowDiskUse, final boolean retryReads) {
62+
public ListSearchIndexesOperation(final MongoNamespace namespace, final Decoder<T> decoder, @Nullable final String indexName,
63+
@Nullable final Integer batchSize, @Nullable final Collation collation,
64+
@Nullable final BsonValue comment,
65+
@Nullable final Boolean allowDiskUse, final boolean retryReads) {
6566
this.namespace = namespace;
6667
this.decoder = decoder;
6768
this.allowDiskUse = allowDiskUse;

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,14 +31,15 @@
3131
*
3232
* <p>This class is not part of the public API and may be removed or changed at any time</p>
3333
*/
34-
final class SearchIndexRequest {
34+
public final class SearchIndexRequest {
3535
private final BsonDocument definition;
3636
@Nullable
3737
private final String indexName;
3838
@Nullable
3939
private final SearchIndexType searchIndexType;
4040

41-
SearchIndexRequest(final BsonDocument definition, @Nullable final String indexName, @Nullable final SearchIndexType searchIndexType) {
41+
public SearchIndexRequest(final BsonDocument definition, @Nullable final String indexName,
42+
@Nullable final SearchIndexType searchIndexType) {
4243
assertNotNull(definition);
4344
this.definition = definition;
4445
this.indexName = indexName;

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy