Add support for upgrades of values in maps.

Since keys in maps and weighted sets must be primitives we should be complete here.

2 files changed, 19 insertions, 6 deletions

diff --git a/docproc/src/main/java/com/yahoo/docproc/AbstractConcreteDocumentFactory.java b/docproc/src/main/java/com/yahoo/docproc/AbstractConcreteDocumentFactory.java
index 522af9e1f84..7e059da6097 100644
--- a/docproc/src/main/java/com/yahoo/docproc/AbstractConcreteDocumentFactory.java
+++ b/docproc/src/main/java/com/yahoo/docproc/AbstractConcreteDocumentFactory.java
@@ -10,6 +10,7 @@ import com.yahoo.document.Field;
 import com.yahoo.document.annotation.Annotation;
 import com.yahoo.document.datatypes.Array;
 import com.yahoo.document.datatypes.FieldValue;
+import com.yahoo.document.datatypes.MapFieldValue;
 import com.yahoo.document.datatypes.Struct;
 import com.yahoo.document.datatypes.StructuredFieldValue;
 
@@ -56,11 +57,16 @@ public abstract class AbstractConcreteDocumentFactory extends com.yahoo.componen
         } else if (fv instanceof Array) {
             Array<FieldValue> array = (Array<FieldValue>) fv;
             DataType nestedType = array.getDataType().getNestedType();
-            for (int i=0; i < array.size(); i++) {
-                array.set(i, optionallyUpgrade(nestedType, array.get(i)));
+            if (nestedType.getPrimitiveType() == null) {
+                array.replaceAll((item) -> optionallyUpgrade(nestedType, item));
+            }
+        } else if (fv instanceof MapFieldValue) {
+            MapFieldValue<FieldValue, FieldValue> map = (MapFieldValue<FieldValue, FieldValue>) fv;
+            DataType valueTypeType = map.getDataType().getValueType();
+            if (valueTypeType.getPrimitiveType() == null) {
+                map.replaceAll((key, value) -> optionallyUpgrade(valueTypeType, value));
             }
         }
-        // TODO We also need specialhandling for weighted set/map. Limiting to array until verified.
         return fv;
     }
 }
diff --git a/documentgen-test/src/test/java/com/yahoo/vespa/config/DocumentGenPluginTest.java b/documentgen-test/src/test/java/com/yahoo/vespa/config/DocumentGenPluginTest.java
index b3aca93cd96..8c20355ce0c 100644
--- a/documentgen-test/src/test/java/com/yahoo/vespa/config/DocumentGenPluginTest.java
+++ b/documentgen-test/src/test/java/com/yahoo/vespa/config/DocumentGenPluginTest.java
@@ -316,9 +316,7 @@ public class DocumentGenPluginTest {
         verifyArrayOfStruct(toBook(copyBySerialization(book)));
     }
 
-    @Test
-    public void testMaps() {
-        Book book = getBook();
+    private void verifyMaps(Book book) {
         assertTrue(book.getField("stringmap").getDataType() instanceof MapDataType);
         MapFieldValue mfv = (MapFieldValue) book.getFieldValue("stringmap");
         assertEquals(mfv.get(new StringFieldValue("Melville")), new StringFieldValue("Moby Dick"));
@@ -328,6 +326,8 @@ public class DocumentGenPluginTest {
         assertEquals(mfv.keySet().size(), 2);
         book.getStringmap().put("Melville", "MD");
         assertEquals(mfv.keySet().size(), 3);
+        book.getStringmap().put("Melville", "Moby Dick");
+        assertEquals(mfv.keySet().size(), 3);
 
         assertEquals(book.getStructmap().get(50).getS1(), "test s1");
         MapFieldValue mfv2 = (MapFieldValue) book.getFieldValue("structmap");
@@ -337,6 +337,13 @@ public class DocumentGenPluginTest {
     }
 
     @Test
+    public void testMaps() {
+        Book book = getBook();
+        verifyMaps(book);
+        verifyMaps(toBook(copyBySerialization(book)));
+    }
+
+    @Test
     public void testWeightedSets() {
         Book book = getBook();
         assertTrue(book.getField("mywsfloat").getDataType() instanceof WeightedSetDataType);