Add BFloat16 raw vector format to bbq_hnsw and bbq_disk

Author: thecoop

qa/vector/src/main/java/org/elasticsearch/test/knn/CmdLineArgs.java

@@ -51,6 +51,7 @@ record CmdLineArgs(
     float filterSelectivity,
     long seed,
     VectorSimilarityFunction vectorSpace,
+    int rawVectorSize,
     int quantizeBits,
     VectorEncoding vectorEncoding,
     int dimensions,
@@ -80,6 +81,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");
@@ -123,6 +125,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);
@@ -161,6 +164,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);
@@ -196,6 +200,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;
@@ -305,6 +310,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;
@@ -380,6 +390,7 @@ public CmdLineArgs build() {
                 filterSelectivity,
                 seed,
                 vectorSpace,
+                rawVectorSize,
                 quantizeBits,
                 vectorEncoding,
                 dimensions,

rest-api-spec/src/main/resources/rest-api-spec/api/cat.segments.json

@@ -93,29 +93,6 @@
           "micros",
           "nanos"
         ]
-      },
-      "ignore_unavailable": {
-        "type": "boolean",
-        "description": "Whether specified concrete indices should be ignored when unavailable (missing or closed). Only allowed when providing an index expression."
-      },
-      "ignore_throttled": {
-        "type": "boolean",
-        "description": "Whether specified concrete, expanded or aliased indices should be ignored when throttled. Only allowed when providing an index expression."
-      },
-      "allow_no_indices": {
-        "type": "boolean",
-        "description": "Whether to ignore if a wildcard indices expression resolves into no concrete indices. (This includes `_all` string or when no indices have been specified). Only allowed when providing an index expression."
-      },
-      "expand_wildcards": {
-        "type": "enum",
-        "options": ["open", "closed", "hidden", "none", "all"],
-        "default": "open",
-        "description": "Whether to expand wildcard expression to concrete indices that are open, closed or both."
-      },
-      "allow_closed": {
-        "type": "boolean",
-        "description": "If true, allow closed indices to be returned in the response otherwise if false, keep the legacy behaviour of throwing an exception if index pattern matches closed indices",
-        "default": false
       }
     }
   }

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/cat.segments/10_basic.yml

@@ -209,77 +209,3 @@ tsdb:
       $body: |
         /^(tsdb \s+ 0 \s+ p \s+ \d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3} \s+ _\d (\s\d){3} \s+
         (\d+|\d+[.]\d+)(kb|b) \s+ \d+ (\s+ (false|true)){2} \s+  \d+\.\d+(\.\d+)? \s+ (false|true) \s? \n?)$/
-
----
-Wildcard Expansion Settings:
-  - requires:
-      capabilities:
-        - method: GET
-          path: /_cat/segments
-          capabilities: [ allow_closed ]
-      test_runner_features: [ capabilities ]
-      reason: Capability required to run test
-
-  - do:
-      indices.create:
-        index: basic-index
-        body:
-          settings:
-            number_of_shards: 1
-            number_of_replicas: 0
-
-  - do:
-      index:
-        index: basic-index
-        id: "1"
-        body:
-          field: "basic doc 1"
-
-  - do:
-      indices.create:
-        index: hidden-index
-        body:
-          settings:
-            number_of_shards: 1
-            number_of_replicas: 0
-            index.hidden: true
-
-  - do:
-      index:
-        index: hidden-index
-        id: "1"
-        body:
-          field: "hidden doc 1"
-
-  - do:
-      indices.create:
-        index: closed-index
-        body:
-          settings:
-            number_of_shards: 1
-            number_of_replicas: 0
-
-  - do:
-      index:
-        index: closed-index
-        id: "1"
-        body:
-          field: "closed doc 1"
-
-  - do:
-      indices.refresh:
-        index: [ basic-index, closed-index, hidden-index ]
-
-  - do:
-      indices.close:
-        index: closed-index
-
-  - do:
-      cat.segments:
-        v: true
-        s: index
-        expand_wildcards: all
-        allow_closed: true
-  - match:
-      $body: |
-        /basic-index(\s)+0(\s)+p.*\nclosed-index(\s)+0(\s)+p.*\nhidden-index(\s)+0(\s)+p.*/

server/src/main/java/org/elasticsearch/cluster/coordination/CoordinationState.java

