Skip to content

Add a flat vector format for bfloat16 vector storage #132533

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.

Sign up for GitHub

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

Open
wants to merge 14 commits into
base: lucene_snapshot
Choose a base branch
from
Open

Add a flat vector format for bfloat16 vector storage #132533

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
f105dc8
Add ES92 bfloat16 vector format
thecoop Aug 7, 2025
b008aea
Add to KnnIndexTester
thecoop Aug 7, 2025
370feb2
Add yaml tests
thecoop Aug 7, 2025
522c8ae
Add test feature for bfloat16
thecoop Aug 7, 2025
0c27c55
Mute org.elasticsearch.xpack.esql.qa.mixed.EsqlClientYamlIT test {p0=…
elasticsearchmachine Aug 8, 2025
f979e89
Define a new element type rather than a special option
thecoop Aug 7, 2025
22e0756
[CI] Auto commit changes from spotless
Aug 8, 2025
21623f5
Merge branch 'main' into bfloat16-vector-format
thecoop Aug 8, 2025
0a43817
Merge branch 'lucene_snapshot' into bfloat16-vector-format
thecoop Aug 8, 2025
6ff4792
Only run bfloat16 upgrade tests if bfloat16 is actually supported
thecoop Aug 8, 2025
d35e0bc
[CI] Auto commit changes from spotless
Aug 8, 2025
b142470
Remove unneeded reader/writer
thecoop Aug 8, 2025
ec74d87
[CI] Auto commit changes from spotless
Aug 8, 2025
2f4c21f
Remove byte vector support
thecoop Aug 8, 2025
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
11 changes: 11 additions & 0 deletions qa/vector/src/main/java/org/elasticsearch/test/knn/CmdLineArgs.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ record CmdLineArgs(
float filterSelectivity,
long seed,
VectorSimilarityFunction vectorSpace,
int rawVectorSize,
int quantizeBits,
VectorEncoding vectorEncoding,
int dimensions,
Expand All @@ -75,6 +76,7 @@ record CmdLineArgs(
static final ParseField FORCE_MERGE_FIELD = new ParseField("force_merge");
static final ParseField VECTOR_SPACE_FIELD = new ParseField("vector_space");
static final ParseField QUANTIZE_BITS_FIELD = new ParseField("quantize_bits");
static final ParseField RAW_VECTOR_SIZE_FIELD = new ParseField("raw_vector_size");
static final ParseField VECTOR_ENCODING_FIELD = new ParseField("vector_encoding");
static final ParseField DIMENSIONS_FIELD = new ParseField("dimensions");
static final ParseField EARLY_TERMINATION_FIELD = new ParseField("early_termination");
Expand Down Expand Up @@ -108,6 +110,7 @@ static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
PARSER.declareBoolean(Builder::setReindex, REINDEX_FIELD);
PARSER.declareBoolean(Builder::setForceMerge, FORCE_MERGE_FIELD);
PARSER.declareString(Builder::setVectorSpace, VECTOR_SPACE_FIELD);
PARSER.declareInt(Builder::setRawVectorSize, RAW_VECTOR_SIZE_FIELD);
PARSER.declareInt(Builder::setQuantizeBits, QUANTIZE_BITS_FIELD);
PARSER.declareString(Builder::setVectorEncoding, VECTOR_ENCODING_FIELD);
PARSER.declareInt(Builder::setDimensions, DIMENSIONS_FIELD);
Expand Down Expand Up @@ -143,6 +146,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
builder.field(REINDEX_FIELD.getPreferredName(), reindex);
builder.field(FORCE_MERGE_FIELD.getPreferredName(), forceMerge);
builder.field(VECTOR_SPACE_FIELD.getPreferredName(), vectorSpace.name().toLowerCase(Locale.ROOT));
builder.field(RAW_VECTOR_SIZE_FIELD.getPreferredName(), rawVectorSize);
builder.field(QUANTIZE_BITS_FIELD.getPreferredName(), quantizeBits);
builder.field(VECTOR_ENCODING_FIELD.getPreferredName(), vectorEncoding.name().toLowerCase(Locale.ROOT));
builder.field(DIMENSIONS_FIELD.getPreferredName(), dimensions);
Expand Down Expand Up @@ -176,6 +180,7 @@ static class Builder {
private boolean reindex = false;
private boolean forceMerge = false;
private VectorSimilarityFunction vectorSpace = VectorSimilarityFunction.EUCLIDEAN;
private int rawVectorSize = 32;
private int quantizeBits = 8;
private VectorEncoding vectorEncoding = VectorEncoding.FLOAT32;
private int dimensions;
Expand Down Expand Up @@ -278,6 +283,11 @@ public Builder setVectorSpace(String vectorSpace) {
return this;
}

public Builder setRawVectorSize(int rawVectorSize) {
this.rawVectorSize = rawVectorSize;
return this;
}

public Builder setQuantizeBits(int quantizeBits) {
this.quantizeBits = quantizeBits;
return this;
Expand Down Expand Up @@ -343,6 +353,7 @@ public CmdLineArgs build() {
filterSelectivity,
seed,
vectorSpace,
rawVectorSize,
quantizeBits,
vectorEncoding,
dimensions,
Expand Down
14 changes: 12 additions & 2 deletions qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
import org.elasticsearch.index.codec.vectors.IVFVectorsFormat;
import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat;
import org.elasticsearch.index.codec.vectors.es92.ES92BinaryQuantizedBFloat16VectorsFormat;
import org.elasticsearch.index.codec.vectors.es92.ES92HnswBinaryQuantizedBFloat16VectorsFormat;
import org.elasticsearch.logging.Level;
import org.elasticsearch.logging.LogManager;
import org.elasticsearch.logging.Logger;
Expand Down Expand Up @@ -105,9 +107,17 @@ static Codec createCodec(CmdLineArgs args) {
} else {
if (args.quantizeBits() == 1) {
if (args.indexType() == IndexType.FLAT) {
format = new ES818BinaryQuantizedVectorsFormat();
if (args.rawVectorSize() == 16) {
format = new ES92BinaryQuantizedBFloat16VectorsFormat();
} else {
format = new ES818BinaryQuantizedVectorsFormat();
}
} else {
format = new ES818HnswBinaryQuantizedVectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, null);
if (args.rawVectorSize() == 16) {
format = new ES92HnswBinaryQuantizedBFloat16VectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, null);
} else {
format = new ES818HnswBinaryQuantizedVectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, null);
}
}
} else if (args.quantizeBits() < 32) {
if (args.indexType() == IndexType.FLAT) {
Expand Down
Loading