Cleanup random stats serialization code (#22223)

Some of our stats serialization code duplicates complicated seriazliation logic
or could use existing building blocks from StreamOutput/Input. This commit
cleans up some of the serialization code.

2 files changed, 18 insertions, 68 deletions

diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionFieldStats.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionFieldStats.java
index e5e1b1b919..8b5761a7e9 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionFieldStats.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionFieldStats.java
@@ -27,6 +27,7 @@ import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Terms;
 import org.apache.lucene.search.suggest.document.CompletionTerms;
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.common.FieldMemoryStats;
 import org.elasticsearch.common.regex.Regex;
 
 import java.io.IOException;
@@ -64,6 +65,6 @@ public class CompletionFieldStats {
                 throw new ElasticsearchException(ioe);
             }
         }
-        return new CompletionStats(sizeInBytes, completionFields);
+        return new CompletionStats(sizeInBytes, completionFields == null ? null : new FieldMemoryStats(completionFields));
     }
 }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionStats.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionStats.java
index efea591576..c123d46fe4 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionStats.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionStats.java
@@ -18,7 +18,7 @@
  */
 package org.elasticsearch.search.suggest.completion;
 
-import com.carrotsearch.hppc.ObjectLongHashMap;
+import org.elasticsearch.common.FieldMemoryStats;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -31,15 +31,19 @@ import java.io.IOException;
 
 public class CompletionStats implements Streamable, ToXContent {
 
-    private long sizeInBytes;
+    private static final String COMPLETION = "completion";
+    private static final String SIZE_IN_BYTES = "size_in_bytes";
+    private static final String SIZE = "size";
+    private static final String FIELDS = "fields";
 
+    private long sizeInBytes;
     @Nullable
-    private ObjectLongHashMap<String> fields;
+    private FieldMemoryStats fields;
 
     public CompletionStats() {
     }
 
-    public CompletionStats(long size, @Nullable ObjectLongHashMap<String> fields) {
+    public CompletionStats(long size, @Nullable FieldMemoryStats fields) {
         this.sizeInBytes = size;
         this.fields = fields;
     }
@@ -52,98 +56,43 @@ public class CompletionStats implements Streamable, ToXContent {
         return new ByteSizeValue(sizeInBytes);
     }
 
-    public ObjectLongHashMap<String> getFields() {
+    public FieldMemoryStats getFields() {
         return fields;
     }
 
     @Override
     public void readFrom(StreamInput in) throws IOException {
         sizeInBytes = in.readVLong();
-        if (in.readBoolean()) {
-            int size = in.readVInt();
-            fields = new ObjectLongHashMap<>(size);
-            for (int i = 0; i < size; i++) {
-                fields.put(in.readString(), in.readVLong());
-            }
-        }
+        fields = in.readOptionalWriteable(FieldMemoryStats::new);
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeVLong(sizeInBytes);
-        if (fields == null) {
-            out.writeBoolean(false);
-        } else {
-            out.writeBoolean(true);
-            out.writeVInt(fields.size());
-
-            assert !fields.containsKey(null);
-            final Object[] keys = fields.keys;
-            final long[] values = fields.values;
-            for (int i = 0; i < keys.length; i++) {
-                if (keys[i] != null) {
-                    out.writeString((String) keys[i]);
-                    out.writeVLong(values[i]);
-                }
-            }
-        }
+        out.writeOptionalWriteable(fields);
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.startObject(Fields.COMPLETION);
-        builder.byteSizeField(Fields.SIZE_IN_BYTES, Fields.SIZE, sizeInBytes);
+        builder.startObject(COMPLETION);
+        builder.byteSizeField(SIZE_IN_BYTES, SIZE, sizeInBytes);
         if (fields != null) {
-            builder.startObject(Fields.FIELDS);
-
-            assert !fields.containsKey(null);
-            final Object[] keys = fields.keys;
-            final long[] values = fields.values;
-            for (int i = 0; i < keys.length; i++) {
-                if (keys[i] != null) {
-                    builder.startObject((String) keys[i]);
-                    builder.byteSizeField(Fields.SIZE_IN_BYTES, Fields.SIZE, values[i]);
-                    builder.endObject();
-                }
-            }
-            builder.endObject();
+            fields.toXContent(builder, FIELDS, SIZE_IN_BYTES, SIZE);
         }
         builder.endObject();
         return builder;
     }
 
-    public static CompletionStats readCompletionStats(StreamInput in) throws IOException {
-        CompletionStats stats = new CompletionStats();
-        stats.readFrom(in);
-        return stats;
-    }
-
-    static final class Fields {
-        static final String COMPLETION = "completion";
-        static final String SIZE_IN_BYTES = "size_in_bytes";
-        static final String SIZE = "size";
-        static final String FIELDS = "fields";
-    }
-
     public void add(CompletionStats completion) {
         if (completion == null) {
             return;
         }
-
         sizeInBytes += completion.getSizeInBytes();
-
         if (completion.fields != null) {
             if (fields == null) {
-                fields = completion.fields.clone();
+                fields = completion.fields.copy();
             } else {
-                assert !completion.fields.containsKey(null);
-                final Object[] keys = completion.fields.keys;
-                final long[] values = completion.fields.values;
-                for (int i = 0; i < keys.length; i++) {
-                    if (keys[i] != null) {
-                        fields.addTo((String) keys[i], values[i]);
-                    }
-                }
+                fields.add(completion.fields);
             }
         }
     }