Skip to content

Commit 042b64f

Browse files
rozzarozza
committed
Add vectorSearch operator for $search pipeline stage
JAVA-6130
1 parent c0f9627 commit 042b64f

7 files changed

Lines changed: 356 additions & 1 deletion

File tree

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -633,6 +633,57 @@ static RegexSearchOperator regex(final Iterable<? extends SearchPath> paths, fin
633633
.append("query", queryIterator.hasNext() ? queries : firstQuery));
634634
}
635635

636+
/**
637+
* Returns a {@link SearchOperator} that performs vector search within the {@code $search} pipeline stage.
638+
* This is the approximate (ANN) variant with {@code numCandidates}.
639+
*
640+
* @param path The indexed vector field to search.
641+
* @param queryVector The query vector. The number of dimensions must match the index field.
642+
* @param limit The number of results to return.
643+
* @param numCandidates The number of nearest neighbors to consider during ANN search.
644+
* Must be greater than or equal to {@code limit}. The server may impose an upper bound.
645+
* @return The requested {@link VectorSearchOperator}.
646+
* @mongodb.atlas.manual atlas-search/vector-search/ vectorSearch operator
647+
* @since 5.8
648+
*/
649+
static VectorSearchOperator vectorSearch(
650+
final FieldSearchPath path,
651+
final Iterable<Double> queryVector,
652+
final int limit,
653+
final int numCandidates) {
654+
notNull("path", path);
655+
notNull("queryVector", queryVector);
656+
isTrueArgument("numCandidates must be >= limit", numCandidates >= limit);
657+
return new VectorSearchOperatorConstructibleBsonElement("vectorSearch",
658+
new Document("path", path.toValue())
659+
.append("queryVector", queryVector)
660+
.append("limit", limit)
661+
.append("numCandidates", numCandidates));
662+
}
663+
664+
/**
665+
* Returns a {@link SearchOperator} that performs exact (ENN) vector search within the {@code $search} pipeline stage.
666+
*
667+
* @param path The indexed vector field to search.
668+
* @param queryVector The query vector. The number of dimensions must match the index field.
669+
* @param limit The number of results to return.
670+
* @return The requested {@link VectorSearchOperator}.
671+
* @mongodb.atlas.manual atlas-search/vector-search/ vectorSearch operator
672+
* @since 5.8
673+
*/
674+
static VectorSearchOperator vectorSearchExact(
675+
final FieldSearchPath path,
676+
final Iterable<Double> queryVector,
677+
final int limit) {
678+
notNull("path", path);
679+
notNull("queryVector", queryVector);
680+
return new VectorSearchOperatorConstructibleBsonElement("vectorSearch",
681+
new Document("path", path.toValue())
682+
.append("queryVector", queryVector)
683+
.append("limit", limit)
684+
.append("exact", true));
685+
}
686+
636687
/**
637688
* Creates a {@link SearchOperator} from a {@link Bson} in situations when there is no builder method that better satisfies your needs.
638689
* This method cannot be used to validate the syntax.

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
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.Reason;
20+
import com.mongodb.annotations.Sealed;
21+
22+
/**
23+
* A {@link SearchOperator} that performs vector search within the {@code $search} pipeline stage.
24+
*
25+
* @mongodb.atlas.manual atlas-search/operators-and-collectors/#operators Search operators
26+
* @since 5.8
27+
*/
28+
@Sealed
29+
@Beta(Reason.CLIENT)
30+
public interface VectorSearchOperator extends SearchOperator {
31+
32+
/**
33+
* Creates a new {@link VectorSearchOperator} with the exact search option.
34+
*
35+
* @param exact Whether to use exact (ENN) search. If {@code true}, runs exact nearest neighbor search.
36+
* @return A new {@link VectorSearchOperator}.
37+
*/
38+
VectorSearchOperator exact(boolean exact);
39+
40+
/**
41+
* Creates a new {@link VectorSearchOperator} with the filter specified.
42+
*
43+
* @param filter A search operator to filter documents.
44+
* @return A new {@link VectorSearchOperator}.
45+
*/
46+
VectorSearchOperator filter(SearchOperator filter);
47+
48+
/**
49+
* Creates a new {@link VectorSearchOperator} with the scoring modifier specified.
50+
*
51+
* @param modifier The scoring modifier.
52+
* @return A new {@link VectorSearchOperator}.
53+
*/
54+
@Override
55+
VectorSearchOperator score(SearchScore modifier);
56+
}

