Skip to content

Commit a50345f

Browse files
rozzastrogiyotec
andauthored
Add typed builder API for vector search index definitions (#1960)
Introduced `SearchIndexDefinition` sealed interface with factory methods, `VectorSearchIndexFields` (`vectorField`, `filterField`, `autoEmbedField` builders), `HnswSearchIndexOptions` for HNSW-specific parameters, and `VectorSearchIndexDefinition` as the concrete implementation. Added Scala companion objects for `SearchIndexDefinition` and `VectorSearchIndexFields` with delegating factory methods, and a `SearchIndexModel.apply` overload for `VectorSearchIndexDefinition`. Annotated mutable builders with `@NotThreadSafe`. Added fail-fast null-element validation using notNullElements in factory methods. Includes unit tests (Java) and spec tests (Scala). JAVA-6112 JAVA-6099 --------- Co-authored-by: Almas Abdrazak <almas337519@gmail.com>
1 parent a1cd54a commit a50345f

12 files changed

Lines changed: 1484 additions & 0 deletions

File tree

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
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+
17+
package com.mongodb.client.model;
18+
19+
import com.mongodb.annotations.NotThreadSafe;
20+
import com.mongodb.lang.Nullable;
21+
import org.bson.BsonDocument;
22+
import org.bson.BsonInt32;
23+
import org.bson.codecs.configuration.CodecRegistry;
24+
import org.bson.conversions.Bson;
25+
26+
import static com.mongodb.assertions.Assertions.isTrueArgument;
27+
28+
/**
29+
* Options for the HNSW (Hierarchical Navigable Small World) indexing method in a vector search index.
30+
*
31+
* <p>This class provides a fluent builder for specifying HNSW-specific parameters when creating
32+
* a vector search index with {@code indexingMethod("hnsw")}.</p>
33+
*
34+
* <p>Since {@link VectorSearchIndexFields.VectorField#hnswOptions(Bson)} accepts any {@link Bson},
35+
* a raw {@link org.bson.Document} may also be passed directly for forward compatibility.</p>
36+
*
37+
* <pre>{@code
38+
* vectorField("embedding")
39+
* .indexingMethod("hnsw")
40+
* .hnswOptions(new HnswSearchIndexOptions().maxEdges(16).numEdgeCandidates(200))
41+
* }</pre>
42+
*
43+
* @see VectorSearchIndexFields.VectorField#hnswOptions(Bson)
44+
* @since 5.8
45+
*/
46+
@NotThreadSafe
47+
public final class HnswSearchIndexOptions implements Bson {
48+
@Nullable
49+
private Integer maxEdges;
50+
@Nullable
51+
private Integer numEdgeCandidates;
52+
53+
/**
54+
* Creates a new instance with default settings.
55+
*
56+
* @since 5.8
57+
*/
58+
public HnswSearchIndexOptions() {
59+
}
60+
61+
/**
62+
* Sets the maximum number of connected neighbors for each node in the HNSW graph.
63+
*
64+
* @param maxEdges the maximum number of edges (connected neighbors)
65+
* @return this
66+
* @since 5.8
67+
*/
68+
public HnswSearchIndexOptions maxEdges(final int maxEdges) {
69+
isTrueArgument("maxEdges > 0", maxEdges > 0);
70+
this.maxEdges = maxEdges;
71+
return this;
72+
}
73+
74+
/**
75+
* Sets the number of nearest neighbor candidates to consider when building the HNSW graph.
76+
*
77+
* @param numEdgeCandidates the number of nearest neighbor candidates
78+
* @return this
79+
* @since 5.8
80+
*/
81+
public HnswSearchIndexOptions numEdgeCandidates(final int numEdgeCandidates) {
82+
isTrueArgument("numEdgeCandidates > 0", numEdgeCandidates > 0);
83+
this.numEdgeCandidates = numEdgeCandidates;
84+
return this;
85+
}
86+
87+
@Override
88+
public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> documentClass, final CodecRegistry codecRegistry) {
89+
BsonDocument doc = new BsonDocument();
90+
if (maxEdges != null) {
91+
doc.append("maxEdges", new BsonInt32(maxEdges));
92+
}
93+
if (numEdgeCandidates != null) {
94+
doc.append("numEdgeCandidates", new BsonInt32(numEdgeCandidates));
95+
}
96+
return doc;
97+
}
98+
99+
@Override
100+
public String toString() {
101+
return "HnswSearchIndexOptions{"
102+
+ "maxEdges=" + maxEdges
103+
+ ", numEdgeCandidates=" + numEdgeCandidates
104+
+ '}';
105+
}
106+
}
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
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+
17+
package com.mongodb.client.model;
18+
19+
import com.mongodb.annotations.Sealed;
20+
import org.bson.conversions.Bson;
21+
22+
import java.util.List;
23+
24+
import static com.mongodb.assertions.Assertions.isTrueArgument;
25+
import static com.mongodb.assertions.Assertions.notNull;
26+
import static com.mongodb.assertions.Assertions.notNullElements;
27+
import static java.util.Arrays.asList;
28+
29+
/**
30+
* A definition for an Atlas Search index.
31+
*
32+
* <p>This interface provides factory methods for creating search index definitions
33+
* that can be passed to {@link SearchIndexModel}.</p>
34+
*
35+
* @see SearchIndexModel
36+
* @see VectorSearchIndexDefinition
37+
* @since 5.8
38+
*/
39+
@Sealed
40+
public interface SearchIndexDefinition extends Bson {
41+
42+
/**
43+
* Creates a vector search index definition with the specified fields.
44+
*
45+
* <p>The resulting definition produces a document of the form {@code {"fields": [...]}},
46+
* suitable for use with {@link SearchIndexType#vectorSearch()}.</p>
47+
*
48+
* @param fields the fields for the vector search index. Each field should be created using
49+
* {@link VectorSearchIndexFields} factory methods, or may be a raw {@link Bson} document.
50+
* @return a new {@link VectorSearchIndexDefinition}
51+
* @see VectorSearchIndexFields#vectorField(String)
52+
* @see VectorSearchIndexFields#filterField(String)
53+
* @see VectorSearchIndexFields#autoEmbedField(String)
54+
* @since 5.8
55+
*/
56+
static VectorSearchIndexDefinition vectorSearch(final Bson... fields) {
57+
List<Bson> fieldList = asList(notNull("fields", fields));
58+
isTrueArgument("fields must not be empty", !fieldList.isEmpty());
59+
notNullElements("fields", fieldList);
60+
return new VectorSearchIndexDefinition(fieldList);
61+
}
62+
63+
/**
64+
* Creates a vector search index definition with the specified fields.
65+
*
66+
* <p>The resulting definition produces a document of the form {@code {"fields": [...]}},
67+
* suitable for use with {@link SearchIndexType#vectorSearch()}.</p>
68+
*
69+
* @param fields the fields for the vector search index. Each field should be created using
70+
* {@link VectorSearchIndexFields} factory methods, or may be a raw {@link Bson} document.
71+
* @return a new {@link VectorSearchIndexDefinition}
72+
* @see VectorSearchIndexFields#vectorField(String)
73+
* @see VectorSearchIndexFields#filterField(String)
74+
* @see VectorSearchIndexFields#autoEmbedField(String)
75+
* @since 5.8
76+
*/
77+
static VectorSearchIndexDefinition vectorSearch(final List<? extends Bson> fields) {
78+
notNullElements("fields", fields);
79+
isTrueArgument("fields must not be empty", !fields.isEmpty());
80+
return new VectorSearchIndexDefinition(fields);
81+
}
82+
}

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,14 @@
2424
/**
2525
* A model describing the creation of a single Atlas Search index.
2626
*
27+
* <p>The {@code definition} parameter accepts any {@link org.bson.conversions.Bson} instance.
28+
* For vector search indexes, use the builders provided by {@link SearchIndexDefinition#vectorSearch(Bson...)}
29+
* and {@link VectorSearchIndexFields} to construct the definition, and pass it to the
30+
* {@linkplain #SearchIndexModel(String, VectorSearchIndexDefinition) vector search constructor}
31+
* which automatically sets the index type to {@link SearchIndexType#vectorSearch()}.</p>
32+
*
33+
* @see SearchIndexDefinition
34+
* @see VectorSearchIndexFields
2735
* @since 4.11
2836
* @mongodb.server.release 6.0
2937
*/
@@ -42,6 +50,7 @@ public final class SearchIndexModel {
4250
* will be used to create the search index.</p>
4351
*
4452
* @param definition the search index mapping definition.
53+
* @see SearchIndexDefinition#vectorSearch(Bson...)
4554
*/
4655
public SearchIndexModel(final Bson definition) {
4756
this(null, definition, null);
@@ -52,17 +61,33 @@ public SearchIndexModel(final Bson definition) {
5261
*
5362
* @param name the search index name.
5463
* @param definition the search index mapping definition.
64+
* @see SearchIndexDefinition#vectorSearch(Bson...)
5565
*/
5666
public SearchIndexModel(final String name, final Bson definition) {
5767
this(name, definition, null);
5868
}
5969

70+
/**
71+
* Construct a vector search index instance with the given name and definition.
72+
*
73+
* <p>The index type is automatically set to {@link SearchIndexType#vectorSearch()}.</p>
74+
*
75+
* @param name the search index name.
76+
* @param definition the vector search index definition.
77+
* @see SearchIndexDefinition#vectorSearch(Bson...)
78+
* @since 5.8
79+
*/
80+
public SearchIndexModel(final String name, final VectorSearchIndexDefinition definition) {
81+
this(name, definition, SearchIndexType.vectorSearch());
82+
}
83+
6084
/**
6185
* Construct an instance with the given Atlas Search name, index definition, and type.
6286
*
6387
* @param name the search index name.
6488
* @param definition the search index mapping definition.
6589
* @param type the search index type.
90+
* @see SearchIndexDefinition#vectorSearch(Bson...)
6691
* @since 5.2
6792
*/
6893
public SearchIndexModel(@Nullable final String name, final Bson definition, @Nullable final SearchIndexType type) {
Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
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+
17+
package com.mongodb.client.model;
18+
19+
import org.bson.BsonArray;
20+
import org.bson.BsonDocument;
21+
import org.bson.codecs.configuration.CodecRegistry;
22+
import org.bson.conversions.Bson;
23+
24+
import java.util.ArrayList;
25+
import java.util.List;
26+
27+
import static com.mongodb.assertions.Assertions.doesNotContainNull;
28+
import static com.mongodb.assertions.Assertions.notNull;
29+
30+
/**
31+
* A vector search index definition, producing a document of the form {@code {"fields": [...]}}.
32+
*
33+
* <p>Instances are created via {@link SearchIndexDefinition#vectorSearch(Bson...)}.</p>
34+
*
35+
* @see SearchIndexDefinition
36+
* @see SearchIndexType#vectorSearch()
37+
* @since 5.8
38+
*/
39+
public final class VectorSearchIndexDefinition implements SearchIndexDefinition {
40+
private final List<? extends Bson> fields;
41+
42+
VectorSearchIndexDefinition(final List<? extends Bson> fields) {
43+
doesNotContainNull("fields", notNull("fields", fields));
44+
this.fields = new ArrayList<>(fields);
45+
}
46+
47+
@Override
48+
public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> documentClass, final CodecRegistry codecRegistry) {
49+
BsonArray fieldArray = new BsonArray();
50+
for (Bson field : fields) {
51+
fieldArray.add(field.toBsonDocument(documentClass, codecRegistry));
52+
}
53+
return new BsonDocument("fields", fieldArray);
54+
}
55+
56+
@Override
57+
public String toString() {
58+
return "VectorSearchIndexDefinition{"
59+
+ "fields=" + fields
60+
+ '}';
61+
}
62+
}

0 commit comments

Comments
 (0)