Add OffHeapBytesSizes for Stats (#138433)

john-wagster · web-flow · commit 943990e7f884 · 2025-11-21T15:37:56.000-06:00
* added in off heaps bytes sizes for stats endpoints
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/220_dense_vector_node_index_stats.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/220_dense_vector_node_index_stats.yml
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsReader.java
@@ -27,7 +27,6 @@
 import org.elasticsearch.simdvec.ESVectorUtil;
 
 import java.io.IOException;
-import java.util.Map;
 
 import static org.apache.lucene.codecs.lucene102.Lucene102BinaryQuantizedVectorsFormat.QUERY_BITS;
 import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
@@ -394,11 +393,6 @@ public PostingVisitor getPostingVisitor(FieldInfo fieldInfo, IndexInput indexInp
         return new MemorySegmentPostingsVisitor(target, indexInput, entry, fieldInfo, acceptDocs);
     }
 
-    @Override
-    public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
-        return Map.of();
-    }
-
     private static class MemorySegmentPostingsVisitor implements PostingVisitor {
         final long quantizedByteLength;
         final IndexInput indexInput;
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsReader.java
@@ -42,6 +42,8 @@
 import java.util.Map;
 
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS;
+import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.CENTROID_EXTENSION;
+import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.CLUSTER_EXTENSION;
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.DYNAMIC_VISIT_RATIO;
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO;
 
@@ -87,20 +89,8 @@ protected IVFVectorsReader(SegmentReadState state, GenericFlatVectorReaders.Load
             } finally {
                 CodecUtil.checkFooter(ivfMeta, priorE);
             }
-            ivfCentroids = openDataInput(
-                state,
-                versionMeta,
-                ES920DiskBBQVectorsFormat.CENTROID_EXTENSION,
-                ES920DiskBBQVectorsFormat.NAME,
-                state.context
-            );
-            ivfClusters = openDataInput(
-                state,
-                versionMeta,
-                ES920DiskBBQVectorsFormat.CLUSTER_EXTENSION,
-                ES920DiskBBQVectorsFormat.NAME,
-                state.context
-            );
+            ivfCentroids = openDataInput(state, versionMeta, CENTROID_EXTENSION, ES920DiskBBQVectorsFormat.NAME, state.context);
+            ivfClusters = openDataInput(state, versionMeta, CLUSTER_EXTENSION, ES920DiskBBQVectorsFormat.NAME, state.context);
             success = true;
         } finally {
             if (success == false) {
@@ -388,12 +378,9 @@ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
             assert fieldInfo.getVectorEncoding() == VectorEncoding.BYTE;
             return raw;
         }
-        return raw;  // for now just return the size of raw
 
-        // TODO: determine desired off off-heap requirements
-        // var centroids = Map.of(EXTENSION, fe.xxxLength());
-        // var clusters = Map.of(EXTENSION, fe.yyyLength());
-        // return KnnVectorsReader.mergeOffHeapByteSizeMaps(raw, centroids, clusters);
+        var centroidsClusters = Map.of(CENTROID_EXTENSION, fe.centroidLength, CLUSTER_EXTENSION, fe.postingListLength);
+        return KnnVectorsReader.mergeOffHeapByteSizeMaps(raw, centroidsClusters);
     }
 
     @Override
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsReader.java
@@ -541,7 +541,8 @@ public PostingVisitor getPostingVisitor(FieldInfo fieldInfo, IndexInput indexInp
 
     @Override
     public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
-        return Map.of();
+        // TODO: override if adding new files
+        return super.getOffHeapByteSize(fieldInfo);
     }
 
     private static class MemorySegmentPostingsVisitor implements PostingVisitor {
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java b/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java
@@ -63,6 +63,7 @@ public class MapperFeatures implements FeatureSpecification {
     static final NodeFeature BASE64_DENSE_VECTORS = new NodeFeature("mapper.base64_dense_vectors");
     public static final NodeFeature GENERIC_VECTOR_FORMAT = new NodeFeature("mapper.vectors.generic_vector_format");
     public static final NodeFeature FIX_DENSE_VECTOR_WRONG_FIELDS = new NodeFeature("mapper.fix_dense_vector_wrong_fields");
+    static final NodeFeature BBQ_DISK_STATS_SUPPORT = new NodeFeature("mapper.bbq_disk_stats_support");
 
     @Override
     public Set<NodeFeature> getTestFeatures() {
@@ -106,7 +107,8 @@ public Set<NodeFeature> getTestFeatures() {
             INDEX_MAPPING_IGNORE_DYNAMIC_BEYOND_FIELD_NAME_LIMIT,
             EXCLUDE_VECTORS_DOCVALUE_BUGFIX,
             BASE64_DENSE_VECTORS,
-            FIX_DENSE_VECTOR_WRONG_FIELDS
+            FIX_DENSE_VECTOR_WRONG_FIELDS,
+            BBQ_DISK_STATS_SUPPORT
         );
         if (ES93GenericFlatVectorsFormat.GENERIC_VECTOR_FORMAT.isEnabled()) {
             features = new HashSet<>(features);
diff --git a/server/src/main/java/org/elasticsearch/index/shard/DenseVectorStats.java b/server/src/main/java/org/elasticsearch/index/shard/DenseVectorStats.java
@@ -112,7 +112,7 @@ public Map<String, Map<String, Long>> offHeapStats() {
 
     private Map<String, Long> getTotalsByCategory() {
         if (offHeapStats == null) {
-            return Map.of("veb", 0L, "vec", 0L, "veq", 0L, "vex", 0L);
+            return Map.of("veb", 0L, "vec", 0L, "veq", 0L, "vex", 0L, "cenivf", 0L, "clivf", 0L);
         } else {
             return offHeapStats.entrySet()
                 .stream()
@@ -140,6 +140,8 @@ private void toXContentWithFields(XContentBuilder builder, Params params) throws
         builder.humanReadableField("total_vec_size_bytes", "total_vec_size", ofBytes(totals.getOrDefault("vec", 0L)));
         builder.humanReadableField("total_veq_size_bytes", "total_veq_size", ofBytes(totals.getOrDefault("veq", 0L)));
         builder.humanReadableField("total_vex_size_bytes", "total_vex_size", ofBytes(totals.getOrDefault("vex", 0L)));
+        builder.humanReadableField("total_cenivf_size_bytes", "total_cenivf_size", ofBytes(totals.getOrDefault("cenivf", 0L)));
+        builder.humanReadableField("total_clivf_size_bytes", "total_clivf_size", ofBytes(totals.getOrDefault("clivf", 0L)));
         if (params.paramAsBoolean(INCLUDE_PER_FIELD_STATS, false) && offHeapStats != null && offHeapStats.size() > 0) {
             toXContentWithPerFieldStats(builder);
         }
diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQBFloat16VectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQBFloat16VectorsFormatTests.java
@@ -30,7 +30,6 @@
 
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_CENTROIDS_PER_PARENT_CLUSTER;
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_VECTORS_PER_CLUSTER;
-import static org.hamcrest.Matchers.anEmptyMap;
 import static org.hamcrest.Matchers.equalTo;
 
 public class ES920DiskBBQBFloat16VectorsFormatTests extends BaseBFloat16KnnVectorsFormatTestCase {
@@ -102,9 +101,8 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx
             }
             var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
             long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum();
-            // IVF doesn't report stats at the moment
-            assertThat(offHeap, anEmptyMap());
-            assertThat(totalByteSize, equalTo(0L));
+            assertThat(offHeap.size(), equalTo(3));
+            assertThat(totalByteSize, equalTo(offHeap.values().stream().mapToLong(Long::longValue).sum()));
         } else {
             throw new AssertionError("unexpected:" + r.getClass());
         }
diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormatTests.java
@@ -49,7 +49,6 @@
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MAX_VECTORS_PER_CLUSTER;
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_CENTROIDS_PER_PARENT_CLUSTER;
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_VECTORS_PER_CLUSTER;
-import static org.hamcrest.Matchers.anEmptyMap;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasToString;
 
@@ -122,9 +121,8 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx
             }
             var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
             long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum();
-            // IVF doesn't report stats at the moment
-            assertThat(offHeap, anEmptyMap());
-            assertThat(totalByteSize, equalTo(0L));
+            assertThat(offHeap.size(), equalTo(3));
+            assertThat(totalByteSize, equalTo(offHeap.values().stream().mapToLong(Long::longValue).sum()));
         } else {
             throw new AssertionError("unexpected:" + r.getClass());
         }
@@ -164,7 +162,7 @@ public void testSimpleOffHeapSize() throws IOException {
                     }
                     var fieldInfo = r.getFieldInfos().fieldInfo("f");
                     var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
-                    assertEquals(0, offHeap.size());
+                    assertEquals(3, offHeap.size());
                 }
             }
         }
diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQBFloat16VectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQBFloat16VectorsFormatTests.java
@@ -31,7 +31,6 @@
 
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_CENTROIDS_PER_PARENT_CLUSTER;
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_VECTORS_PER_CLUSTER;
-import static org.hamcrest.Matchers.anEmptyMap;
 import static org.hamcrest.Matchers.equalTo;
 
 public class ESNextDiskBBQBFloat16VectorsFormatTests extends BaseBFloat16KnnVectorsFormatTestCase {
@@ -108,9 +107,8 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx
             }
             var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
             long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum();
-            // IVF doesn't report stats at the moment
-            assertThat(offHeap, anEmptyMap());
-            assertThat(totalByteSize, equalTo(0L));
+            assertThat(offHeap.size(), equalTo(3));
+            assertThat(totalByteSize, equalTo(offHeap.values().stream().mapToLong(Long::longValue).sum()));
         } else {
             throw new AssertionError("unexpected:" + r.getClass());
         }
diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsFormatTests.java
@@ -53,7 +53,6 @@
 import static org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat.MAX_VECTORS_PER_CLUSTER;
 import static org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat.MIN_CENTROIDS_PER_PARENT_CLUSTER;
 import static org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat.MIN_VECTORS_PER_CLUSTER;
-import static org.hamcrest.Matchers.anEmptyMap;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.oneOf;
@@ -127,9 +126,8 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx
             }
             var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
             long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum();
-            // IVF doesn't report stats at the moment
-            assertThat(offHeap, anEmptyMap());
-            assertThat(totalByteSize, equalTo(0L));
+            assertThat(offHeap.size(), equalTo(3));
+            assertThat(totalByteSize, equalTo(offHeap.values().stream().mapToLong(Long::longValue).sum()));
         } else {
             throw new AssertionError("unexpected:" + r.getClass());
         }
@@ -177,7 +175,7 @@ public void testSimpleOffHeapSize() throws IOException {
                     }
                     var fieldInfo = r.getFieldInfos().fieldInfo("f");
                     var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
-                    assertEquals(0, offHeap.size());
+                    assertEquals(3, offHeap.size());
                 }
             }
         }
diff --git a/server/src/test/java/org/elasticsearch/index/shard/DenseVectorStatsTests.java b/server/src/test/java/org/elasticsearch/index/shard/DenseVectorStatsTests.java
@@ -87,7 +87,10 @@ public void testBasicEquality() {
     }
 
     public void testBasicXContent() throws IOException {
-        var stats = new DenseVectorStats(5L, Map.of("foo", Map.of("vec", 9L), "bar", Map.of("vec", 14L, "vex", 1L, "veb", 3L)));
+        var stats = new DenseVectorStats(
+            5L,
+            Map.of("foo", Map.of("vec", 9L), "bar", Map.of("vec", 14L, "vex", 1L, "veb", 3L, "cenivf", 7L, "clivf", 2L))
+        );
 
         XContentBuilder builder = XContentFactory.jsonBuilder().prettyPrint();
         builder.startObject();
@@ -110,11 +113,13 @@ public void testBasicXContent() throws IOException {
               "dense_vector" : {
                 "value_count" : 5,
                 "off_heap" : {
-                  "total_size_bytes" : 27,
+                  "total_size_bytes" : 36,
                   "total_veb_size_bytes" : 3,
                   "total_vec_size_bytes" : 23,
                   "total_veq_size_bytes" : 0,
-                  "total_vex_size_bytes" : 1
+                  "total_vex_size_bytes" : 1,
+                  "total_cenivf_size_bytes" : 7,
+                  "total_clivf_size_bytes" : 2
                 }
               }
             }""";
@@ -129,13 +134,17 @@ public void testBasicXContent() throws IOException {
               "dense_vector" : {
                 "value_count" : 5,
                 "off_heap" : {
-                  "total_size_bytes" : 27,
+                  "total_size_bytes" : 36,
                   "total_veb_size_bytes" : 3,
                   "total_vec_size_bytes" : 23,
                   "total_veq_size_bytes" : 0,
                   "total_vex_size_bytes" : 1,
+                  "total_cenivf_size_bytes" : 7,
+                  "total_clivf_size_bytes" : 2,
                   "fielddata" : {
                     "bar" : {
+                      "cenivf_size_bytes" : 7,
+                      "clivf_size_bytes" : 2,
                       "veb_size_bytes" : 3,
                       "vec_size_bytes" : 14,
                       "vex_size_bytes" : 1
@@ -168,7 +177,9 @@ public void testBasicXContent() throws IOException {
                           "total_veb_size_bytes" : 0,
                           "total_vec_size_bytes" : 0,
                           "total_veq_size_bytes" : 0,
-                          "total_vex_size_bytes" : 0
+                          "total_vex_size_bytes" : 0,
+                          "total_cenivf_size_bytes" : 0,
+                          "total_clivf_size_bytes" : 0
                         }
                       }
                     }""";
@@ -202,6 +213,10 @@ public void testXContentHumanReadable() throws IOException {
                   "total_veq_size_bytes" : 1099511627776,
                   "total_vex_size" : "190.7mb",
                   "total_vex_size_bytes" : 200000000,
+                  "total_cenivf_size" : "0b",
+                  "total_cenivf_size_bytes" : 0,
+                  "total_clivf_size" : "0b",
+                  "total_clivf_size_bytes" : 0,
                   "fielddata" : {
                     "bar" : {
                       "veb_size" : "1kb",

Original file line number	Diff line number	Diff line change
`@@ -541,7 +541,8 @@ public PostingVisitor getPostingVisitor(FieldInfo fieldInfo, IndexInput indexInp`
`541`	`541`
`542`	`542`	`@Override`
`543`	`543`	`public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {`
`544`		`- return Map.of();`
	`544`	`+ // TODO: override if adding new files`
	`545`	`+ return super.getOffHeapByteSize(fieldInfo);`
`545`	`546`	`}`
`546`	`547`
`547`	`548`	`private static class MemorySegmentPostingsVisitor implements PostingVisitor {`
Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,6 @@`
`49`	`49`	`import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MAX_VECTORS_PER_CLUSTER;`
`50`	`50`	`import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_CENTROIDS_PER_PARENT_CLUSTER;`
`51`	`51`	`import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_VECTORS_PER_CLUSTER;`
`52`		`-import static org.hamcrest.Matchers.anEmptyMap;`
`53`	`52`	`import static org.hamcrest.Matchers.equalTo;`
`54`	`53`	`import static org.hamcrest.Matchers.hasToString;`
`55`	`54`
`@@ -122,9 +121,8 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx`
`122`	`121`	`}`
`123`	`122`	`var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);`
`124`	`123`	`long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum();`
`125`		`- // IVF doesn't report stats at the moment`
`126`		`- assertThat(offHeap, anEmptyMap());`
`127`		`- assertThat(totalByteSize, equalTo(0L));`
	`124`	`+ assertThat(offHeap.size(), equalTo(3));`
	`125`	`+ assertThat(totalByteSize, equalTo(offHeap.values().stream().mapToLong(Long::longValue).sum()));`
`128`	`126`	`} else {`
`129`	`127`	`throw new AssertionError("unexpected:" + r.getClass());`
`130`	`128`	`}`
`@@ -164,7 +162,7 @@ public void testSimpleOffHeapSize() throws IOException {`
`164`	`162`	`}`
`165`	`163`	`var fieldInfo = r.getFieldInfos().fieldInfo("f");`
`166`	`164`	`var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);`
`167`		`- assertEquals(0, offHeap.size());`
	`165`	`+ assertEquals(3, offHeap.size());`
`168`	`166`	`}`
`169`	`167`	`}`
`170`	`168`	`}`
Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,6 @@`
`53`	`53`	`import static org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat.MAX_VECTORS_PER_CLUSTER;`
`54`	`54`	`import static org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat.MIN_CENTROIDS_PER_PARENT_CLUSTER;`
`55`	`55`	`import static org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat.MIN_VECTORS_PER_CLUSTER;`
`56`		`-import static org.hamcrest.Matchers.anEmptyMap;`
`57`	`56`	`import static org.hamcrest.Matchers.equalTo;`
`58`	`57`	`import static org.hamcrest.Matchers.is;`
`59`	`58`	`import static org.hamcrest.Matchers.oneOf;`
`@@ -127,9 +126,8 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx`
`127`	`126`	`}`
`128`	`127`	`var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);`
`129`	`128`	`long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum();`
`130`		`- // IVF doesn't report stats at the moment`
`131`		`- assertThat(offHeap, anEmptyMap());`
`132`		`- assertThat(totalByteSize, equalTo(0L));`
	`129`	`+ assertThat(offHeap.size(), equalTo(3));`
	`130`	`+ assertThat(totalByteSize, equalTo(offHeap.values().stream().mapToLong(Long::longValue).sum()));`
`133`	`131`	`} else {`
`134`	`132`	`throw new AssertionError("unexpected:" + r.getClass());`
`135`	`133`	`}`
`@@ -177,7 +175,7 @@ public void testSimpleOffHeapSize() throws IOException {`
`177`	`175`	`}`
`178`	`176`	`var fieldInfo = r.getFieldInfos().fieldInfo("f");`
`179`	`177`	`var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);`
`180`		`- assertEquals(0, offHeap.size());`
	`178`	`+ assertEquals(3, offHeap.size());`
`181`	`179`	`}`
`182`	`180`	`}`
`183`	`181`	`}`