Skip to content

Commit 5455ee9

Browse files
stIncMalestIncMale
authored
Introduce builder for $vectorSearch aggregation stage (#1200)
JAVA-5117
1 parent e477257 commit 5455ee9

File tree

18 files changed

+720
-13
lines changed

18 files changed

+720
-13
lines changed

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

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,17 @@
1717
package com.mongodb.client.model;
1818

1919
import com.mongodb.MongoNamespace;
20+
import com.mongodb.annotations.Beta;
2021
import com.mongodb.client.model.densify.DensifyOptions;
2122
import com.mongodb.client.model.densify.DensifyRange;
2223
import com.mongodb.client.model.fill.FillOptions;
2324
import com.mongodb.client.model.fill.FillOutputField;
2425
import com.mongodb.client.model.geojson.Point;
26+
import com.mongodb.client.model.search.FieldSearchPath;
2527
import com.mongodb.client.model.search.SearchCollector;
2628
import com.mongodb.client.model.search.SearchOperator;
2729
import com.mongodb.client.model.search.SearchOptions;
30+
import com.mongodb.client.model.search.VectorSearchOptions;
2831
import com.mongodb.lang.Nullable;
2932
import org.bson.BsonArray;
3033
import org.bson.BsonBoolean;
@@ -34,6 +37,7 @@
3437
import org.bson.BsonString;
3538
import org.bson.BsonType;
3639
import org.bson.BsonValue;
40+
import org.bson.Document;
3741
import org.bson.codecs.configuration.CodecRegistry;
3842
import org.bson.conversions.Bson;
3943

@@ -47,6 +51,7 @@
4751
import static com.mongodb.client.model.GeoNearOptions.geoNearOptions;
4852
import static com.mongodb.client.model.densify.DensifyOptions.densifyOptions;
4953
import static com.mongodb.client.model.search.SearchOptions.searchOptions;
54+
import static com.mongodb.client.model.search.VectorSearchOptions.vectorSearchOptions;
5055
import static com.mongodb.internal.Iterables.concat;
5156
import static com.mongodb.internal.client.model.Util.sizeAtLeast;
5257
import static java.util.Arrays.asList;
@@ -933,6 +938,90 @@ public static Bson searchMeta(final SearchCollector collector, final SearchOptio
933938
return new SearchStage("$searchMeta", notNull("collector", collector), notNull("options", options));
934939
}
935940

941+
/**
942+
* Creates a {@code $vectorSearch} pipeline stage supported by MongoDB Atlas.
943+
* You may use the {@code $meta: "vectorSearchScore"} expression, e.g., via {@link Projections#metaVectorSearchScore(String)},
944+
* to extract the relevance score assigned to each found document.
945+
*
946+
* @param queryVector The query vector. The number of dimensions must match that of the {@code index}.
947+
* @param path The field to be searched.
948+
* @param index The name of the index to use.
949+
* @param numCandidates The number of candidates.
950+
* @param limit The limit on the number of documents produced by the pipeline stage.
951+
* @return The {@code $vectorSearch} pipeline stage.
952+
*
953+
* @mongodb.atlas.manual atlas-vector-search/vector-search-stage/ $vectorSearch
954+
* @mongodb.atlas.manual atlas-search/scoring/ Scoring
955+
* @mongodb.server.release 6.0.10
956+
* @since 4.11
957+
*/
958+
@Beta(Beta.Reason.SERVER)
959+
public static Bson vectorSearch(
960+
final FieldSearchPath path,
961+
final Iterable<Double> queryVector,
962+
final String index,
963+
final long numCandidates,
964+
final long limit) {
965+
return vectorSearch(notNull("path", path), notNull("queryVector", queryVector), notNull("index", index), numCandidates, limit,
966+
vectorSearchOptions());
967+
}
968+
969+
/**
970+
* Creates a {@code $vectorSearch} pipeline stage supported by MongoDB Atlas.
971+
* You may use the {@code $meta: "vectorSearchScore"} expression, e.g., via {@link Projections#metaVectorSearchScore(String)},
972+
* to extract the relevance score assigned to each found document.
973+
*
974+
* @param queryVector The query vector. The number of dimensions must match that of the {@code index}.
975+
* @param path The field to be searched.
976+
* @param index The name of the index to use.
977+
* @param numCandidates The number of candidates.
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+
*
982+
* @mongodb.atlas.manual atlas-vector-search/vector-search-stage/ $vectorSearch
983+
* @mongodb.atlas.manual atlas-search/scoring/ Scoring
984+
* @mongodb.server.release 6.0.10
985+
* @since 4.11
986+
*/
987+
@Beta(Beta.Reason.SERVER)
988+
public static Bson vectorSearch(
989+
final FieldSearchPath path,
990+
final Iterable<Double> queryVector,
991+
final String index,
992+
final long numCandidates,
993+
final long limit,
994+
final VectorSearchOptions options) {
995+
notNull("path", path);
996+
notNull("queryVector", queryVector);
997+
notNull("index", index);
998+
notNull("options", options);
999+
return new Bson() {
1000+
@Override
1001+
public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> documentClass, final CodecRegistry codecRegistry) {
1002+
Document specificationDoc = new Document("path", path.toValue())
1003+
.append("queryVector", queryVector)
1004+
.append("index", index)
1005+
.append("numCandidates", numCandidates)
1006+
.append("limit", limit);
1007+
specificationDoc.putAll(options.toBsonDocument(documentClass, codecRegistry));
1008+
return new Document("$vectorSearch", specificationDoc).toBsonDocument(documentClass, codecRegistry);
1009+
}
1010+
1011+
@Override
1012+
public String toString() {
1013+
return "Stage{name=$vectorSearch"
1014+
+ ", path=" + path
1015+
+ ", queryVector=" + queryVector
1016+
+ ", index=" + index
1017+
+ ", numCandidates=" + numCandidates
1018+
+ ", limit=" + limit
1019+
+ ", options=" + options
1020+
+ '}';
1021+
}
1022+
};
1023+
}
1024+
9361025
/**
9371026
* Creates an $unset pipeline stage that removes/excludes fields from documents
9381027
*

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.mongodb.client.model;
1818

19+
import com.mongodb.annotations.Beta;
1920
import com.mongodb.client.model.geojson.Geometry;
2021
import com.mongodb.client.model.geojson.Point;
2122
import com.mongodb.client.model.search.SearchCollector;
@@ -84,11 +85,30 @@ public static <TItem> Bson eq(@Nullable final TItem value) {
8485
* @param <TItem> the value type
8586
* @return the filter
8687
* @mongodb.driver.manual reference/operator/query/eq $eq
88+
* @see #eqFull(String, Object)
8789
*/
8890
public static <TItem> Bson eq(final String fieldName, @Nullable final TItem value) {
8991
return new SimpleEncodingFilter<>(fieldName, value);
9092
}
9193

