Skip to content

Commit 12ef20d

Browse files
LielNagarLielNagar
committed
RDBC-921: add missing code for vector search API
1 parent ccdd8de commit 12ef20d

20 files changed

+1428
-4
lines changed

‎src/main/java/net/ravendb/client/documents/indexes/AbstractGenericIndexCreationTask.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ public abstract class AbstractGenericIndexCreationTask extends AbstractIndexCrea
2121
protected final Set<String> indexSuggestions;
2222
protected final Map<String, FieldTermVector> termVectorsStrings;
2323
protected final Map<String, SpatialOptions> spatialOptionsStrings;
24-
protected final Map<String, VectorFieldOptions> vectorOptionsStrings;
24+
protected Map<String, VectorFieldOptions> vectorOptionsStrings;
2525

2626
protected String outputReduceToCollection;
2727
protected String patternForOutputReduceToCollectionReferences;

‎src/main/java/net/ravendb/client/documents/indexes/AbstractIndexCreationTask.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ public IndexDefinition createIndexDefinition() {
3535
indexDefinitionBuilder.setSuggestionsOptions(indexSuggestions);
3636
indexDefinitionBuilder.setTermVectorsStrings(termVectorsStrings);
3737
indexDefinitionBuilder.setSpatialIndexesStrings(spatialOptionsStrings);
38+
indexDefinitionBuilder.setVectorOptionsStrings(vectorOptionsStrings);
3839
indexDefinitionBuilder.setOutputReduceToCollection(outputReduceToCollection);
3940
indexDefinitionBuilder.setPatternForOutputReduceToCollectionReferences(patternForOutputReduceToCollectionReferences);
4041
indexDefinitionBuilder.setPatternReferencesCollectionName(patternReferencesCollectionName);
@@ -50,6 +51,9 @@ public IndexDefinition createIndexDefinition() {
5051

5152
if (searchEngineType != null) {
5253
indexDefinitionBuilder.getConfiguration().put(Constants.Configuration.Indexes.INDEXING_STATIC_SEARCH_ENGINE_TYPE, SharpEnum.value(searchEngineType));
54+
} else if (vectorOptionsStrings != null && !vectorOptionsStrings.isEmpty()) {
55+
indexDefinitionBuilder.getConfiguration()
56+
.put(Constants.Configuration.Indexes.INDEXING_STATIC_SEARCH_ENGINE_TYPE, "Corax");
5357
}
5458

5559
return indexDefinitionBuilder.toIndexDefinition(conventions);

‎src/main/java/net/ravendb/client/documents/indexes/AbstractIndexDefinitionBuilder.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ public abstract class AbstractIndexDefinitionBuilder<TIndexDefinition extends In
1717
private Map<String, FieldStorage> storesStrings;
1818
private Map<String, FieldIndexing> indexesStrings;
1919
private Map<String, String> analyzersStrings;
20-
public Map<String, VectorFieldOptions> vectorFieldStrings = new HashMap<>();
20+
private Map<String, VectorFieldOptions> vectorFieldStrings = new HashMap<>();
2121
private Set<String> suggestionsOptions;
2222
private Map<String, FieldTermVector> termVectorsStrings;
2323
private Map<String, SpatialOptions> spatialIndexesStrings;
@@ -114,6 +114,10 @@ public void setReduce(String reduce) {
114114
this.reduce = reduce;
115115
}
116116

117+
public void setVectorOptionsStrings(Map<String, VectorFieldOptions> vectorOptionsStrings) {
118+
this.vectorFieldStrings = vectorOptionsStrings;
119+
}
120+
117121
public Map<String, FieldStorage> getStoresStrings() {
118122
return storesStrings;
119123
}

‎src/main/java/net/ravendb/client/documents/indexes/AbstractMultiMapIndexCreationTask.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ public IndexDefinition createIndexDefinition() {
4646
indexDefinitionBuilder.setPriority(getPriority());
4747
indexDefinitionBuilder.setState(getState());
4848
indexDefinitionBuilder.setDeploymentMode(getDeploymentMode());
49+
indexDefinitionBuilder.setVectorFieldStrings(vectorOptionsStrings);
4950

5051
if (searchEngineType != null) {
5152
indexDefinitionBuilder.getConfiguration().put(Constants.Configuration.Indexes.INDEXING_STATIC_SEARCH_ENGINE_TYPE, SharpEnum.value(searchEngineType));

‎src/main/java/net/ravendb/client/documents/queries/vectorSearch/IVectorOptions.java

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package net.ravendb.client.documents.queries.vectorSearch;
2+
3+
public class IVectorOptions {
4+
private Integer numberOfCandidates;
5+
private Double similarity;
6+
private Boolean isExact;
7+
8+
// Getters and setters
9+
public Integer getNumberOfCandidates() { return numberOfCandidates; }
10+
public void setNumberOfCandidates(Integer numberOfCandidates) { this.numberOfCandidates = numberOfCandidates; }
11+
12+
public Double getSimilarity() { return similarity; }
13+
public void setSimilarity(Double similarity) { this.similarity = similarity; }
14+
15+
public Boolean getIsExact() { return isExact; }
16+
public void setIsExact(Boolean isExact) { this.isExact = isExact; }
17+
}

‎src/main/java/net/ravendb/client/documents/queries/vectorSearch/VectorEmbeddingType.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
package net.ravendb.client.documents.queries.vectorSearch;
22

3+
import net.ravendb.client.primitives.UseSharpEnum;
4+
35
/**
46
* Represents the type of vector embedding.
57
*/
8+
@UseSharpEnum
69
public enum VectorEmbeddingType {
710
/**
811
* Single precision floating point (32-bit) vector

‎src/main/java/net/ravendb/client/documents/queries/vectorSearch/fields/IVectorOptionsJson.java

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package net.ravendb.client.documents.queries.vectorSearch.fields;
2+
3+
import net.ravendb.client.documents.queries.vectorSearch.VectorEmbeddingType;
4+
5+
public class IVectorOptionsJson {
6+
private VectorEmbeddingType sourceEmbeddingType;
7+
private VectorEmbeddingType destinationEmbeddingType;
8+
private Integer dimensions;
9+
private Integer numberOfEdges;
10+
private Integer numberOfCandidates;
11+
private Double similarity;
12+
private Boolean isExact;
13+
private String sourceFieldName;
14+
15+
// Getters and setters
16+
17+
public VectorEmbeddingType getSourceEmbeddingType() {
18+
return sourceEmbeddingType;
19+
}
20+
21+
public void setSourceEmbeddingType(VectorEmbeddingType sourceEmbeddingType) {
22+
this.sourceEmbeddingType = sourceEmbeddingType;
23+
}
24+
25+
public VectorEmbeddingType getDestinationEmbeddingType() {
26+
return destinationEmbeddingType;
27+
}
28+
29+
public void setDestinationEmbeddingType(VectorEmbeddingType destinationEmbeddingType) {
30+
this.destinationEmbeddingType = destinationEmbeddingType;
31+
}
32+
33+
public Integer getDimensions() {
34+
return dimensions;
35+
}
36+
37+
public void setDimensions(Integer dimensions) {
38+
this.dimensions = dimensions;
39+
}
40+
41+
public Integer getNumberOfEdges() {
42+
return numberOfEdges;
43+
}
44+
45+
public void setNumberOfEdges(Integer numberOfEdges) {
46+
this.numberOfEdges = numberOfEdges;
47+
}
48+
49+
public Integer getNumberOfCandidates() {
50+
return numberOfCandidates;
51+
}
52+
53+
public void setNumberOfCandidates(Integer numberOfCandidates) {
54+
this.numberOfCandidates = numberOfCandidates;
55+
}
56+
57+
public Double getSimilarity() {
58+
return similarity;
59+
}
60+
61+
public void setSimilarity(Double similarity) {
62+
this.similarity = similarity;
63+
}
64+
65+
public Boolean getIsExact() {
66+
return isExact;
67+
}
68+
69+
public void setIsExact(Boolean isExact) {
70+
this.isExact = isExact;
71+
}
72+
73+
public String getSourceFieldName() {
74+
return sourceFieldName;
75+
}
76+
77+
public void setSourceFieldName(String sourceFieldName) {
78+
this.sourceFieldName = sourceFieldName;
79+
}
80+
}

‎src/main/java/net/ravendb/client/documents/session/AbstractDocumentQuery.java

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
import net.ravendb.client.documents.queries.timeSeries.ITimeSeriesQueryBuilder;
2828
import net.ravendb.client.documents.queries.timeSeries.TimeSeriesQueryBuilder;
2929
import net.ravendb.client.documents.queries.timings.QueryTimings;
30+
import net.ravendb.client.documents.queries.vectorSearch.IVectorOptions;
31+
import net.ravendb.client.documents.queries.vectorSearch.VectorEmbeddingFieldFactory;
3032
import net.ravendb.client.documents.session.loaders.IncludeBuilderBase;
3133
import net.ravendb.client.documents.session.operations.QueryOperation;
3234
import net.ravendb.client.documents.session.operations.lazy.LazyQueryOperation;
@@ -46,6 +48,7 @@
4648
import java.time.Duration;
4749
import java.util.*;
4850
import java.util.function.Consumer;
51+
import java.util.function.Function;
4952

5053
/**
5154
* A query against a Raven index
@@ -969,6 +972,65 @@ protected CleanCloseable setFilterMode(boolean on) {
969972
return new FilterModeScope(filterModeStack, on);
970973
}
971974

975+
protected void _vectorSearch(Object fieldName, Object valueOrFactory, IVectorOptions options){
976+
this.assertMethodIsCurrentlySupported("vectorSearch");
977+
VectorEmbeddingFieldFactory vectorFactory = new VectorEmbeddingFieldFactory();
978+
IVectorField fieldAccessor;
979+
980+
if (fieldName instanceof String) {
981+
fieldAccessor = vectorFactory.withField((String) fieldName);
982+
} else if (fieldName instanceof Function) {
983+
Function<VectorEmbeddingFieldFactory, IVectorEmbeddingFieldFactoryAccessor<T>> func =
984+
(Function<VectorEmbeddingFieldFactory, IVectorEmbeddingFieldFactoryAccessor<T>>) fieldName;
985+
fieldAccessor = func.apply(vectorFactory);
986+
} else {
987+
throw new IllegalArgumentException(
988+
"fieldName must be either a string or a function that selects a vector field"
989+
);
990+
}
991+
992+
WhereParams whereParams = new WhereParams();
993+
whereParams.setFieldName(fieldAccessor.getFieldName());
994+
995+
if (valueOrFactory instanceof Consumer) {
996+
// Function case
997+
VectorEmbeddingFieldValueFactory fieldValueFactory = new VectorEmbeddingFieldValueFactory();
998+
@SuppressWarnings("unchecked")
999+
Consumer<VectorEmbeddingFieldValueFactory> func =
1000+
(Consumer<VectorEmbeddingFieldValueFactory>) valueOrFactory;
1001+
func.accept(fieldValueFactory);
1002+
1003+
if (fieldValueFactory.getEmbeddings() != null) {
1004+
whereParams.setValue(fieldValueFactory.getEmbeddings());
1005+
} else if (fieldValueFactory.getEmbedding() != null) {
1006+
whereParams.setValue(fieldValueFactory.getEmbedding());
1007+
} else if (fieldValueFactory.getText() != null) {
1008+
whereParams.setValue(fieldValueFactory.getText());
1009+
} else if (fieldValueFactory.getTexts() != null) {
1010+
whereParams.setValue(fieldValueFactory.getTexts());
1011+
} else {
1012+
throw new IllegalStateException("No value was provided in the valueFactory");
1013+
}
1014+
} else {
1015+
whereParams.setValue(valueOrFactory);
1016+
}
1017+
1018+
whereParams.setAllowWildcards(true);
1019+
Object transformToEqualValue = transformValue(whereParams);
1020+
List<QueryToken> tokens = getCurrentWhereTokens();
1021+
appendOperatorIfNeeded(tokens);
1022+
negateIfNeeded(tokens, whereParams.getFieldName());
1023+
1024+
WhereToken whereToken = WhereToken.create(
1025+
WhereOperator.VECTOR_SEARCH,
1026+
whereParams.getFieldName(),
1027+
this.addQueryParameter(transformToEqualValue),
1028+
new WhereToken.WhereOptions(options)
1029+
);
1030+
1031+
tokens.add(whereToken);
1032+
}
1033+
9721034
private static class FilterModeScope implements CleanCloseable {
9731035
private final Stack<Boolean> _modeStack;
9741036

‎src/main/java/net/ravendb/client/documents/session/DocumentQuery.java

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import net.ravendb.client.documents.queries.suggestions.*;
2222
import net.ravendb.client.documents.queries.timeSeries.ITimeSeriesQueryBuilder;
2323
import net.ravendb.client.documents.queries.timings.QueryTimings;
24+
import net.ravendb.client.documents.queries.vectorSearch.IVectorOptions;
2425
import net.ravendb.client.documents.session.loaders.IQueryIncludeBuilder;
2526
import net.ravendb.client.documents.session.loaders.QueryIncludeBuilder;
2627
import net.ravendb.client.documents.session.querying.sharding.IQueryShardedContextBuilder;
@@ -947,4 +948,24 @@ public IDocumentQuery<T> shardContext(Consumer<IQueryShardedContextBuilder> buil
947948
_shardContext(builder);
948949
return this;
949950
}
951+
952+
@Override
953+
public IDocumentQuery<T> vectorSearch(
954+
Object fieldName,
955+
Object valueFactory,
956+
IVectorOptions options
957+
) {
958+
this._vectorSearch(fieldName, valueFactory, options);
959+
return this;
960+
}
961+
962+
@Override
963+
public IDocumentQuery<T> vectorSearch(
964+
Object fieldName,
965+
Runnable valueFactory,
966+
IVectorOptions options
967+
) {
968+
this._vectorSearch(fieldName, valueFactory, options);
969+
return this;
970+
}
950971
}

‎src/main/java/net/ravendb/client/documents/session/IDocumentQuery.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import net.ravendb.client.documents.queries.suggestions.ISuggestionDocumentQuery;
1010
import net.ravendb.client.documents.queries.suggestions.SuggestionBase;
1111
import net.ravendb.client.documents.queries.timeSeries.ITimeSeriesQueryBuilder;
12+
import net.ravendb.client.documents.queries.vectorSearch.IVectorOptions;
1213
import net.ravendb.client.documents.session.querying.sharding.IQueryShardedContextBuilder;
1314

1415
import java.util.function.Consumer;
@@ -139,4 +140,16 @@ public interface IDocumentQuery<T> extends IDocumentQueryBase<T, IDocumentQuery<
139140
ISuggestionDocumentQuery<T> suggestUsing(Consumer<ISuggestionBuilder<T>> builder);
140141

141142
IDocumentQuery<T> shardContext(Consumer<IQueryShardedContextBuilder> builder);
143+
144+
IDocumentQuery<T> vectorSearch(
145+
Object fieldName,
146+
Object value,
147+
IVectorOptions options
148+
);
149+
150+
IDocumentQuery<T> vectorSearch(
151+
Object fieldName,
152+
Runnable factory,
153+
IVectorOptions options
154+
);
142155
}

0 commit comments

Comments
 (0)