SOLR-17780: Add support for scalar quantized dense vectors #3385
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,107 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one or more | ||
| * contributor license agreements. See the NOTICE file distributed with | ||
| * this work for additional information regarding copyright ownership. | ||
| * The ASF licenses this file to You under the Apache License, Version 2.0 | ||
| * (the "License"); you may not use this file except in compliance with | ||
| * the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package org.apache.solr.schema; | ||
|
|
||
| import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsFormat; | ||
| import org.apache.lucene.index.VectorEncoding; | ||
| import org.apache.lucene.index.VectorSimilarityFunction; | ||
|
|
||
| import java.util.Map; | ||
|
|
||
| import static java.util.Optional.ofNullable; | ||
|
|
||
| public class ScalarQuantizedDenseVectorField extends DenseVectorField { | ||
| public static final String BITS = "bits"; // | ||
| public static final String CONFIDENCE_INTERVAL = "confidenceInterval"; | ||
| public static final String DYNAMIC_CONFIDENCE_INTERVAL = "dynamicConfidenceInterval"; | ||
| public static final String COMPRESS = "compress"; // can only be enabled when bits = 4 per Lucene codec spec | ||
|
There was a problem hiding this comment. these first four lines are param names if I'm not mistaken. Maybe we can call them '_PARAM' e.g. 'BITS_PARAM' Or add a comment line at the beginning that clearly group them as param names, There was a problem hiding this comment. to be fair, checking the original DenseVectorField, the same renaming could help there as well I suspect |
||
|
|
||
| private static final int DEFAULT_BITS = 7; // use signed byte as default when unspecified | ||
| private static final Float DEFAULT_CONFIDENCE_INTERVAL = null; // use dimension scaled confidence interval | ||
|
|
||
| /** | ||
| * Number of bits to use for storage | ||
| * Must be 4 (half-byte) or 7 (signed-byte) per Lucene codec spec | ||
| */ | ||
| private int bits; | ||
|
|
||
| /** | ||
| * Confidence interval to use for scalar quantization | ||
| * Default is calculated as `1-1/(vector_dimensions + 1)` | ||
| */ | ||
| private Float confidenceInterval; | ||
|
|
||
| /** | ||
| * When enabled, in conjunction with 4 bit size, will pair values into single bytes for 50% reduction in memory usage | ||
| * (comes at the cost of some decode speed penalty) | ||
| */ | ||
| private boolean compress; | ||
|
|
||
| public ScalarQuantizedDenseVectorField(int dimension, | ||
| VectorSimilarityFunction similarityFunction, | ||
| VectorEncoding vectorEncoding, | ||
| int bits, | ||
| Float confidenceInterval, | ||
| boolean compress) { | ||
| super(dimension, similarityFunction, vectorEncoding); | ||
| this.bits = bits; | ||
| this.confidenceInterval = confidenceInterval; | ||
| this.compress = compress; | ||
| } | ||
|
|
||
| @Override | ||
| public void init(IndexSchema schema, Map<String, String> args) { | ||
| super.init(schema, args); | ||
|
|
||
| this.bits = ofNullable(args.get(BITS)) | ||
| .map(Integer::parseInt) | ||
| .orElse(DEFAULT_BITS); | ||
| args.remove(BITS); | ||
liangkaiwen marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| this.compress = ofNullable(args.get(COMPRESS)) | ||
| .map(Boolean::parseBoolean) | ||
| .orElse(false); | ||
| args.remove(COMPRESS); | ||
|
|
||
| boolean useDynamicConfidenceInterval = ofNullable(args.get(DYNAMIC_CONFIDENCE_INTERVAL)) | ||
| .map(Boolean::parseBoolean) | ||
| .orElse(false); | ||
| args.remove(DYNAMIC_CONFIDENCE_INTERVAL); | ||
|
|
||
| if (useDynamicConfidenceInterval) { | ||
| this.confidenceInterval = Lucene99ScalarQuantizedVectorsFormat.DYNAMIC_CONFIDENCE_INTERVAL; | ||
| } | ||
|
|
||
| this.confidenceInterval = ofNullable(args.get(CONFIDENCE_INTERVAL)) | ||
|
There was a problem hiding this comment. this overwrites the confidenceInterval set when useDynamicConfidenceInterval==true above There was a problem hiding this comment. good catch. I was hoping to make any explicit confidence interval override the dynamic behavior, however with both absence of arg and the default being null I think I will flip the behavior (use dynamic if specified otherwise confidence interval if specified otherwise default confidence interval if none provided) There was a problem hiding this comment. Okay but then let's check if both are set and throw an exception. There was a problem hiding this comment. If neither are set, it should default to the default behavior (scaled confidence interval) - which according to lucene spec is passing a null confidence interval |
||
| .map(Float::parseFloat) | ||
| .orElse(DEFAULT_CONFIDENCE_INTERVAL); | ||
| args.remove(CONFIDENCE_INTERVAL); | ||
| } | ||
|
|
||
| public int getBits() { | ||
| return bits; | ||
| } | ||
|
|
||
| public boolean useCompression() { | ||
| return compress; | ||
| } | ||
|
|
||
| public Float getConfidenceInterval() { | ||
| return confidenceInterval; | ||
| } | ||
|
|
||
| } | ||
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.
The instanceof checks can be a code smell. I look at all this and wonder: wouldn't it make sense for a method DenseVectorField.buildKnnVectorsFormat() to exist? Or just a getter if it's built in the field type's init()... which would also mean no need to have each of those fields (e.g. no bits, ...) with their getters either.
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.
reworked it. Let me know if this is along the lines of what you were thinking
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.
Much better. Did you consider init() also creating the codec to thus avoid the need corresponding fields, getters, setters?
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'm not quite sure I follow. You mean the ScalarQuantizedDenseVectorField init() or the SchemaCodecFactory init()?
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.
ScalarQuantizedDenseVectorField.init()(and its subclass, DenseVectorField) creating aKnnVectorsFormatthat later can simply be returned.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.
That's probably true. However it does look like the getters are used in the DenseVectorField tests as KnnVectorsFormat doesn't expose the internal setting values. Personally I think it's ok to leave patterned as-is and having getters affords some level of flexibility and transparency.
But I can make this change if you feel strongly about it