This adds a new experimental IVF format behind a feature flag (elastic#128631)

* Adding new bbq_ivf format behind a feature flag

* adding tests

* [CI] Auto commit changes from spotless

* addressing pr comments

* fixing flagging for yaml tests

* adjust ivf search to utilize num candidates as approximation measure

---------

Co-authored-by: elasticsearchmachine <[email protected]>

Author: 2 people

rest-api-spec/src/yamlRestTest/java/org/elasticsearch/test/rest/ClientYamlTestSuiteIT.java

@@ -38,6 +38,7 @@ public class ClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
         .feature(FeatureFlag.SUB_OBJECTS_AUTO_ENABLED)
         .feature(FeatureFlag.DOC_VALUES_SKIPPER)
         .feature(FeatureFlag.USE_LUCENE101_POSTINGS_FORMAT)
+        .feature(FeatureFlag.IVF_FORMAT)
         .build();
 
     public ClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) {

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/46_knn_search_bbq_ivf.yml

@@ -61,14 +61,26 @@ public class IVFVectorsFormat extends KnnVectorsFormat {
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
 
-    private static final int DEFAULT_VECTORS_PER_CLUSTER = 1000;
+    // This dynamically sets the cluster probe based on the `k` requested and the number of clusters.
+    // useful when searching with 'efSearch' type parameters instead of requiring a specific nprobe.
+    public static final int DYNAMIC_NPROBE = -1;
+    public static final int DEFAULT_VECTORS_PER_CLUSTER = 384;
+    public static final int MIN_VECTORS_PER_CLUSTER = 64;
+    public static final int MAX_VECTORS_PER_CLUSTER = 1 << 16; // 65536
 
     private final int vectorPerCluster;
 
     public IVFVectorsFormat(int vectorPerCluster) {
         super(NAME);
-        if (vectorPerCluster <= 0) {
-            throw new IllegalArgumentException("vectorPerCluster must be > 0");
+        if (vectorPerCluster < MIN_VECTORS_PER_CLUSTER || vectorPerCluster > MAX_VECTORS_PER_CLUSTER) {
+            throw new IllegalArgumentException(
+                "vectorsPerCluster must be between "
+                    + MIN_VECTORS_PER_CLUSTER
+                    + " and "
+                    + MAX_VECTORS_PER_CLUSTER
+                    + ", got: "
+                    + vectorPerCluster
+            );
         }
         this.vectorPerCluster = vectorPerCluster;
     }
@@ -90,12 +102,12 @@ public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException
 
     @Override
     public int getMaxDimensions(String fieldName) {
-        return 1024;
+        return 4096;
     }
 
     @Override
     public String toString() {
-        return "IVFVectorFormat";
+        return "IVFVectorsFormat(" + "vectorPerCluster=" + vectorPerCluster + ')';
     }
 
     static IVFVectorsReader getIVFReader(KnnVectorsReader vectorsReader, String fieldName) {

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

@@ -38,6 +38,7 @@
 import java.util.function.IntPredicate;
 
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS;
+import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.DYNAMIC_NPROBE;
 
 /**
  * Reader for IVF vectors. This reader is used to read the IVF vectors from the index.
@@ -226,17 +227,6 @@ public final ByteVectorValues getByteVectorValues(String field) throws IOExcepti
         return rawVectorsReader.getByteVectorValues(field);
     }
 
-    protected float[] getGlobalCentroid(FieldInfo info) {
-        if (info == null || info.getVectorEncoding().equals(VectorEncoding.BYTE)) {
-            return null;
-        }
-        FieldEntry entry = fields.get(info.number);
-        if (entry == null) {
-            return null;
-        }
-        return entry.globalCentroid();
-    }
-
     @Override
     public final void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
         final FieldInfo fieldInfo = state.fieldInfos.fieldInfo(field);
@@ -261,12 +251,9 @@ public final void search(String field, float[] target, KnnCollector knnCollector
             }
             return visitedDocs.getAndSet(docId) == false;
         };
-        final int nProbe;
+        int nProbe = DYNAMIC_NPROBE;
         if (knnCollector.getSearchStrategy() instanceof IVFKnnSearchStrategy ivfSearchStrategy) {
             nProbe = ivfSearchStrategy.getNProbe();
-        } else {
-            // TODO calculate nProbe given the number of centroids vs. number of vectors for given `k`
-            nProbe = 10;
         }
 
         FieldEntry entry = fields.get(fieldInfo.number);
@@ -277,17 +264,27 @@ public final void search(String field, float[] target, KnnCollector knnCollector
             target,
             ivfClusters
         );
+        if (nProbe == DYNAMIC_NPROBE) {
+            // empirically based, and a good dynamic to get decent recall while scaling a la "efSearch"
+            // scaling by the number of centroids vs. the nearest neighbors requested
+            // not perfect, but a comparative heuristic.
+            // we might want to utilize the total vector count as well, but this is a good start
+            nProbe = (int) Math.round(Math.log10(centroidQueryScorer.size()) * Math.sqrt(knnCollector.k()));
+            // clip to be between 1 and the number of centroids
+            nProbe = Math.max(Math.min(nProbe, centroidQueryScorer.size()), 1);
+        }
         final NeighborQueue centroidQueue = scorePostingLists(fieldInfo, knnCollector, centroidQueryScorer, nProbe);
         PostingVisitor scorer = getPostingVisitor(fieldInfo, ivfClusters, target, needsScoring);
         int centroidsVisited = 0;
         long expectedDocs = 0;
         long actualDocs = 0;
         // initially we visit only the "centroids to search"
-        while (centroidQueue.size() > 0 && centroidsVisited < nProbe) {
+        while (centroidQueue.size() > 0 && centroidsVisited < nProbe && actualDocs < knnCollector.k()) {
             ++centroidsVisited;
             // todo do we actually need to know the score???
             int centroidOrdinal = centroidQueue.pop();
-            // todo do we need direct access to the raw centroid???
+            // todo do we need direct access to the raw centroid???, this is used for quantizing, maybe hydrating and quantizing
+            // is enough?
             expectedDocs += scorer.resetPostingsScorer(centroidOrdinal, centroidQueryScorer.centroid(centroidOrdinal));
             actualDocs += scorer.visit(knnCollector);
         }

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

@@ -41,6 +41,7 @@ public class MapperFeatures implements FeatureSpecification {
         "mapper.unknown_field_mapping_update_error_message"
     );
     static final NodeFeature NPE_ON_DIMS_UPDATE_FIX = new NodeFeature("mapper.npe_on_dims_update_fix");
+    static final NodeFeature IVF_FORMAT_CLUSTER_FEATURE = new NodeFeature("mapper.ivf_format_cluster_feature");
 
     @Override
     public Set<NodeFeature> getTestFeatures() {
@@ -68,7 +69,8 @@ public Set<NodeFeature> getTestFeatures() {
             DateFieldMapper.INVALID_DATE_FIX,
             NPE_ON_DIMS_UPDATE_FIX,
             RESCORE_ZERO_VECTOR_QUANTIZED_VECTOR_MAPPING,
-            USE_DEFAULT_OVERSAMPLE_VALUE_FOR_BBQ
+            USE_DEFAULT_OVERSAMPLE_VALUE_FOR_BBQ,
+            IVF_FORMAT_CLUSTER_FEATURE
         );
     }
 }

server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java

@@ -39,6 +39,7 @@
 import org.apache.lucene.util.VectorUtil;
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.settings.Setting;
+import org.elasticsearch.common.util.FeatureFlag;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.index.IndexVersion;
@@ -48,6 +49,7 @@
 import org.elasticsearch.index.codec.vectors.ES814HnswScalarQuantizedVectorsFormat;
 import org.elasticsearch.index.codec.vectors.ES815BitFlatVectorFormat;
 import org.elasticsearch.index.codec.vectors.ES815HnswBitVectorsFormat;
+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.fielddata.FieldDataContext;
@@ -62,6 +64,7 @@
 import org.elasticsearch.index.mapper.Mapper;
 import org.elasticsearch.index.mapper.MapperBuilderContext;
 import org.elasticsearch.index.mapper.MapperParsingException;
+import org.elasticsearch.index.mapper.MappingLookup;
 import org.elasticsearch.index.mapper.MappingParser;
 import org.elasticsearch.index.mapper.NumberFieldMapper;
 import org.elasticsearch.index.mapper.SimpleMappedFieldType;
@@ -78,6 +81,7 @@
 import org.elasticsearch.search.vectors.ESDiversifyingChildrenFloatKnnVectorQuery;
 import org.elasticsearch.search.vectors.ESKnnByteVectorQuery;
 import org.elasticsearch.search.vectors.ESKnnFloatVectorQuery;
+import org.elasticsearch.search.vectors.IVFKnnFloatVectorQuery;
 import org.elasticsearch.search.vectors.RescoreKnnVectorQuery;
 import org.elasticsearch.search.vectors.VectorData;
 import org.elasticsearch.search.vectors.VectorSimilarityQuery;
@@ -106,6 +110,8 @@
 import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_INDEX_VERSION_CREATED;
 import static org.elasticsearch.common.Strings.format;
 import static org.elasticsearch.common.xcontent.XContentParserUtils.ensureExpectedToken;
+import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.MAX_VECTORS_PER_CLUSTER;
+import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.MIN_VECTORS_PER_CLUSTER;
 
 /**
  * A {@link FieldMapper} for indexing a dense vector of floats.
@@ -115,6 +121,8 @@ public class DenseVectorFieldMapper extends FieldMapper {
     private static final float EPS = 1e-3f;
     public static final int BBQ_MIN_DIMS = 64;
 
+    public static final FeatureFlag IVF_FORMAT = new FeatureFlag("ivf_format");
+
     public static boolean isNotUnitVector(float magnitude) {
         return Math.abs(magnitude - 1.0f) > EPS;
     }
@@ -1594,14 +1602,63 @@ public boolean supportsElementType(ElementType elementType) {
                 return elementType == ElementType.FLOAT;
             }
 
+            @Override
+            public boolean supportsDimension(int dims) {
+                return dims >= BBQ_MIN_DIMS;
+            }
+        },
+        BBQ_IVF("bbq_ivf", true) {
+            @Override
+            public IndexOptions parseIndexOptions(String fieldName, Map<String, ?> indexOptionsMap, IndexVersion indexVersion) {
+                Object clusterSizeNode = indexOptionsMap.remove("cluster_size");
+                int clusterSize = IVFVectorsFormat.DEFAULT_VECTORS_PER_CLUSTER;
+                if (clusterSizeNode != null) {
+                    clusterSize = XContentMapValues.nodeIntegerValue(clusterSizeNode);
+                    if (clusterSize < MIN_VECTORS_PER_CLUSTER || clusterSize > MAX_VECTORS_PER_CLUSTER) {
+                        throw new IllegalArgumentException(
+                            "cluster_size must be between "
+                                + MIN_VECTORS_PER_CLUSTER
+                                + " and "
+                                + MAX_VECTORS_PER_CLUSTER
+                                + ", got: "
+                                + clusterSize
+                        );
+                    }
+                }
+                RescoreVector rescoreVector = RescoreVector.fromIndexOptions(indexOptionsMap, indexVersion);
+                if (rescoreVector == null) {
+                    rescoreVector = new RescoreVector(DEFAULT_OVERSAMPLE);
+                }
+                Object nProbeNode = indexOptionsMap.remove("default_n_probe");
+                int nProbe = -1;
+                if (nProbeNode != null) {
+                    nProbe = XContentMapValues.nodeIntegerValue(nProbeNode);
+                    if (nProbe < 1 && nProbe != -1) {
+                        throw new IllegalArgumentException(
+                            "default_n_probe must be at least 1 or exactly -1, got: " + nProbe + " for field [" + fieldName + "]"
+                        );
+                    }
+                }
+                MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap);
+                return new BBQIVFIndexOptions(clusterSize, nProbe, rescoreVector);
+            }
+
+            @Override
+            public boolean supportsElementType(ElementType elementType) {
+                return elementType == ElementType.FLOAT;
+            }
+
             @Override
             public boolean supportsDimension(int dims) {
                 return dims >= BBQ_MIN_DIMS;
             }
         };
 
         static Optional<VectorIndexType> fromString(String type) {
-            return Stream.of(VectorIndexType.values()).filter(vectorIndexType -> vectorIndexType.name.equals(type)).findFirst();
+            return Stream.of(VectorIndexType.values())
+                .filter(vectorIndexType -> vectorIndexType != VectorIndexType.BBQ_IVF || IVF_FORMAT.isEnabled())
+                .filter(vectorIndexType -> vectorIndexType.name.equals(type))
+                .findFirst();
         }
 
         private final String name;
@@ -2100,6 +2157,54 @@ public boolean validateDimension(int dim, boolean throwOnError) {
 
     }
 
+    static class BBQIVFIndexOptions extends QuantizedIndexOptions {
+        final int clusterSize;
+        final int defaultNProbe;
+
+        BBQIVFIndexOptions(int clusterSize, int defaultNProbe, RescoreVector rescoreVector) {
+            super(VectorIndexType.BBQ_IVF, rescoreVector);
+            this.clusterSize = clusterSize;
+            this.defaultNProbe = defaultNProbe;
+        }
+
+        @Override
+        KnnVectorsFormat getVectorsFormat(ElementType elementType) {
+            assert elementType == ElementType.FLOAT;
+            return new IVFVectorsFormat(clusterSize);
+        }
+
+        @Override
+        boolean updatableTo(IndexOptions update) {
+            return update.type.equals(this.type);
+        }
+
+        @Override
+        boolean doEquals(IndexOptions other) {
+            BBQIVFIndexOptions that = (BBQIVFIndexOptions) other;
+            return clusterSize == that.clusterSize
+                && defaultNProbe == that.defaultNProbe
+                && Objects.equals(rescoreVector, that.rescoreVector);
+        }
+
+        @Override
+        int doHashCode() {
+            return Objects.hash(clusterSize, defaultNProbe, rescoreVector);
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            builder.startObject();
+            builder.field("type", type);
+            builder.field("cluster_size", clusterSize);
+            builder.field("default_n_probe", defaultNProbe);
+            if (rescoreVector != null) {
+                rescoreVector.toXContent(builder, params);
+            }
+            builder.endObject();
+            return builder;
+        }
+    }
+
     public record RescoreVector(float oversample) implements ToXContentObject {
         static final String NAME = "rescore_vector";
         static final String OVERSAMPLE = "oversample";
@@ -2411,17 +2516,25 @@ && isNotUnitVector(squaredMagnitude)) {
                 adjustedK = Math.min((int) Math.ceil(k * oversample), OVERSAMPLE_LIMIT);
                 numCands = Math.max(adjustedK, numCands);
             }
-            Query knnQuery = parentFilter != null
-                ? new ESDiversifyingChildrenFloatKnnVectorQuery(
-                    name(),
-                    queryVector,
-                    filter,
-                    adjustedK,
-                    numCands,
-                    parentFilter,
-                    knnSearchStrategy
-                )
-                : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, knnSearchStrategy);
+            if (parentFilter != null && indexOptions instanceof BBQIVFIndexOptions) {
+                throw new IllegalArgumentException("IVF index does not support nested queries");
+            }
+            Query knnQuery;
+            if (indexOptions instanceof BBQIVFIndexOptions bbqIndexOptions) {
+                knnQuery = new IVFKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, bbqIndexOptions.defaultNProbe);
+            } else {
+                knnQuery = parentFilter != null
+                    ? new ESDiversifyingChildrenFloatKnnVectorQuery(
+                        name(),
+                        queryVector,
+                        filter,
+                        adjustedK,
+                        numCands,
+                        parentFilter,
+                        knnSearchStrategy
+                    )
+                    : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, knnSearchStrategy);
+            }
             if (rescore) {
                 knnQuery = new RescoreKnnVectorQuery(
                     name(),
@@ -2651,6 +2764,19 @@ public FieldMapper.Builder getMergeBuilder() {
         return new Builder(leafName(), indexCreatedVersion).init(this);
     }
 
+    @Override
+    public void doValidate(MappingLookup mappers) {
+        if (indexOptions instanceof BBQIVFIndexOptions && mappers.nestedLookup().getNestedParent(fullPath()) != null) {
+            throw new IllegalArgumentException(
+                "["
+                    + CONTENT_TYPE
+                    + "] fields with index type ["
+                    + indexOptions.type
+                    + "] cannot be indexed if they're within [nested] mappings"
+            );
+        }
+    }
+
     private static IndexOptions parseIndexOptions(String fieldName, Object propNode, IndexVersion indexVersion) {
         @SuppressWarnings("unchecked")
         Map<String, ?> indexOptionsMap = (Map<String, ?>) propNode;