Skip to content

Commit e07d22e

Browse files
bkimmigmorganlutz
bkimmig
authored andcommitted
Adds support for dense_vector type
A dense_vector field stores dense vectors of float values. The maximum number of dimensions that can be in a vector should not exceed 2048. A dense_vector field is a single-valued field. Closes spring-projects#1700
1 parent 1c549b7 commit e07d22e

File tree

7 files changed

+86
-3
lines changed

7 files changed

+86
-3
lines changed

Diff for: src/main/java/org/springframework/data/elasticsearch/annotations/Field.java

+9-2
Original file line numberDiff line numberDiff line change
@@ -157,14 +157,14 @@
157157

158158
/**
159159
* if true, the field will be stored in Elasticsearch even if it has a null value
160-
*
160+
*
161161
* @since 4.1
162162
*/
163163
boolean storeNullValue() default false;
164164

165165
/**
166166
* to be used in combination with {@link FieldType#Rank_Feature}
167-
*
167+
*
168168
* @since 4.1
169169
*/
170170
boolean positiveScoreImpact() default true;
@@ -185,4 +185,11 @@
185185
* @since 4.1
186186
*/
187187
NullValueType nullValueType() default NullValueType.String;
188+
189+
/**
190+
* to be used in combination with {@link FieldType#Dense_Vector}
191+
*
192+
* @since 4.2
193+
*/
194+
int dims() default -1;
188195
}

Diff for: src/main/java/org/springframework/data/elasticsearch/annotations/FieldType.java

+3-1
Original file line numberDiff line numberDiff line change
@@ -57,5 +57,7 @@ public enum FieldType {
5757
/** @since 4.1 */
5858
Rank_Features, //
5959
/** since 4.2 */
60-
Wildcard //
60+
Wildcard, //
61+
/** @since 4.2 */
62+
Dense_Vector //
6163
}

Diff for: src/main/java/org/springframework/data/elasticsearch/annotations/InnerField.java

+7
Original file line numberDiff line numberDiff line change
@@ -140,4 +140,11 @@
140140
* @since 4.1
141141
*/
142142
NullValueType nullValueType() default NullValueType.String;
143+
144+
/**
145+
* to be used in combination with {@link FieldType#Dense_Vector}
146+
*
147+
* @since 4.2
148+
*/
149+
int dims() default -1;
143150
}

Diff for: src/main/java/org/springframework/data/elasticsearch/core/index/MappingParameters.java

+10
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ public final class MappingParameters {
6565
static final String FIELD_PARAM_NULL_VALUE = "null_value";
6666
static final String FIELD_PARAM_POSITION_INCREMENT_GAP = "position_increment_gap";
6767
static final String FIELD_PARAM_POSITIVE_SCORE_IMPACT = "positive_score_impact";
68+
static final String FIELD_PARAM_DIMS = "dims";
6869
static final String FIELD_PARAM_SCALING_FACTOR = "scaling_factor";
6970
static final String FIELD_PARAM_SEARCH_ANALYZER = "search_analyzer";
7071
static final String FIELD_PARAM_STORE = "store";
@@ -94,6 +95,7 @@ public final class MappingParameters {
9495
private final NullValueType nullValueType;
9596
private final Integer positionIncrementGap;
9697
private final boolean positiveScoreImpact;
98+
private final Integer dims;
9799
private final String searchAnalyzer;
98100
private final double scalingFactor;
99101
private final Similarity similarity;
@@ -153,6 +155,8 @@ private MappingParameters(Field field) {
153155
|| (maxShingleSize >= 2 && maxShingleSize <= 4), //
154156
"maxShingleSize must be in inclusive range from 2 to 4 for field type search_as_you_type");
155157
positiveScoreImpact = field.positiveScoreImpact();
158+
dims = field.dims();
159+
Assert.isTrue(dims <= 2048, "The maximum number of dimensions that can be in a vector should not exceed 2048.");
156160
Assert.isTrue(field.enabled() || type == FieldType.Object, "enabled false is only allowed for field type object");
157161
enabled = field.enabled();
158162
eagerGlobalOrdinals = field.eagerGlobalOrdinals();
@@ -191,6 +195,8 @@ private MappingParameters(InnerField field) {
191195
|| (maxShingleSize >= 2 && maxShingleSize <= 4), //
192196
"maxShingleSize must be in inclusive range from 2 to 4 for field type search_as_you_type");
193197
positiveScoreImpact = field.positiveScoreImpact();
198+
dims = field.dims();
199+
Assert.isTrue(dims <= 2048, "The maximum number of dimensions that can be in a vector should not exceed 2048.");
194200
enabled = true;
195201
eagerGlobalOrdinals = field.eagerGlobalOrdinals();
196202
}
@@ -323,6 +329,10 @@ public void writeTypeAndParametersTo(XContentBuilder builder) throws IOException
323329
builder.field(FIELD_PARAM_POSITIVE_SCORE_IMPACT, positiveScoreImpact);
324330
}
325331

332+
if (dims >= 1) {
333+
builder.field(FIELD_PARAM_DIMS, dims);
334+
}
335+
326336
if (!enabled) {
327337
builder.field(FIELD_PARAM_ENABLED, enabled);
328338
}

Diff for: src/test/java/org/springframework/data/elasticsearch/core/index/MappingBuilderIntegrationTests.java

+17
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,16 @@ void shouldWriteWildcardFieldMapping() {
271271
indexOps.putMapping();
272272
}
273273

274+
@Test // #1700
275+
@DisplayName("should write dense_vector field mapping")
276+
void shouldWriteDenseVectorFieldMapping() {
277+
278+
IndexOperations indexOps = operations.indexOps(DenseVectorEntity.class);
279+
indexOps.create();
280+
indexOps.putMapping();
281+
indexOps.delete();
282+
}
283+
274284
@Test // #1370
275285
@DisplayName("should write mapping for disabled entity")
276286
void shouldWriteMappingForDisabledEntity() {
@@ -657,4 +667,11 @@ static class DisabledMappingProperty {
657667
@Field(type = Text) private String text;
658668
@Mapping(enabled = false) @Field(type = Object) private Object object;
659669
}
670+
671+
@Data
672+
@Document(indexName = "densevector-test")
673+
static class DenseVectorEntity {
674+
@Id private String id;
675+
@Field(type = Dense_Vector, dims = 3) private float[] dense_vector;
676+
}
660677
}

Diff for: src/test/java/org/springframework/data/elasticsearch/core/index/MappingBuilderUnitTests.java

+25
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242

4343
import org.elasticsearch.search.suggest.completion.context.ContextMapping;
4444
import org.json.JSONException;
45+
import org.junit.jupiter.api.Assertions;
4546
import org.junit.jupiter.api.DisplayName;
4647
import org.junit.jupiter.api.Test;
4748
import org.springframework.data.annotation.Id;
@@ -506,6 +507,23 @@ void shouldWriteRankFeatureProperties() throws JSONException {
506507
assertEquals(expected, mapping, false);
507508
}
508509

510+
@Test // #1700
511+
@DisplayName("should write dense_vector properties")
512+
void shouldWriteDenseVectorProperties() throws JSONException {
513+
String expected = "{\n" + //
514+
" \"properties\": {\n" + //
515+
" \"my_vector\": {\n" + //
516+
" \"type\": \"dense_vector\",\n" + //
517+
" \"dims\": 16\n" + //
518+
" }\n" + //
519+
" }\n" + //
520+
"}\n"; //
521+
522+
String mapping = getMappingBuilder().buildPropertyMapping(DenseVectorEntity.class);
523+
524+
assertEquals(expected, mapping, false);
525+
}
526+
509527
@Test // #1370
510528
@DisplayName("should not write mapping when enabled is false on entity")
511529
void shouldNotWriteMappingWhenEnabledIsFalseOnEntity() throws JSONException {
@@ -963,6 +981,13 @@ static class RankFeatureEntity {
963981
@Field(type = FieldType.Rank_Features) private Map<String, Integer> topics;
964982
}
965983

984+
@Data
985+
static class DenseVectorEntity {
986+
987+
@Id private String id;
988+
@Field(type = FieldType.Dense_Vector, dims = 16) private float[] my_vector;
989+
}
990+
966991
@Data
967992
@Mapping(enabled = false)
968993
static class DisabledMappingEntity {

Diff for: src/test/java/org/springframework/data/elasticsearch/core/index/MappingParametersTest.java

+15
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.springframework.data.elasticsearch.core.index;
22

33
import static org.assertj.core.api.Assertions.*;
4+
import static org.springframework.data.elasticsearch.annotations.FieldType.Dense_Vector;
45
import static org.springframework.data.elasticsearch.annotations.FieldType.Object;
56

67
import java.lang.annotation.Annotation;
@@ -66,6 +67,16 @@ void shouldAllowEnabledFalseOnlyOnObjectFields() {
6667
assertThatThrownBy(() -> MappingParameters.from(annotation)).isInstanceOf(IllegalArgumentException.class);
6768
}
6869

70+
@Test // #1700
71+
@DisplayName("should not allow dims length greater than 2048")
72+
void shouldNotAllowDimsLengthGreaterThan2048() {
73+
ElasticsearchPersistentEntity<?> failEntity = elasticsearchConverter.get().getMappingContext()
74+
.getRequiredPersistentEntity(InvalidDenseVectorClass.class);
75+
Annotation annotation = failEntity.getRequiredPersistentProperty("dense_vector").findAnnotation(Field.class);
76+
77+
assertThatThrownBy(() -> MappingParameters.from(annotation)).isInstanceOf(IllegalArgumentException.class);
78+
}
79+
6980
static class AnnotatedClass {
7081
@Nullable @Field private String field;
7182
@Nullable @MultiField(mainField = @Field,
@@ -79,4 +90,8 @@ static class AnnotatedClass {
7990
static class InvalidEnabledFieldClass {
8091
@Nullable @Field(type = FieldType.Text, enabled = false) private String disabledObject;
8192
}
93+
94+
static class InvalidDenseVectorClass {
95+
@Field(type = Dense_Vector, dims = 2049) private float[] dense_vector;
96+
}
8297
}

0 commit comments

Comments
 (0)