Add totalPostingBytes to ShardFieldStats

Author: jordan-powers

server/src/main/java/org/elasticsearch/common/lucene/Lucene.java

@@ -739,15 +739,26 @@ public static Version parseVersionLenient(String toParse, Version defaultValue)
      * If no SegmentReader can be extracted an {@link IllegalStateException} is thrown.
      */
     public static SegmentReader segmentReader(LeafReader reader) {
+        SegmentReader segmentReader = tryUnwrapSegmentReader(reader);
+        if (segmentReader == null) {
+            throw new IllegalStateException("Can not extract segment reader from given index reader [" + reader + "]");
+        }
+        return segmentReader;
+    }
+
+    /**
+     * Tries to extract a segment reader from the given index reader. Unlike {@link #segmentReader(LeafReader)} this method returns
+     * null if no SegmentReader can be unwrapped instead of throwing an exception.
+     */
+    public static SegmentReader tryUnwrapSegmentReader(LeafReader reader) {
         if (reader instanceof SegmentReader) {
             return (SegmentReader) reader;
         } else if (reader instanceof final FilterLeafReader fReader) {
-            return segmentReader(FilterLeafReader.unwrap(fReader));
+            return tryUnwrapSegmentReader(FilterLeafReader.unwrap(fReader));
         } else if (reader instanceof final FilterCodecReader fReader) {
-            return segmentReader(FilterCodecReader.unwrap(fReader));
+            return tryUnwrapSegmentReader(FilterCodecReader.unwrap(fReader));
         }
-        // hard fail - we can't get a SegmentReader
-        throw new IllegalStateException("Can not extract segment reader from given index reader [" + reader + "]");
+        return null;
     }
 
     @SuppressForbidden(reason = "Version#parseLeniently() used in a central place")

server/src/main/java/org/elasticsearch/index/engine/Engine.java

@@ -62,6 +62,7 @@
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.VersionType;
 import org.elasticsearch.index.codec.FieldInfosWithUsages;
+import org.elasticsearch.index.codec.TrackingPostingsInMemoryBytesCodec;
 import org.elasticsearch.index.codec.vectors.reflect.OffHeapByteSizeUtils;
 import org.elasticsearch.index.mapper.DocumentParser;
 import org.elasticsearch.index.mapper.LuceneDocument;
@@ -275,6 +276,7 @@ protected static ShardFieldStats shardFieldStats(List<LeafReaderContext> leaves)
         int numSegments = 0;
         int totalFields = 0;
         long usages = 0;
+        long totalPostingBytes = 0;
         for (LeafReaderContext leaf : leaves) {
             numSegments++;
             var fieldInfos = leaf.reader().getFieldInfos();
@@ -286,8 +288,17 @@ protected static ShardFieldStats shardFieldStats(List<LeafReaderContext> leaves)
             } else {
                 usages = -1;
             }
+            SegmentReader segmentReader = Lucene.tryUnwrapSegmentReader(leaf.reader());
+            if (segmentReader != null) {
+                String postingBytes = segmentReader.getSegmentInfo().info.getAttribute(
+                    TrackingPostingsInMemoryBytesCodec.IN_MEMORY_POSTINGS_BYTES_KEY
+                );
+                if (postingBytes != null) {
+                    totalPostingBytes += Long.parseLong(postingBytes);
+                }
+            }
         }
-        return new ShardFieldStats(numSegments, totalFields, usages);
+        return new ShardFieldStats(numSegments, totalFields, usages, totalPostingBytes);
     }
 
     /**

server/src/main/java/org/elasticsearch/index/engine/InternalEngine.java

@@ -79,6 +79,7 @@
 import org.elasticsearch.index.IndexVersions;
 import org.elasticsearch.index.VersionType;
 import org.elasticsearch.index.cache.query.TrivialQueryCachingPolicy;
+import org.elasticsearch.index.codec.TrackingPostingsInMemoryBytesCodec;
 import org.elasticsearch.index.mapper.DocumentParser;
 import org.elasticsearch.index.mapper.IdFieldMapper;
 import org.elasticsearch.index.mapper.LuceneDocument;
@@ -2777,7 +2778,7 @@ private IndexWriterConfig getIndexWriterConfig() {
         iwc.setMaxFullFlushMergeWaitMillis(-1);
         iwc.setSimilarity(engineConfig.getSimilarity());
         iwc.setRAMBufferSizeMB(engineConfig.getIndexingBufferSize().getMbFrac());
-        iwc.setCodec(engineConfig.getCodec());
+        iwc.setCodec(new TrackingPostingsInMemoryBytesCodec(engineConfig.getCodec()));
         boolean useCompoundFile = engineConfig.getUseCompoundFile();
         iwc.setUseCompoundFile(useCompoundFile);
         if (useCompoundFile == false) {

server/src/main/java/org/elasticsearch/index/shard/ShardFieldStats.java

@@ -17,7 +17,8 @@
  * @param totalFields the total number of fields across the segments
  * @param fieldUsages the number of usages for segment-level fields (e.g., doc_values, postings, norms, points)
  *                    -1 if unavailable
+ * @param postingsInMemoryBytes the total bytes in memory used for postings across all fields
  */
-public record ShardFieldStats(int numSegments, int totalFields, long fieldUsages) {
+public record ShardFieldStats(int numSegments, int totalFields, long fieldUsages, long postingsInMemoryBytes) {
 
 }