-
Notifications
You must be signed in to change notification settings - Fork 1.5k
Add vectorSearch operator for $search pipeline stage #1962
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
base: main
Are you sure you want to change the base?
Changes from all commits
042b64f
45914b1
34f9c32
cbcb01c
4245fbb
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| /* | ||
| * Copyright 2008-present MongoDB, Inc. | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package com.mongodb.client.model.search; | ||
|
|
||
| import com.mongodb.annotations.Beta; | ||
| import com.mongodb.annotations.Reason; | ||
| import com.mongodb.annotations.Sealed; | ||
|
|
||
| /** | ||
| * A {@link SearchOperator} that performs vector search within the {@code $search} pipeline stage. | ||
| * | ||
| * @mongodb.atlas.manual atlas-search/operators-and-collectors/#operators Search operators | ||
| * @since 5.8 | ||
| */ | ||
| @Sealed | ||
| @Beta(Reason.CLIENT) | ||
| public interface VectorSearchOperator extends SearchOperator { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. what about Should we add a fluent index builder that can set
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The index work was covered in #1960 - this is for creating the query eg: Which can be done via: I added This follows the existing |
||
|
|
||
| /** | ||
| * Creates a new {@link VectorSearchOperator} with the filter specified. | ||
| * | ||
| * @param filter A search operator to filter documents. | ||
| * @return A new {@link VectorSearchOperator}. | ||
| */ | ||
| VectorSearchOperator filter(SearchOperator filter); | ||
|
|
||
| /** | ||
| * Creates a new {@link VectorSearchOperator} with the scoring modifier specified. | ||
| * | ||
| * @param modifier The scoring modifier. | ||
| * @return A new {@link VectorSearchOperator}. | ||
| */ | ||
| @Override | ||
| VectorSearchOperator score(SearchScore modifier); | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,49 @@ | ||
| /* | ||
| * Copyright 2008-present MongoDB, Inc. | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package com.mongodb.client.model.search; | ||
|
|
||
| import com.mongodb.internal.client.model.AbstractConstructibleBsonElement; | ||
| import org.bson.conversions.Bson; | ||
|
|
||
| import static com.mongodb.assertions.Assertions.notNull; | ||
|
|
||
| final class VectorSearchOperatorConstructibleBsonElement | ||
| extends AbstractConstructibleBsonElement<VectorSearchOperatorConstructibleBsonElement> | ||
| implements VectorSearchOperator { | ||
|
|
||
| VectorSearchOperatorConstructibleBsonElement(final String name, final Bson value) { | ||
| super(name, value); | ||
| } | ||
|
|
||
| private VectorSearchOperatorConstructibleBsonElement(final Bson baseElement, final Bson appendedElementValue) { | ||
| super(baseElement, appendedElementValue); | ||
| } | ||
|
|
||
| @Override | ||
| protected VectorSearchOperatorConstructibleBsonElement newSelf(final Bson baseElement, final Bson appendedElementValue) { | ||
| return new VectorSearchOperatorConstructibleBsonElement(baseElement, appendedElementValue); | ||
| } | ||
|
|
||
| @Override | ||
| public VectorSearchOperator filter(final SearchOperator filter) { | ||
| return newWithAppendedValue("filter", notNull("filter", filter)); | ||
| } | ||
|
|
||
| @Override | ||
| public VectorSearchOperatorConstructibleBsonElement score(final SearchScore modifier) { | ||
| return newWithAppendedValue("score", notNull("modifier", modifier)); | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -16,6 +16,7 @@ | |
| package com.mongodb.client.model.search; | ||
|
|
||
| import com.mongodb.MongoClientSettings; | ||
| import com.mongodb.client.model.Aggregates; | ||
| import com.mongodb.client.model.geojson.Point; | ||
| import com.mongodb.client.model.geojson.Position; | ||
| import org.bson.BsonArray; | ||
|
|
@@ -1002,6 +1003,129 @@ void regex() { | |
| ); | ||
| } | ||
|
|
||
| @Test | ||
| void vectorSearch() { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. one of the acceptance criteria from the doc is to
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We don't have e2e testing yet at a drivers level, I'm expecting this work flow should be convered by a Drivers rather than the doc. I'll make sure to add a DRIVERS ticket to ensure its documented in the Specs repo and / or unified tests are added. |
||
| assertAll( | ||
| () -> assertThrows(IllegalArgumentException.class, () -> | ||
| // path must not be null | ||
| SearchOperator.vectorSearch(null, asList(1.0), 10, 50) | ||
| ), | ||
| () -> assertThrows(IllegalArgumentException.class, () -> | ||
| // queryVector must not be null | ||
| SearchOperator.vectorSearch(fieldPath("embedding"), null, 10, 50) | ||
| ), | ||
| () -> assertThrows(IllegalArgumentException.class, () -> | ||
| // numCandidates must be >= limit | ||
| SearchOperator.vectorSearch(fieldPath("embedding"), asList(1.0), 100, 50) | ||
| ), | ||
| () -> assertEquals( | ||
| new BsonDocument("vectorSearch", | ||
| new BsonDocument("path", new BsonString("embedding")) | ||
| .append("queryVector", new BsonArray(asList( | ||
| new BsonDouble(1.0), new BsonDouble(2.0), new BsonDouble(3.0)))) | ||
| .append("limit", new BsonInt32(10)) | ||
| .append("numCandidates", new BsonInt32(100))), | ||
| SearchOperator.vectorSearch( | ||
| fieldPath("embedding"), | ||
| asList(1.0, 2.0, 3.0), | ||
| 10, | ||
| 100 | ||
| ).toBsonDocument() | ||
| ), | ||
| () -> assertEquals( | ||
| new BsonDocument("vectorSearch", | ||
| new BsonDocument("path", new BsonString("embedding")) | ||
| .append("queryVector", new BsonArray(asList( | ||
| new BsonDouble(1.0), new BsonDouble(2.0)))) | ||
| .append("limit", new BsonInt32(10)) | ||
| .append("numCandidates", new BsonInt32(50)) | ||
| .append("filter", new BsonDocument("text", | ||
| new BsonDocument("query", new BsonString("hello")) | ||
| .append("path", new BsonString("title")))) | ||
| .append("score", new BsonDocument("boost", | ||
| new BsonDocument("value", new BsonDouble(2.0))))), | ||
| SearchOperator.vectorSearch( | ||
| fieldPath("embedding"), | ||
| asList(1.0, 2.0), | ||
| 10, | ||
| 50 | ||
| ).filter(SearchOperator.text(fieldPath("title"), "hello")) | ||
| .score(boost(2f)) | ||
| .toBsonDocument() | ||
| ) | ||
| ); | ||
| } | ||
|
|
||
| @Test | ||
| void vectorSearchExact() { | ||
| assertAll( | ||
| () -> assertThrows(IllegalArgumentException.class, () -> | ||
| // path must not be null | ||
| SearchOperator.vectorSearchExact(null, asList(1.0), 10) | ||
| ), | ||
| () -> assertThrows(IllegalArgumentException.class, () -> | ||
| // queryVector must not be null | ||
| SearchOperator.vectorSearchExact(fieldPath("embedding"), null, 10) | ||
| ), | ||
| () -> assertEquals( | ||
| new BsonDocument("vectorSearch", | ||
| new BsonDocument("path", new BsonString("embedding")) | ||
| .append("queryVector", new BsonArray(asList( | ||
| new BsonDouble(1.0), new BsonDouble(2.0), new BsonDouble(3.0)))) | ||
| .append("limit", new BsonInt32(5)) | ||
| .append("exact", BsonBoolean.TRUE)), | ||
| SearchOperator.vectorSearchExact( | ||
| fieldPath("embedding"), | ||
| asList(1.0, 2.0, 3.0), | ||
| 5 | ||
| ).toBsonDocument() | ||
| ), | ||
| () -> assertEquals( | ||
| new BsonDocument("vectorSearch", | ||
| new BsonDocument("path", new BsonString("embedding")) | ||
| .append("queryVector", new BsonArray(asList( | ||
| new BsonDouble(1.0), new BsonDouble(2.0)))) | ||
| .append("limit", new BsonInt32(10)) | ||
| .append("exact", BsonBoolean.TRUE) | ||
| .append("filter", new BsonDocument("text", | ||
| new BsonDocument("query", new BsonString("hello")) | ||
| .append("path", new BsonString("title")))) | ||
| .append("score", new BsonDocument("boost", | ||
| new BsonDocument("value", new BsonDouble(2.0))))), | ||
| SearchOperator.vectorSearchExact( | ||
| fieldPath("embedding"), | ||
| asList(1.0, 2.0), | ||
| 10 | ||
| ).filter(SearchOperator.text(fieldPath("title"), "hello")) | ||
| .score(boost(2f)) | ||
| .toBsonDocument() | ||
| ) | ||
| ); | ||
| } | ||
|
|
||
| @Test | ||
| void vectorSearchInsideSearchStage() { | ||
| assertEquals( | ||
| new BsonDocument("$search", | ||
| new BsonDocument("index", new BsonString("myIndex")) | ||
| .append("vectorSearch", | ||
| new BsonDocument("path", new BsonString("embedding")) | ||
| .append("queryVector", new BsonArray(asList( | ||
| new BsonDouble(1.0), new BsonDouble(2.0), new BsonDouble(3.0)))) | ||
| .append("limit", new BsonInt32(10)) | ||
| .append("numCandidates", new BsonInt32(100)))), | ||
| Aggregates.search( | ||
| SearchOperator.vectorSearch( | ||
| fieldPath("embedding"), | ||
| asList(1.0, 2.0, 3.0), | ||
| 10, | ||
| 100 | ||
| ), | ||
| SearchOptions.searchOptions().index("myIndex") | ||
| ).toBsonDocument() | ||
| ); | ||
| } | ||
|
|
||
| private static SearchOperator docExamplePredefined() { | ||
| return SearchOperator.exists( | ||
| fieldPath("fieldName")); | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.