‎driver-core/src/main/com/mongodb/client/model/search/VectorSearchOperatorConstructibleBsonElement.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.internal.client.model.AbstractConstructibleBsonElement;
19+
import org.bson.conversions.Bson;
20+
21+
import static com.mongodb.assertions.Assertions.notNull;
22+
23+
final class VectorSearchOperatorConstructibleBsonElement
24+
extends AbstractConstructibleBsonElement<VectorSearchOperatorConstructibleBsonElement>
25+
implements VectorSearchOperator {
26+
27+
VectorSearchOperatorConstructibleBsonElement(final String name, final Bson value) {
28+
super(name, value);
29+
}
30+
31+
private VectorSearchOperatorConstructibleBsonElement(final Bson baseElement, final Bson appendedElementValue) {
32+
super(baseElement, appendedElementValue);
33+
}
34+
35+
@Override
36+
protected VectorSearchOperatorConstructibleBsonElement newSelf(final Bson baseElement, final Bson appendedElementValue) {
37+
return new VectorSearchOperatorConstructibleBsonElement(baseElement, appendedElementValue);
38+
}
39+
40+
@Override
41+
public VectorSearchOperator exact(final boolean exact) {
42+
return newWithAppendedValue("exact", exact);
43+
}
44+
45+
@Override
46+
public VectorSearchOperator filter(final SearchOperator filter) {
47+
return newWithAppendedValue("filter", notNull("filter", filter));
48+
}
49+
50+
@Override
51+
public VectorSearchOperatorConstructibleBsonElement score(final SearchScore modifier) {
52+
return newWithAppendedValue("score", notNull("modifier", modifier));
53+
}
54+
}

‎driver-core/src/test/unit/com/mongodb/client/model/search/SearchOperatorTest.java‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
package com.mongodb.client.model.search;
1717

1818
import com.mongodb.MongoClientSettings;
19+
import com.mongodb.client.model.Aggregates;
1920
import com.mongodb.client.model.geojson.Point;
2021
import com.mongodb.client.model.geojson.Position;
2122
import org.bson.BsonArray;
@@ -1002,6 +1003,125 @@ void regex() {
10021003
);
10031004
}
10041005