@@ -211,9 +211,9 @@ public Join handleStartJoin(StartJoinRequest startJoinRequest) {
     }
 
     /**
-     * May be called on receipt of a {@link Join}, which is effectively a vote for the receiving node to be the elected master.
+     * May be called on receipt of a Join.
      *
-     * @param join The {@link Join} received.
+     * @param join The Join received.
      * @return true iff this instance does not already have a join vote from the given source node for this term
      * @throws CoordinationStateRejectedException if the arguments were incompatible with the current state of this object.
      */
@@ -234,9 +234,6 @@ public boolean handleJoin(Join join) {
 
         final long lastAcceptedTerm = getLastAcceptedTerm();
         if (join.lastAcceptedTerm() > lastAcceptedTerm) {
-            // Note that this is running on the receiving node, so it must reject joins from nodes with fresher state. This is unlike a
-            // real-world election where candidates will accept every vote they receive and it's the voter's responsibility to be selective
-            // about the votes they cast.
             logger.debug(
                 "handleJoin: ignored join as joiner has a better last accepted term (expected: <=[{}], actual: [{}])",
                 lastAcceptedTerm,
@@ -251,9 +248,6 @@ public boolean handleJoin(Join join) {
         }
 
         if (join.lastAcceptedTerm() == lastAcceptedTerm && join.lastAcceptedVersion() > getLastAcceptedVersion()) {
-            // Note that this is running on the receiving node, so it must reject joins from nodes with fresher state. This is unlike a
-            // real-world election where candidates will accept every vote they receive and it's the voter's responsibility to be selective
-            // about the votes they cast.
             logger.debug(
                 "handleJoin: ignored join as joiner has a better last accepted version (expected: <=[{}], actual: [{}]) in term {}",
                 getLastAcceptedVersion(),

server/src/main/java/org/elasticsearch/index/codec/vectors/BFloat16.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.codec.vectors;
+
+import org.apache.lucene.util.BitUtil;
+
+import java.nio.ByteOrder;
+import java.nio.ShortBuffer;
+
+public class BFloat16 {
+
+    public static final int BYTES = Short.BYTES;
+
+    public static short floatToBFloat16(float f) {
+        // this rounds towards 0
+        // zero - zero exp, zero fraction
+        // denormal - zero exp, non-zero fraction
+        // infinity - all-1 exp, zero fraction
+        // NaN - all-1 exp, non-zero fraction
+        // the Float.NaN constant is 0x7fc0_0000, so this won't turn the most common NaN values into
+        // infinities
+        return (short) (Float.floatToIntBits(f) >>> 16);
+    }
+
+    public static float bFloat16ToFloat(short bf) {
+        return Float.intBitsToFloat(bf << 16);
+    }
+
+    public static void floatToBFloat16(float[] floats, ShortBuffer bFloats) {
+        assert bFloats.remaining() == floats.length;
+        assert bFloats.order() == ByteOrder.LITTLE_ENDIAN;
+        for (float v : floats) {
+            bFloats.put(floatToBFloat16(v));
+        }
+    }
+
+    public static void bFloat16ToFloat(byte[] bfBytes, float[] floats) {
+        assert floats.length * 2 == bfBytes.length;
+        for (int i = 0; i < floats.length; i++) {
+            floats[i] = bFloat16ToFloat((short) BitUtil.VH_LE_SHORT.get(bfBytes, i * 2));
+        }
+    }
+
+    public static void bFloat16ToFloat(ShortBuffer bFloats, float[] floats) {
+        assert floats.length == bFloats.remaining();
+        assert bFloats.order() == ByteOrder.LITTLE_ENDIAN;
+        for (int i = 0; i < floats.length; i++) {
+            floats[i] = bFloat16ToFloat(bFloats.get());
+        }
+    }
+
+    private BFloat16() {}
+}

server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java

@@ -11,18 +11,81 @@
 
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.store.FlushInfo;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.MergeInfo;
+import org.elasticsearch.common.util.set.Sets;
+import org.elasticsearch.index.codec.vectors.es818.DirectIOHint;
+import org.elasticsearch.index.store.FsDirectoryFactory;
 
 import java.io.IOException;
+import java.util.Set;
 
 public abstract class DirectIOCapableFlatVectorsFormat extends AbstractFlatVectorsFormat {
     protected DirectIOCapableFlatVectorsFormat(String name) {
         super(name);
     }
 
+    protected abstract FlatVectorsReader createReader(SegmentReadState state) throws IOException;
+
+    static boolean canUseDirectIO(SegmentReadState state) {
+        return FsDirectoryFactory.isHybridFs(state.directory);
+    }
+
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
         return fieldsReader(state, false);
     }
 
-    public abstract FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException;
+    public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException {
+        if (state.context.context() == IOContext.Context.DEFAULT && useDirectIO && canUseDirectIO(state)) {
+            // only override the context for the random-access use case
+            SegmentReadState directIOState = new SegmentReadState(
+                state.directory,
+                state.segmentInfo,
+                state.fieldInfos,
+                new DirectIOContext(state.context.hints()),
+                state.segmentSuffix
+            );
+            // Use mmap for merges and direct I/O for searches.
+            return new MergeReaderWrapper(createReader(directIOState), createReader(state));
+        } else {
+            return createReader(state);
+        }
+    }
+
+    static class DirectIOContext implements IOContext {
+
+        final Set<FileOpenHint> hints;
+
+        DirectIOContext(Set<FileOpenHint> hints) {
+            // always add DirectIOHint to the hints given
+            this.hints = Sets.union(hints, Set.of(DirectIOHint.INSTANCE));
+        }
+
+        @Override
+        public Context context() {
+            return Context.DEFAULT;
+        }
+
+        @Override
+        public MergeInfo mergeInfo() {
+            return null;
+        }
+
+        @Override
+        public FlushInfo flushInfo() {
+            return null;
+        }
+
+        @Override
+        public Set<FileOpenHint> hints() {
+            return hints;
+        }
+
+        @Override
+        public IOContext withHints(FileOpenHint... hints) {
+            return new DirectIOContext(Set.of(hints));
+        }
+    }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormat.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
 import org.elasticsearch.index.codec.vectors.es93.DirectIOCapableLucene99FlatVectorsFormat;
+import org.elasticsearch.index.codec.vectors.es93.ES93BFloat16FlatVectorsFormat;
 
 import java.io.IOException;
 import java.util.Map;
@@ -58,12 +59,17 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
     public static final int VERSION_DIRECT_IO = 1;
     public static final int VERSION_CURRENT = VERSION_DIRECT_IO;
 
-    private static final DirectIOCapableFlatVectorsFormat rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
+    private static final DirectIOCapableFlatVectorsFormat float32VectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
+        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
+    );
+    private static final DirectIOCapableFlatVectorsFormat bfloat16VectorFormat = new ES93BFloat16FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
     private static final Map<String, DirectIOCapableFlatVectorsFormat> supportedFormats = Map.of(
-        rawVectorFormat.getName(),
-        rawVectorFormat
+        float32VectorFormat.getName(),
+        float32VectorFormat,
+        bfloat16VectorFormat.getName(),
+        bfloat16VectorFormat
     );
 
     // This dynamically sets the cluster probe based on the `k` requested and the number of clusters.
@@ -79,12 +85,13 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
     private final int vectorPerCluster;
     private final int centroidsPerParentCluster;
     private final boolean useDirectIO;
+    private final DirectIOCapableFlatVectorsFormat rawVectorFormat;
 
     public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster) {
-        this(vectorPerCluster, centroidsPerParentCluster, false);
+        this(vectorPerCluster, centroidsPerParentCluster, false, false);
     }
 
-    public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster, boolean useDirectIO) {
+    public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster, boolean useDirectIO, boolean useBFloat16) {
         super(NAME);
         if (vectorPerCluster < MIN_VECTORS_PER_CLUSTER || vectorPerCluster > MAX_VECTORS_PER_CLUSTER) {
             throw new IllegalArgumentException(
@@ -109,6 +116,7 @@ public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentClu
         this.vectorPerCluster = vectorPerCluster;
         this.centroidsPerParentCluster = centroidsPerParentCluster;
         this.useDirectIO = useDirectIO;
+        this.rawVectorFormat = useBFloat16 ? bfloat16VectorFormat : float32VectorFormat;
     }
 
     /** Constructs a format using the given graph construction parameters and scalar quantization. */