Add dims parameter to dense_vector mapping #43444
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -12,10 +12,11 @@ | |
| import org.apache.lucene.index.IndexableField; | ||
| import org.apache.lucene.search.DocValuesFieldExistsQuery; | ||
| import org.apache.lucene.search.Query; | ||
| import org.apache.lucene.util.BytesRef; | ||
| import org.elasticsearch.common.settings.Settings; | ||
| import org.elasticsearch.common.xcontent.XContentBuilder; | ||
| import org.elasticsearch.common.xcontent.XContentParser.Token; | ||
| import org.elasticsearch.common.xcontent.support.XContentMapValues; | ||
| import org.elasticsearch.index.fielddata.IndexFieldData; | ||
| import org.elasticsearch.index.mapper.ArrayValueMapperParser; | ||
| import org.elasticsearch.index.mapper.FieldMapper; | ||
|
|
@@ -56,12 +57,28 @@ public static class Defaults { | |
| } | ||
|
|
||
| public static class Builder extends FieldMapper.Builder<Builder, DenseVectorFieldMapper> { | ||
| private int dims = 0; | ||
|
|
||
| public Builder(String name) { | ||
| super(name, Defaults.FIELD_TYPE, Defaults.FIELD_TYPE); | ||
| builder = this; | ||
| } | ||
|
|
||
| public Builder dims(int dims) { | ||
| if ((dims > MAX_DIMS_COUNT) || (dims < 1)) { | ||
| throw new MapperParsingException("The number of dimensions for field [" + name + | ||
| "] should be in the range [1, " + MAX_DIMS_COUNT + "]"); | ||
| } | ||
| this.dims = dims; | ||
| return this; | ||
| } | ||
|
|
||
| @Override | ||
| protected void setupFieldType(BuilderContext context) { | ||
| super.setupFieldType(context); | ||
| fieldType().setDims(dims); | ||
| } | ||
|
|
||
| @Override | ||
| public DenseVectorFieldType fieldType() { | ||
| return (DenseVectorFieldType) super.fieldType(); | ||
|
|
@@ -80,11 +97,17 @@ public static class TypeParser implements Mapper.TypeParser { | |
| @Override | ||
| public Mapper.Builder<?,?> parse(String name, Map<String, Object> node, ParserContext parserContext) throws MapperParsingException { | ||
| DenseVectorFieldMapper.Builder builder = new DenseVectorFieldMapper.Builder(name); | ||
| Object dimsField = node.remove("dims"); | ||
jtibshirani marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| if (dimsField == null) { | ||
|
There was a problem hiding this comment. Using the same reasoning as above, it would be nice to move this check to There was a problem hiding this comment. @jtibshirani Thanks, but we need this check here to ensure we don't get Also, I was thinking that There was a problem hiding this comment. One approach would be to check if the node is null, and only parse it and call |
||
| throw new MapperParsingException("The [dims] property must be specified for field [" + name + "]."); | ||
| } | ||
| int dims = XContentMapValues.nodeIntegerValue(dimsField); | ||
| return builder.dims(dims); | ||
| } | ||
| } | ||
|
|
||
| public static final class DenseVectorFieldType extends MappedFieldType { | ||
| private int dims; | ||
|
|
||
| public DenseVectorFieldType() {} | ||
|
|
||
|
|
@@ -96,6 +119,14 @@ public DenseVectorFieldType clone() { | |
| return new DenseVectorFieldType(this); | ||
| } | ||
|
|
||
| int dims() { | ||
| return dims; | ||
| } | ||
|
|
||
| void setDims(int dims) { | ||
| this.dims = dims; | ||
| } | ||
|
|
||
| @Override | ||
| public String typeName() { | ||
| return CONTENT_TYPE; | ||
|
|
@@ -143,39 +174,47 @@ public DenseVectorFieldType fieldType() { | |
| @Override | ||
| public void parse(ParseContext context) throws IOException { | ||
| if (context.externalValueSet()) { | ||
| throw new MapperParsingException("Field [" + name() + "] of type [" + typeName() + "] can't be used in multi-fields"); | ||
| } | ||
| int dims = fieldType().dims(); //number of vector dimensions | ||
|
|
||
| // encode array of floats as array of integers and store into buf | ||
| // this code is here and not int the VectorEncoderDecoder so not to create extra arrays | ||
| byte[] buf = new byte[dims * INT_BYTES]; | ||
| int offset = 0; | ||
| int dim = 0; | ||
| for (Token token = context.parser().nextToken(); token != Token.END_ARRAY; token = context.parser().nextToken()) { | ||
| if (dim++ >= dims) { | ||
| throw new MapperParsingException("Field [" + name() + "] of type [" + typeName() + "] of doc [" + | ||
|
There was a problem hiding this comment. It seems like There was a problem hiding this comment. @jtibshirani I agree with you. This should be |
||
| context.sourceToParse().id() + "] has exceeded the number of dimensions [" + dims + "] defined in mapping"); | ||
| } | ||
| ensureExpectedToken(Token.VALUE_NUMBER, token, context.parser()::getTokenLocation); | ||
| float value = context.parser().floatValue(true); | ||
| int intValue = Float.floatToIntBits(value); | ||
| buf[offset++] = (byte) (intValue >> 24); | ||
| buf[offset++] = (byte) (intValue >> 16); | ||
| buf[offset++] = (byte) (intValue >> 8); | ||
| buf[offset++] = (byte) intValue; | ||
| } | ||
| if (dim != dims) { | ||
| throw new MapperParsingException("Field [" + name() + "] of type [" + typeName() + "] of doc [" + | ||
| context.sourceToParse().id() + "] has number of dimensions [" + dim + | ||
| "] less than defined in the mapping [" + dims +"]"); | ||
| } | ||
| BinaryDocValuesField field = new BinaryDocValuesField(fieldType().name(), new BytesRef(buf, 0, offset)); | ||
| if (context.doc().getByKey(fieldType().name()) != null) { | ||
| throw new MapperParsingException("Field [" + name() + "] of type [" + typeName() + | ||
| "] doesn't not support indexing multiple values for the same field in the same document"); | ||
| } | ||
| context.doc().addWithKey(fieldType().name(), field); | ||
| } | ||
|
|
||
| @Override | ||
| protected void doXContentBody(XContentBuilder builder, boolean includeDefaults, Params params) throws IOException { | ||
| super.doXContentBody(builder, includeDefaults, params); | ||
| builder.field("dims", fieldType().dims()); | ||
| } | ||
|
|
||
| @Override | ||
| protected void parseCreateField(ParseContext context, List<IndexableField> fields) { | ||
| throw new AssertionError("parse is implemented directly"); | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this callout
<1>should go on the next line.