1006+
@Test
1007+
void vectorSearch() {
1008+
assertAll(
1009+
() -> assertThrows(IllegalArgumentException.class, () ->
1010+
// path must not be null
1011+
SearchOperator.vectorSearch(null, asList(1.0), 10, 50)
1012+
),
1013+
() -> assertThrows(IllegalArgumentException.class, () ->
1014+
// queryVector must not be null
1015+
SearchOperator.vectorSearch(fieldPath("embedding"), null, 10, 50)
1016+
),
1017+
() -> assertThrows(IllegalArgumentException.class, () ->
1018+
// numCandidates must be >= limit
1019+
SearchOperator.vectorSearch(fieldPath("embedding"), asList(1.0), 100, 50)
1020+
),
1021+
() -> assertEquals(
1022+
new BsonDocument("vectorSearch",
1023+
new BsonDocument("path", new BsonString("embedding"))
1024+
.append("queryVector", new BsonArray(asList(
1025+
new BsonDouble(1.0), new BsonDouble(2.0), new BsonDouble(3.0))))
1026+
.append("limit", new BsonInt32(10))
1027+
.append("numCandidates", new BsonInt32(100))),
1028+
SearchOperator.vectorSearch(
1029+
fieldPath("embedding"),
1030+
asList(1.0, 2.0, 3.0),
1031+
10,
1032+
100
1033+
).toBsonDocument()
1034+
),
1035+
() -> assertEquals(
1036+
new BsonDocument("vectorSearch",
1037+
new BsonDocument("path", new BsonString("embedding"))
1038+
.append("queryVector", new BsonArray(asList(
1039+
new BsonDouble(1.0), new BsonDouble(2.0))))
1040+
.append("limit", new BsonInt32(10))
1041+
.append("numCandidates", new BsonInt32(50))
1042+
.append("filter", new BsonDocument("text",
1043+
new BsonDocument("query", new BsonString("hello"))
1044+
.append("path", new BsonString("title"))))
1045+
.append("score", new BsonDocument("boost",
1046+
new BsonDocument("value", new BsonDouble(2.0))))),
1047+
SearchOperator.vectorSearch(
1048+
fieldPath("embedding"),
1049+
asList(1.0, 2.0),
1050+
10,
1051+
50
1052+
).filter(SearchOperator.text(fieldPath("title"), "hello"))
1053+
.score(boost(2f))
1054+
.toBsonDocument()
1055+
)
1056+
);
1057+
}
1058+
1059+
@Test
1060+
void vectorSearchExact() {
1061+
assertAll(
1062+
() -> assertThrows(IllegalArgumentException.class, () ->
1063+
// path must not be null
1064+
SearchOperator.vectorSearchExact(null, asList(1.0), 10)
1065+
),
1066+
() -> assertThrows(IllegalArgumentException.class, () ->
1067+
// queryVector must not be null
1068+
SearchOperator.vectorSearchExact(fieldPath("embedding"), null, 10)
1069+
),
1070+
() -> assertEquals(
1071+
new BsonDocument("vectorSearch",
1072+
new BsonDocument("path", new BsonString("embedding"))
1073+
.append("queryVector", new BsonArray(asList(
1074+
new BsonDouble(1.0), new BsonDouble(2.0), new BsonDouble(3.0))))
1075+
.append("limit", new BsonInt32(5))
1076+
.append("exact", BsonBoolean.TRUE)),
1077+
SearchOperator.vectorSearchExact(
1078+
fieldPath("embedding"),
1079+
asList(1.0, 2.0, 3.0),
1080+
5
1081+
).toBsonDocument()
1082+
),
1083+
() -> assertEquals(
1084+
new BsonDocument("vectorSearch",
1085+
new BsonDocument("path", new BsonString("embedding"))
1086+
.append("queryVector", new BsonArray(asList(
1087+
new BsonDouble(1.0), new BsonDouble(2.0))))
1088+
.append("limit", new BsonInt32(10))
1089+
.append("numCandidates", new BsonInt32(50))
1090+
.append("exact", BsonBoolean.FALSE)),
1091+
SearchOperator.vectorSearch(
1092+
fieldPath("embedding"),
1093+
asList(1.0, 2.0),
1094+
10,
1095+
50
1096+
).exact(false)
1097+
.toBsonDocument()
1098+
)
1099+
);
1100+
}
1101+
1102+
@Test
1103+
void vectorSearchInsideSearchStage() {
1104+
assertEquals(
1105+
new BsonDocument("$search",
1106+
new BsonDocument("index", new BsonString("myIndex"))
1107+
.append("vectorSearch",
1108+
new BsonDocument("path", new BsonString("embedding"))
1109+
.append("queryVector", new BsonArray(asList(
1110+
new BsonDouble(1.0), new BsonDouble(2.0), new BsonDouble(3.0))))
1111+
.append("limit", new BsonInt32(10))
1112+
.append("numCandidates", new BsonInt32(100)))),
1113+
Aggregates.search(
1114+
SearchOperator.vectorSearch(
1115+
fieldPath("embedding"),
1116+
asList(1.0, 2.0, 3.0),
1117+
10,
1118+
100
1119+
),
1120+
SearchOptions.searchOptions().index("myIndex")
1121+
).toBsonDocument()
1122+
);
1123+
}
1124+
10051125
private static SearchOperator docExamplePredefined() {
10061126
return SearchOperator.exists(
10071127
fieldPath("fieldName"));

‎driver-scala/src/main/scala/org/mongodb/scala/model/search/SearchOperator.scala‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -495,6 +495,42 @@ object SearchOperator {
495495
def regex(paths: Iterable[_ <: SearchPath], queries: Iterable[String]): RegexSearchOperator =
496496
JSearchOperator.regex(paths.asJava, queries.asJava)
497497

498+
/**
499+
* Returns a `SearchOperator` that performs vector search within the `\$search` pipeline stage.
500+
* This is the approximate (ANN) variant with `numCandidates`.
501+
*
502+
* @param path The indexed vector field to search.
503+
* @param queryVector The query vector. The number of dimensions must match the index field.
504+
* @param limit The number of results to return.
505+
* @param numCandidates The number of nearest neighbors to consider during ANN search.
506+
* Must be greater than or equal to `limit`. The server may impose an upper bound.
507+
* @return The requested `VectorSearchOperator`.
508+
* @see [[https://www.mongodb.com/docs/atlas/atlas-search/vector-search/ vectorSearch operator]]
509+
* @since 5.8
510+
*/
511+
def vectorSearch(
512+
path: FieldSearchPath,
513+
queryVector: Iterable[Double],
514+
limit: Int,
515+
numCandidates: Int): VectorSearchOperator =
516+
JSearchOperator.vectorSearch(path, queryVector.map(Double.box).asJava, limit, numCandidates)
517+
518+
/**
519+
* Returns a `SearchOperator` that performs exact (ENN) vector search within the `\$search` pipeline stage.
520+
*
521+
* @param path The indexed vector field to search.
522+
* @param queryVector The query vector. The number of dimensions must match the index field.
523+
* @param limit The number of results to return.
524+
* @return The requested `VectorSearchOperator`.
525+
* @see [[https://www.mongodb.com/docs/atlas/atlas-search/vector-search/ vectorSearch operator]]
526+
* @since 5.8
527+
*/
528+
def vectorSearchExact(
529+
path: FieldSearchPath,
530+
queryVector: Iterable[Double],
531+
limit: Int): VectorSearchOperator =
532+
JSearchOperator.vectorSearchExact(path, queryVector.map(Double.box).asJava, limit)
533+
498534
/**
499535
* Creates a `SearchOperator` from a `Bson` in situations when there is no builder method that better satisfies your needs.
500536
* This method cannot be used to validate the syntax.

‎driver-scala/src/main/scala/org/mongodb/scala/model/search/package.scala‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,16 @@ package object search {
234234
@Beta(Array(Reason.CLIENT))
235235
type QueryStringSearchOperator = com.mongodb.client.model.search.QueryStringSearchOperator
236236

237+
/**
238+
* A `SearchOperator` that performs vector search within the `\$search` pipeline stage.
239+
*
240+
* @see `SearchOperator.vectorSearch`
241+
* @since 5.8
242+
*/
243+
@Sealed
244+
@Beta(Array(Reason.CLIENT, Reason.SERVER))
245+
type VectorSearchOperator = com.mongodb.client.model.search.VectorSearchOperator
246+
237247
/**
238248
* Fuzzy search options that may be used with some [[SearchOperator]]s.
239249
*

0 commit comments

Comments
 (0)