This adds a new experimental IVF format behind a feature flag

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

@@ -61,14 +61,24 @@ public class IVFVectorsFormat extends KnnVectorsFormat {
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
 
-    private static final int DEFAULT_VECTORS_PER_CLUSTER = 1000;
+    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 +100,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/IVFVectorsReader.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/mapper/MapperFeatures.java

@@ -11,7 +11,9 @@
 
 import org.elasticsearch.features.FeatureSpecification;
 import org.elasticsearch.features.NodeFeature;
+import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
 
+import java.util.HashSet;
 import java.util.Set;
 
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.RESCORE_VECTOR_QUANTIZED_VECTOR_MAPPING;
@@ -41,34 +43,41 @@ 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() {
-        return Set.of(
-            RangeFieldMapper.DATE_RANGE_INDEXING_FIX,
-            IgnoredSourceFieldMapper.DONT_EXPAND_DOTS_IN_IGNORED_SOURCE,
-            SourceFieldMapper.REMOVE_SYNTHETIC_SOURCE_ONLY_VALIDATION,
-            SourceFieldMapper.SOURCE_MODE_FROM_INDEX_SETTING,
-            IgnoredSourceFieldMapper.IGNORED_SOURCE_AS_TOP_LEVEL_METADATA_ARRAY_FIELD,
-            IgnoredSourceFieldMapper.ALWAYS_STORE_OBJECT_ARRAYS_IN_NESTED_OBJECTS,
-            MapperService.LOGSDB_DEFAULT_IGNORE_DYNAMIC_BEYOND_LIMIT,
-            DocumentParser.FIX_PARSING_SUBOBJECTS_FALSE_DYNAMIC_FALSE,
-            CONSTANT_KEYWORD_SYNTHETIC_SOURCE_WRITE_FIX,
-            META_FETCH_FIELDS_ERROR_CODE_CHANGED,
-            SPARSE_VECTOR_STORE_SUPPORT,
-            COUNTED_KEYWORD_SYNTHETIC_SOURCE_NATIVE_SUPPORT,
-            SORT_FIELDS_CHECK_FOR_NESTED_OBJECT_FIX,
-            DYNAMIC_HANDLING_IN_COPY_TO,
-            TSDB_NESTED_FIELD_SUPPORT,
-            SourceFieldMapper.SYNTHETIC_RECOVERY_SOURCE,
-            ObjectMapper.SUBOBJECTS_FALSE_MAPPING_UPDATE_FIX,
-            UKNOWN_FIELD_MAPPING_UPDATE_ERROR_MESSAGE,
-            DOC_VALUES_SKIPPER,
-            RESCORE_VECTOR_QUANTIZED_VECTOR_MAPPING,
-            DateFieldMapper.INVALID_DATE_FIX,
-            NPE_ON_DIMS_UPDATE_FIX,
-            RESCORE_ZERO_VECTOR_QUANTIZED_VECTOR_MAPPING,
-            USE_DEFAULT_OVERSAMPLE_VALUE_FOR_BBQ
+        HashSet<NodeFeature> features = new HashSet<>(
+            Set.of(
+                RangeFieldMapper.DATE_RANGE_INDEXING_FIX,
+                IgnoredSourceFieldMapper.DONT_EXPAND_DOTS_IN_IGNORED_SOURCE,
+                SourceFieldMapper.REMOVE_SYNTHETIC_SOURCE_ONLY_VALIDATION,
+                SourceFieldMapper.SOURCE_MODE_FROM_INDEX_SETTING,
+                IgnoredSourceFieldMapper.IGNORED_SOURCE_AS_TOP_LEVEL_METADATA_ARRAY_FIELD,
+                IgnoredSourceFieldMapper.ALWAYS_STORE_OBJECT_ARRAYS_IN_NESTED_OBJECTS,
+                MapperService.LOGSDB_DEFAULT_IGNORE_DYNAMIC_BEYOND_LIMIT,
+                DocumentParser.FIX_PARSING_SUBOBJECTS_FALSE_DYNAMIC_FALSE,
+                CONSTANT_KEYWORD_SYNTHETIC_SOURCE_WRITE_FIX,
+                META_FETCH_FIELDS_ERROR_CODE_CHANGED,
+                SPARSE_VECTOR_STORE_SUPPORT,
+                COUNTED_KEYWORD_SYNTHETIC_SOURCE_NATIVE_SUPPORT,
+                SORT_FIELDS_CHECK_FOR_NESTED_OBJECT_FIX,
+                DYNAMIC_HANDLING_IN_COPY_TO,
+                TSDB_NESTED_FIELD_SUPPORT,
+                SourceFieldMapper.SYNTHETIC_RECOVERY_SOURCE,
+                ObjectMapper.SUBOBJECTS_FALSE_MAPPING_UPDATE_FIX,
+                UKNOWN_FIELD_MAPPING_UPDATE_ERROR_MESSAGE,
+                DOC_VALUES_SKIPPER,
+                RESCORE_VECTOR_QUANTIZED_VECTOR_MAPPING,
+                DateFieldMapper.INVALID_DATE_FIX,
+                NPE_ON_DIMS_UPDATE_FIX,
+                RESCORE_ZERO_VECTOR_QUANTIZED_VECTOR_MAPPING,
+                USE_DEFAULT_OVERSAMPLE_VALUE_FOR_BBQ
+            )
         );
+        if (DenseVectorFieldMapper.IVF_FORMAT.isEnabled()) {
+            features.add(IVF_FORMAT_CLUSTER_FEATURE);
+        }
+        return features;
     }
 }