Skip to content

Add support for vector search with BSON Vector #1549

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 13 commits into from
Oct 31, 2024
2 changes: 1 addition & 1 deletion .evergreen/run-atlas-search-tests.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
91 changes: 70 additions & 21 deletions driver-core/src/main/com/mongodb/client/model/Aggregates.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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);
}

/**
Expand Down Expand Up @@ -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() {
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading