Skip to content

Adds support for dense_vector type #1708

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

Merged
merged 1 commit into from
Feb 26, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@
* @author Peter-Josef Meisch
* @author Xiao Yu
* @author Aleksei Arsenev
* @author Brian Kimmig
* @author Morgan Lutz
*/
@Retention(RetentionPolicy.RUNTIME)
@Target({ ElementType.FIELD, ElementType.ANNOTATION_TYPE })
Expand Down Expand Up @@ -185,4 +187,11 @@
* @since 4.1
*/
NullValueType nullValueType() default NullValueType.String;

/**
* to be used in combination with {@link FieldType#Dense_Vector}
*
* @since 4.2
*/
int dims() default -1;
}
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
* @author Zeng Zetang
* @author Peter-Josef Meisch
* @author Aleksei Arsenev
* @author Brian Kimmig
* @author Morgan Lutz
*/
public enum FieldType {
Auto, //
Expand Down Expand Up @@ -57,5 +59,7 @@ public enum FieldType {
/** @since 4.1 */
Rank_Features, //
/** since 4.2 */
Wildcard //
Wildcard, //
/** @since 4.2 */
Dense_Vector //
}
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
* @author Xiao Yu
* @author Peter-Josef Meisch
* @author Aleksei Arsenev
* @author Brian Kimmig
* @author Morgan Lutz
*/
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.ANNOTATION_TYPE)
Expand Down Expand Up @@ -140,4 +142,11 @@
* @since 4.1
*/
NullValueType nullValueType() default NullValueType.String;

/**
* to be used in combination with {@link FieldType#Dense_Vector}
*
* @since 4.2
*/
int dims() default -1;
}
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@
*
* @author Peter-Josef Meisch
* @author Aleksei Arsenev
* @author Brian Kimmig
* @author Morgan Lutz
* @since 4.0
*/
public final class MappingParameters {
Expand All @@ -65,6 +67,7 @@ public final class MappingParameters {
static final String FIELD_PARAM_NULL_VALUE = "null_value";
static final String FIELD_PARAM_POSITION_INCREMENT_GAP = "position_increment_gap";
static final String FIELD_PARAM_POSITIVE_SCORE_IMPACT = "positive_score_impact";
static final String FIELD_PARAM_DIMS = "dims";
static final String FIELD_PARAM_SCALING_FACTOR = "scaling_factor";
static final String FIELD_PARAM_SEARCH_ANALYZER = "search_analyzer";
static final String FIELD_PARAM_STORE = "store";
Expand Down Expand Up @@ -94,6 +97,7 @@ public final class MappingParameters {
private final NullValueType nullValueType;
private final Integer positionIncrementGap;
private final boolean positiveScoreImpact;
private final Integer dims;
private final String searchAnalyzer;
private final double scalingFactor;
private final Similarity similarity;
Expand Down Expand Up @@ -153,6 +157,10 @@ private MappingParameters(Field field) {
|| (maxShingleSize >= 2 && maxShingleSize <= 4), //
"maxShingleSize must be in inclusive range from 2 to 4 for field type search_as_you_type");
positiveScoreImpact = field.positiveScoreImpact();
dims = field.dims();
if (type == FieldType.Dense_Vector) {
Assert.isTrue(dims >= 1 && dims <= 2048, "Invalid required parameter! Dense_Vector value \"dims\" must be between 1 and 2048.");
}
Assert.isTrue(field.enabled() || type == FieldType.Object, "enabled false is only allowed for field type object");
enabled = field.enabled();
eagerGlobalOrdinals = field.eagerGlobalOrdinals();
Expand Down Expand Up @@ -191,6 +199,10 @@ private MappingParameters(InnerField field) {
|| (maxShingleSize >= 2 && maxShingleSize <= 4), //
"maxShingleSize must be in inclusive range from 2 to 4 for field type search_as_you_type");
positiveScoreImpact = field.positiveScoreImpact();
dims = field.dims();
if (type == FieldType.Dense_Vector) {
Assert.isTrue(dims >= 1 && dims <= 2048, "Invalid required parameter! Dense_Vector value \"dims\" must be between 1 and 2048.");
}
enabled = true;
eagerGlobalOrdinals = field.eagerGlobalOrdinals();
}
Expand Down Expand Up @@ -323,6 +335,10 @@ public void writeTypeAndParametersTo(XContentBuilder builder) throws IOException
builder.field(FIELD_PARAM_POSITIVE_SCORE_IMPACT, positiveScoreImpact);
}

if (type == FieldType.Dense_Vector) {
builder.field(FIELD_PARAM_DIMS, dims);
}

if (!enabled) {
builder.field(FIELD_PARAM_ENABLED, enabled);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,8 @@
* @author Peter-Josef Meisch
* @author Xiao Yu
* @author Roman Puchkovskiy
* @author Brian Kimmig
* @author Morgan Lutz
*/
@SpringIntegrationTest
@ContextConfiguration(classes = { ElasticsearchRestTemplateConfiguration.class })
Expand Down Expand Up @@ -271,6 +273,16 @@ void shouldWriteWildcardFieldMapping() {
indexOps.putMapping();
}

@Test // #1700
@DisplayName("should write dense_vector field mapping")
void shouldWriteDenseVectorFieldMapping() {

IndexOperations indexOps = operations.indexOps(DenseVectorEntity.class);
indexOps.create();
indexOps.putMapping();
indexOps.delete();
}

@Test // #1370
@DisplayName("should write mapping for disabled entity")
void shouldWriteMappingForDisabledEntity() {
Expand Down Expand Up @@ -657,4 +669,11 @@ static class DisabledMappingProperty {
@Field(type = Text) private String text;
@Mapping(enabled = false) @Field(type = Object) private Object object;
}

@Data
@Document(indexName = "densevector-test")
static class DenseVectorEntity {
@Id private String id;
@Field(type = Dense_Vector, dims = 3) private float[] dense_vector;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@
* @author Peter-Josef Meisch
* @author Xiao Yu
* @author Roman Puchkovskiy
* @author Brian Kimmig
* @author Morgan Lutz
*/
public class MappingBuilderUnitTests extends MappingContextBaseTests {

Expand Down Expand Up @@ -506,6 +508,23 @@ void shouldWriteRankFeatureProperties() throws JSONException {
assertEquals(expected, mapping, false);
}

@Test // #1700
@DisplayName("should write dense_vector properties")
void shouldWriteDenseVectorProperties() throws JSONException {
String expected = "{\n" + //
" \"properties\": {\n" + //
" \"my_vector\": {\n" + //
" \"type\": \"dense_vector\",\n" + //
" \"dims\": 16\n" + //
" }\n" + //
" }\n" + //
"}\n"; //

String mapping = getMappingBuilder().buildPropertyMapping(DenseVectorEntity.class);

assertEquals(expected, mapping, false);
}

@Test // #1370
@DisplayName("should not write mapping when enabled is false on entity")
void shouldNotWriteMappingWhenEnabledIsFalseOnEntity() throws JSONException {
Expand Down Expand Up @@ -963,6 +982,13 @@ static class RankFeatureEntity {
@Field(type = FieldType.Rank_Features) private Map<String, Integer> topics;
}

@Data
static class DenseVectorEntity {

@Id private String id;
@Field(type = FieldType.Dense_Vector, dims = 16) private float[] my_vector;
}

@Data
@Mapping(enabled = false)
static class DisabledMappingEntity {
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package org.springframework.data.elasticsearch.core.index;

import static org.assertj.core.api.Assertions.*;
import static org.springframework.data.elasticsearch.annotations.FieldType.Dense_Vector;
import static org.springframework.data.elasticsearch.annotations.FieldType.Object;

import java.lang.annotation.Annotation;
Expand All @@ -17,6 +18,8 @@

/**
* @author Peter-Josef Meisch
* @author Brian Kimmig
* @author Morgan Lutz
*/
public class MappingParametersTest extends MappingContextBaseTests {

Expand Down Expand Up @@ -66,6 +69,26 @@ void shouldAllowEnabledFalseOnlyOnObjectFields() {
assertThatThrownBy(() -> MappingParameters.from(annotation)).isInstanceOf(IllegalArgumentException.class);
}

@Test // #1700
@DisplayName("should not allow dims length greater than 2048 for dense_vector type")
void shouldNotAllowDimsLengthGreaterThan2048ForDenseVectorType() {
ElasticsearchPersistentEntity<?> failEntity = elasticsearchConverter.get().getMappingContext()
.getRequiredPersistentEntity(DenseVectorInvalidDimsClass.class);
Annotation annotation = failEntity.getRequiredPersistentProperty("dense_vector").findAnnotation(Field.class);

assertThatThrownBy(() -> MappingParameters.from(annotation)).isInstanceOf(IllegalArgumentException.class);
}

@Test // #1700
@DisplayName("should require dims parameter for dense_vector type")
void shouldRequireDimsParameterForDenseVectorType() {
ElasticsearchPersistentEntity<?> failEntity = elasticsearchConverter.get().getMappingContext()
.getRequiredPersistentEntity(DenseVectorMissingDimsClass.class);
Annotation annotation = failEntity.getRequiredPersistentProperty("dense_vector").findAnnotation(Field.class);

assertThatThrownBy(() -> MappingParameters.from(annotation)).isInstanceOf(IllegalArgumentException.class);
}

static class AnnotatedClass {
@Nullable @Field private String field;
@Nullable @MultiField(mainField = @Field,
Expand All @@ -79,4 +102,12 @@ static class AnnotatedClass {
static class InvalidEnabledFieldClass {
@Nullable @Field(type = FieldType.Text, enabled = false) private String disabledObject;
}

static class DenseVectorInvalidDimsClass {
@Field(type = Dense_Vector, dims = 2049) private float[] dense_vector;
}

static class DenseVectorMissingDimsClass {
@Field(type = Dense_Vector) private float[] dense_vector;
}
}