Fix counted_keyword fields for arrays of objects

Author: jordan-powers

x-pack/plugin/mapper-counted-keyword/src/main/java/org/elasticsearch/xpack/countedkeyword/CountedKeywordFieldMapper.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.countedkeyword;
 
-import org.apache.lucene.document.BinaryDocValuesField;
 import org.apache.lucene.document.FieldType;
 import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.index.DocValues;
@@ -19,6 +18,7 @@
 import org.apache.lucene.index.TermsEnum;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.io.stream.ByteArrayStreamInput;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
@@ -31,6 +31,7 @@
 import org.elasticsearch.index.fielddata.plain.AbstractIndexOrdinalsFieldData;
 import org.elasticsearch.index.fielddata.plain.AbstractLeafOrdinalsFieldData;
 import org.elasticsearch.index.mapper.BinaryFieldMapper;
+import org.elasticsearch.index.mapper.CustomDocValuesField;
 import org.elasticsearch.index.mapper.DocumentParserContext;
 import org.elasticsearch.index.mapper.FieldMapper;
 import org.elasticsearch.index.mapper.KeywordFieldMapper;
@@ -434,15 +435,17 @@ protected void parseCreateField(DocumentParserContext context) throws IOExceptio
             return;
         }
 
-        int i = 0;
-        int[] counts = new int[values.size()];
-        for (Map.Entry<String, Integer> value : values.entrySet()) {
-            context.doc().add(new KeywordFieldMapper.KeywordField(fullPath(), new BytesRef(value.getKey()), fieldType));
-            counts[i++] = value.getValue();
+        for (String value : values.keySet()) {
+            context.doc().add(new KeywordFieldMapper.KeywordField(fullPath(), new BytesRef(value), fieldType));
+        }
+        CountsBinaryDocValuesField field = (CountsBinaryDocValuesField) context.doc().getByKey(countFieldMapper.fieldType().name());
+        if (field == null) {
+            field = new CountsBinaryDocValuesField(countFieldMapper.fieldType().name());
+            field.add(values);
+            context.doc().addWithKey(countFieldMapper.fieldType().name(), field);
+        } else {
+            field.add(values);
         }
-        BytesStreamOutput streamOutput = new BytesStreamOutput();
-        streamOutput.writeVIntArray(counts);
-        context.doc().add(new BinaryDocValuesField(countFieldMapper.fullPath(), streamOutput.bytes().toBytesRef()));
     }
 
     private void parseArray(DocumentParserContext context, SortedMap<String, Integer> values) throws IOException {
@@ -509,4 +512,36 @@ protected SyntheticSourceSupport syntheticSourceSupport() {
         );
     }
 
+    private class CountsBinaryDocValuesField extends CustomDocValuesField {
+        private final SortedMap<String, Integer> counts;
+
+        CountsBinaryDocValuesField(String name) {
+            super(name);
+            counts = new TreeMap<>();
+        }
+
+        public void add(SortedMap<String, Integer> newCounts) {
+            for (Map.Entry<String, Integer> currCount : newCounts.entrySet()) {
+                this.counts.put(currCount.getKey(), this.counts.getOrDefault(currCount.getKey(), 0) + currCount.getValue());
+            }
+        }
+
+        @Override
+        public BytesRef binaryValue() {
+            try {
+                int bytesSize = (counts.size() + 1) * 5;
+                BytesStreamOutput out = new BytesStreamOutput(bytesSize);
+                int countsArr[] = new int[counts.size()];
+                int i = 0;
+                for (Integer currCount : counts.values()) {
+                    countsArr[i++] = currCount;
+                }
+                out.writeVIntArray(countsArr);
+                return out.bytes().toBytesRef();
+            } catch (IOException e) {
+                throw new ElasticsearchException("Failed to get binary value", e);
+            }
+        }
+    }
+
 }