94+
/**
95+
* Creates a filter that matches all documents where the value of the field name equals the specified value.
96+
* Unlike {@link #eq(String, Object)}, this method creates a full form of {@code $eq}.
97+
* This method exists temporarily until Atlas starts supporting the short form of {@code $eq}.
98+
* It will likely be removed in the next driver release.
99+
*
100+
* @param fieldName the field name
101+
* @param value the value, which may be null
102+
* @param <TItem> the value type
103+
* @return the filter
104+
* @mongodb.driver.manual reference/operator/query/eq $eq
105+
* @since 4.11
106+
*/
107+
@Beta(Beta.Reason.SERVER)
108+
public static <TItem> Bson eqFull(final String fieldName, @Nullable final TItem value) {
109+
return new OperatorFilter<>("$eq", fieldName, value);
110+
}
111+
92112
/**
93113
* Creates a filter that matches all documents where the value of the field name does not equal the specified value.
94114
*

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,13 @@
1616

1717
package com.mongodb.client.model;
1818

19+
import com.mongodb.annotations.Beta;
20+
import com.mongodb.client.model.search.FieldSearchPath;
1921
import com.mongodb.client.model.search.SearchCollector;
2022
import com.mongodb.client.model.search.SearchCount;
2123
import com.mongodb.client.model.search.SearchOperator;
2224
import com.mongodb.client.model.search.SearchOptions;
25+
import com.mongodb.client.model.search.VectorSearchOptions;
2326
import org.bson.BsonArray;
2427
import org.bson.BsonDocument;
2528
import org.bson.BsonInt32;
@@ -163,6 +166,7 @@ public static Bson elemMatch(final String fieldName, final Bson filter) {
163166
* @since 4.1
164167
* @see #metaTextScore(String)
165168
* @see #metaSearchScore(String)
169+
* @see #metaVectorSearchScore(String)
166170
* @see #metaSearchHighlights(String)
167171
*/
168172
public static Bson meta(final String fieldName, final String metaFieldName) {
@@ -196,6 +200,22 @@ public static Bson metaSearchScore(final String fieldName) {
196200
return meta(fieldName, "searchScore");
197201
}
198202

203+
/**
204+
* Creates a projection to the given field name of the vectorSearchScore,
205+
* for use with {@link Aggregates#vectorSearch(FieldSearchPath, Iterable, String, long, long, VectorSearchOptions)}.
206+
* Calling this method is equivalent to calling {@link #meta(String, String)} with {@code "vectorSearchScore"} as the second argument.
207+
*
208+
* @param fieldName the field name
209+
* @return the projection
210+
* @mongodb.atlas.manual atlas-search/scoring/ Scoring
211+
* @mongodb.server.release 6.0.10
212+
* @since 4.11
213+
*/
214+
@Beta(Beta.Reason.SERVER)
215+
public static Bson metaVectorSearchScore(final String fieldName) {
216+
return meta(fieldName, "vectorSearchScore");
217+
}
218+
199219
/**
200220
* Creates a projection to the given field name of the searchHighlights,
201221
* for use with {@link Aggregates#search(SearchOperator, SearchOptions)} / {@link Aggregates#search(SearchCollector, SearchOptions)}.

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright 2008-present MongoDB, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.mongodb.client.model.search;
17+
18+
import com.mongodb.annotations.Immutable;
19+
import com.mongodb.internal.client.model.AbstractConstructibleBson;
20+
import org.bson.BsonDocument;
21+
import org.bson.Document;
22+
import org.bson.conversions.Bson;
23+
24+
import static com.mongodb.assertions.Assertions.notNull;
25+
26+
final class VectorSearchConstructibleBson extends AbstractConstructibleBson<VectorSearchConstructibleBson> implements VectorSearchOptions {
27+
/**
28+
* An {@linkplain Immutable immutable} {@link BsonDocument#isEmpty() empty} instance.
29+
*/
30+
static final VectorSearchConstructibleBson EMPTY_IMMUTABLE = new VectorSearchConstructibleBson(AbstractConstructibleBson.EMPTY_IMMUTABLE);
31+
32+
VectorSearchConstructibleBson(final Bson base) {
33+
super(base);
34+
}
35+
36+
private VectorSearchConstructibleBson(final Bson base, final Document appended) {
37+
super(base, appended);
38+
}
39+
40+
@Override
41+
protected VectorSearchConstructibleBson newSelf(final Bson base, final Document appended) {
42+
return new VectorSearchConstructibleBson(base, appended);
43+
}
44+
45+
@Override
46+
public VectorSearchOptions filter(final Bson filter) {
47+
return newAppended("filter", notNull("name", filter));
48+
}
49+
50+
@Override
51+
public VectorSearchOptions option(final String name, final Object value) {
52+
return newAppended(notNull("name", name), notNull("value", value));
53+
}
54+
}

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Copyright 2008-present MongoDB, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.mongodb.client.model.search;
17+
18+
import com.mongodb.annotations.Beta;
19+
import com.mongodb.annotations.Sealed;
20+
import com.mongodb.client.model.Aggregates;
21+
import com.mongodb.client.model.Filters;
22+
import org.bson.conversions.Bson;
23+
24+
/**
25+
* Represents optional fields of the {@code $vectorSearch} pipeline stage of an aggregation pipeline.
26+
*
27+
* @see Aggregates#vectorSearch(FieldSearchPath, Iterable, String, long, long, VectorSearchOptions)
28+
* @mongodb.atlas.manual atlas-vector-search/vector-search-stage/ $vectorSearch
29+
* @mongodb.server.release 6.0.10
30+
* @since 4.11
31+
*/
32+
@Sealed
33+
@Beta(Beta.Reason.SERVER)
34+
public interface VectorSearchOptions extends Bson {
35+
/**
36+
* Creates a new {@link VectorSearchOptions} with the filter specified.
37+
*
38+
* @param filter A filter that is applied before applying the
39+
* {@link Aggregates#vectorSearch(FieldSearchPath, Iterable, String, long, long, VectorSearchOptions) queryVector}.
40+
* One may use {@link Filters} to create this filter, though not all filters may be supported.
41+
* See the MongoDB documentation for the list of supported filters.
42+
* <p>
43+
* Note that for now one has to use {@link Filters#eqFull(String, Object)} instead of {@link Filters#eq(String, Object)}.</p>
44+
* @return A new {@link VectorSearchOptions}.
45+
*/
46+
VectorSearchOptions filter(Bson filter);
47+
48+
/**
49+
* Creates a new {@link VectorSearchOptions} with the specified option in situations when there is no builder method
50+
* that better satisfies your needs.
51+
* This method cannot be used to validate the syntax.
52+
* <p>
53+
* <i>Example</i><br>
54+
* The following code creates two functionally equivalent {@link VectorSearchOptions} objects,
55+
* though they may not be {@linkplain Object#equals(Object) equal}.
56+
* <pre>{@code
57+
* VectorSearchOptions options1 = VectorSearchOptions.vectorSearchOptions()
58+
* .filter(Filters.lt("fieldName", 1));
59+
* VectorSearchOptions options2 = VectorSearchOptions.vectorSearchOptions()
60+
* .option("filter", Filters.lt("fieldName", 1));
61+
* }</pre>
62+
*
63+
* @param name The option name.
64+
* @param value The option value.
65+
* @return A new {@link VectorSearchOptions}.
66+
*/
67+
VectorSearchOptions option(String name, Object value);
68+
69+
/**
70+
* Returns {@link VectorSearchOptions} that represents server defaults.
71+
*
72+
* @return {@link VectorSearchOptions} that represents server defaults.
73+
*/
74+
static VectorSearchOptions vectorSearchOptions() {
75+
return VectorSearchConstructibleBson.EMPTY_IMMUTABLE;
76+
}
77+
}

‎driver-core/src/main/com/mongodb/client/model/search/package-info.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
*
2626
* @see com.mongodb.client.model.Aggregates#search(SearchOperator, SearchOptions)
2727
* @see com.mongodb.client.model.Aggregates#search(SearchCollector, SearchOptions)
28+
* @see com.mongodb.client.model.Aggregates#vectorSearch(FieldSearchPath, java.lang.Iterable, java.lang.String, long, long, VectorSearchOptions)
2829
* @mongodb.atlas.manual atlas-search/ Atlas Search
2930
* @mongodb.atlas.manual atlas-search/query-syntax/ Atlas Search aggregation pipeline stages
3031
* @since 4.7

‎driver-core/src/test/functional/com/mongodb/client/model/FiltersFunctionalSpecification.groovy‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ class FiltersFunctionalSpecification extends OperationFunctionalSpecification {
8080
def 'eq'() {
8181
expect:
8282
find(eq('x', 1)) == [a]
83+
find(eq('_id', 2)) == [b]
8384
find(eq(2)) == [b]
8485
}
8586

0 commit comments

Comments
 (0)