-
Notifications
You must be signed in to change notification settings - Fork 1.5k
Add typed builder API for vector search index definitions #1960
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
Changes from 2 commits
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,106 @@ | ||
| /* | ||
| * 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; | ||
|
|
||
| import com.mongodb.annotations.NotThreadSafe; | ||
| import com.mongodb.lang.Nullable; | ||
| import org.bson.BsonDocument; | ||
| import org.bson.BsonInt32; | ||
| import org.bson.codecs.configuration.CodecRegistry; | ||
| import org.bson.conversions.Bson; | ||
|
|
||
| import static com.mongodb.assertions.Assertions.isTrueArgument; | ||
|
|
||
| /** | ||
| * Options for the HNSW (Hierarchical Navigable Small World) indexing method in a vector search index. | ||
| * | ||
| * <p>This class provides a fluent builder for specifying HNSW-specific parameters when creating | ||
| * a vector search index with {@code indexingMethod("hnsw")}.</p> | ||
| * | ||
| * <p>Since {@link VectorSearchIndexFields.VectorField#hnswOptions(Bson)} accepts any {@link Bson}, | ||
| * a raw {@link org.bson.Document} may also be passed directly for forward compatibility.</p> | ||
| * | ||
| * <pre>{@code | ||
| * vectorField("embedding") | ||
| * .indexingMethod("hnsw") | ||
| * .hnswOptions(new HnswSearchIndexOptions().maxEdges(16).numEdgeCandidates(200)) | ||
| * }</pre> | ||
| * | ||
| * @see VectorSearchIndexFields.VectorField#hnswOptions(Bson) | ||
| * @since 5.8 | ||
| */ | ||
| @NotThreadSafe | ||
| public final class HnswSearchIndexOptions implements Bson { | ||
| @Nullable | ||
| private Integer maxEdges; | ||
| @Nullable | ||
| private Integer numEdgeCandidates; | ||
|
|
||
| /** | ||
| * Creates a new instance with default settings. | ||
| * | ||
| * @since 5.8 | ||
| */ | ||
| public HnswSearchIndexOptions() { | ||
| } | ||
|
|
||
| /** | ||
| * Sets the maximum number of connected neighbors for each node in the HNSW graph. | ||
| * | ||
| * @param maxEdges the maximum number of edges (connected neighbors) | ||
| * @return this | ||
| * @since 5.8 | ||
| */ | ||
| public HnswSearchIndexOptions maxEdges(final int maxEdges) { | ||
| isTrueArgument("maxEdges > 0", maxEdges > 0); | ||
| this.maxEdges = maxEdges; | ||
| return this; | ||
| } | ||
|
|
||
| /** | ||
| * Sets the number of nearest neighbor candidates to consider when building the HNSW graph. | ||
| * | ||
| * @param numEdgeCandidates the number of nearest neighbor candidates | ||
| * @return this | ||
| * @since 5.8 | ||
| */ | ||
| public HnswSearchIndexOptions numEdgeCandidates(final int numEdgeCandidates) { | ||
| isTrueArgument("numEdgeCandidates > 0", numEdgeCandidates > 0); | ||
| this.numEdgeCandidates = numEdgeCandidates; | ||
| return this; | ||
| } | ||
|
|
||
| @Override | ||
| public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> documentClass, final CodecRegistry codecRegistry) { | ||
| BsonDocument doc = new BsonDocument(); | ||
| if (maxEdges != null) { | ||
| doc.append("maxEdges", new BsonInt32(maxEdges)); | ||
| } | ||
| if (numEdgeCandidates != null) { | ||
| doc.append("numEdgeCandidates", new BsonInt32(numEdgeCandidates)); | ||
| } | ||
| return doc; | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return "HnswSearchIndexOptions{" | ||
| + "maxEdges=" + maxEdges | ||
| + ", numEdgeCandidates=" + numEdgeCandidates | ||
| + '}'; | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,82 @@ | ||
| /* | ||
| * 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; | ||
|
|
||
| import com.mongodb.annotations.Sealed; | ||
| import org.bson.conversions.Bson; | ||
|
|
||
| import java.util.List; | ||
|
|
||
| import static com.mongodb.assertions.Assertions.isTrueArgument; | ||
| import static com.mongodb.assertions.Assertions.notNull; | ||
| import static com.mongodb.assertions.Assertions.notNullElements; | ||
| import static java.util.Arrays.asList; | ||
|
|
||
| /** | ||
| * A definition for an Atlas Search index. | ||
| * | ||
| * <p>This interface provides factory methods for creating search index definitions | ||
| * that can be passed to {@link SearchIndexModel}.</p> | ||
| * | ||
| * @see SearchIndexModel | ||
| * @see VectorSearchIndexDefinition | ||
| * @since 5.8 | ||
| */ | ||
| @Sealed | ||
| public interface SearchIndexDefinition extends Bson { | ||
|
|
||
| /** | ||
| * Creates a vector search index definition with the specified fields. | ||
| * | ||
| * <p>The resulting definition produces a document of the form {@code {"fields": [...]}}, | ||
| * suitable for use with {@link SearchIndexType#vectorSearch()}.</p> | ||
| * | ||
| * @param fields the fields for the vector search index. Each field should be created using | ||
| * {@link VectorSearchIndexFields} factory methods, or may be a raw {@link Bson} document. | ||
| * @return a new {@link VectorSearchIndexDefinition} | ||
| * @see VectorSearchIndexFields#vectorField(String) | ||
| * @see VectorSearchIndexFields#filterField(String) | ||
| * @see VectorSearchIndexFields#autoEmbedField(String) | ||
| * @since 5.8 | ||
| */ | ||
| static VectorSearchIndexDefinition vectorSearch(final Bson... fields) { | ||
| List<Bson> fieldList = asList(notNull("fields", fields)); | ||
| isTrueArgument("fields must not be empty", !fieldList.isEmpty()); | ||
| notNullElements("fields", fieldList); | ||
| return new VectorSearchIndexDefinition(fieldList); | ||
| } | ||
|
|
||
| /** | ||
| * Creates a vector search index definition with the specified fields. | ||
| * | ||
| * <p>The resulting definition produces a document of the form {@code {"fields": [...]}}, | ||
| * suitable for use with {@link SearchIndexType#vectorSearch()}.</p> | ||
| * | ||
| * @param fields the fields for the vector search index. Each field should be created using | ||
| * {@link VectorSearchIndexFields} factory methods, or may be a raw {@link Bson} document. | ||
| * @return a new {@link VectorSearchIndexDefinition} | ||
| * @see VectorSearchIndexFields#vectorField(String) | ||
| * @see VectorSearchIndexFields#filterField(String) | ||
| * @see VectorSearchIndexFields#autoEmbedField(String) | ||
| * @since 5.8 | ||
| */ | ||
| static VectorSearchIndexDefinition vectorSearch(final List<? extends Bson> fields) { | ||
| notNullElements("fields", fields); | ||
| isTrueArgument("fields must not be empty", !fields.isEmpty()); | ||
| return new VectorSearchIndexDefinition(fields); | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,62 @@ | ||
| /* | ||
| * 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; | ||
|
|
||
| import org.bson.BsonArray; | ||
| import org.bson.BsonDocument; | ||
| import org.bson.codecs.configuration.CodecRegistry; | ||
| import org.bson.conversions.Bson; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
|
|
||
| import static com.mongodb.assertions.Assertions.doesNotContainNull; | ||
| import static com.mongodb.assertions.Assertions.notNull; | ||
|
|
||
| /** | ||
| * A vector search index definition, producing a document of the form {@code {"fields": [...]}}. | ||
| * | ||
| * <p>Instances are created via {@link SearchIndexDefinition#vectorSearch(Bson...)}.</p> | ||
| * | ||
| * @see SearchIndexDefinition | ||
| * @see SearchIndexType#vectorSearch() | ||
| * @since 5.8 | ||
| */ | ||
| public final class VectorSearchIndexDefinition implements SearchIndexDefinition { | ||
| private final List<? extends Bson> fields; | ||
|
|
||
| VectorSearchIndexDefinition(final List<? extends Bson> fields) { | ||
| doesNotContainNull("fields", notNull("fields", fields)); | ||
| this.fields = new ArrayList<>(fields); | ||
|
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. NIT: does this list have to be mutable ? I can only see it being used in
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. List.of will be usable in the future once we drop Java 8 :) |
||
| } | ||
|
rozza marked this conversation as resolved.
|
||
|
|
||
| @Override | ||
| public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> documentClass, final CodecRegistry codecRegistry) { | ||
| BsonArray fieldArray = new BsonArray(); | ||
|
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. NIT: we already know the size for this array , let's use another constructor |
||
| for (Bson field : fields) { | ||
| fieldArray.add(field.toBsonDocument(documentClass, codecRegistry)); | ||
| } | ||
| return new BsonDocument("fields", fieldArray); | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return "VectorSearchIndexDefinition{" | ||
| + "fields=" + fields | ||
| + '}'; | ||
| } | ||
| } | ||
Uh oh!
There was an error while loading. Please reload this page.