Replace hashmap by a counter in TrackingPostingsInMemoryBytesCodec

Author: martijnvg

server/src/main/java/org/elasticsearch/index/codec/TrackingPostingsInMemoryBytesCodec.java

@@ -15,14 +15,14 @@
 import org.apache.lucene.codecs.FilterCodec;
 import org.apache.lucene.codecs.NormsProducer;
 import org.apache.lucene.codecs.PostingsFormat;
+import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.Fields;
 import org.apache.lucene.index.FilterLeafReader;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
 import org.apache.lucene.index.Terms;
 import org.apache.lucene.index.TermsEnum;
-import org.apache.lucene.internal.hppc.IntIntHashMap;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.util.FeatureFlag;
 
@@ -63,22 +63,18 @@ public FieldsProducer fieldsProducer(SegmentReadState state) throws IOException
     static final class TrackingLengthFieldsConsumer extends FieldsConsumer {
         final SegmentWriteState state;
         final FieldsConsumer in;
-        final IntIntHashMap termsBytesPerField;
+        final long[] totalBytes;
 
         TrackingLengthFieldsConsumer(SegmentWriteState state, FieldsConsumer in) {
             this.state = state;
             this.in = in;
-            this.termsBytesPerField = new IntIntHashMap(state.fieldInfos.size());
+            this.totalBytes = new long[1];
         }
 
         @Override
         public void write(Fields fields, NormsProducer norms) throws IOException {
-            in.write(new TrackingLengthFields(fields, termsBytesPerField, state.fieldInfos), norms);
-            long totalBytes = 0;
-            for (int bytes : termsBytesPerField.values) {
-                totalBytes += bytes;
-            }
-            state.segmentInfo.putAttribute(IN_MEMORY_POSTINGS_BYTES_KEY, Long.toString(totalBytes));
+            in.write(new TrackingLengthFields(fields, totalBytes, state.fieldInfos), norms);
+            state.segmentInfo.putAttribute(IN_MEMORY_POSTINGS_BYTES_KEY, Long.toString(totalBytes[0]));
         }
 
         @Override
@@ -88,12 +84,18 @@ public void close() throws IOException {
     }
 
     static final class TrackingLengthFields extends FilterLeafReader.FilterFields {
-        final IntIntHashMap termsBytesPerField;
+        final long[] totalBytes;
+        final boolean[] seenFields;
         final FieldInfos fieldInfos;
 
-        TrackingLengthFields(Fields in, IntIntHashMap termsBytesPerField, FieldInfos fieldInfos) {
+        TrackingLengthFields(Fields in, long[] totalBytes, FieldInfos fieldInfos) {
             super(in);
-            this.termsBytesPerField = termsBytesPerField;
+            this.totalBytes = totalBytes;
+            int maxFieldNumber = 0;
+            for (FieldInfo info : fieldInfos) {
+                maxFieldNumber = Math.max(info.getFieldNumber(), maxFieldNumber);
+            }
+            this.seenFields = new boolean[maxFieldNumber + 1];
             this.fieldInfos = fieldInfos;
         }
 
@@ -104,10 +106,12 @@ public Terms terms(String field) throws IOException {
                 return null;
             }
             int fieldNum = fieldInfos.fieldInfo(field).number;
-            return new TrackingLengthTerms(
-                terms,
-                bytes -> termsBytesPerField.put(fieldNum, Math.max(termsBytesPerField.getOrDefault(fieldNum, 0), bytes))
-            );
+            return new TrackingLengthTerms(terms, bytes -> {
+                if (seenFields[fieldNum] == false) {
+                    totalBytes[0] += bytes;
+                    seenFields[fieldNum] = true;
+                }
+            });
         }
     }