Merge pull request #19420 from vespa-engine/arnej/wrap-all-maps-in-json-renderer

Arnej/wrap all maps in json renderer

3 files changed, 107 insertions, 7 deletions

diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index e4c1a561dee..4913161d955 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -7076,6 +7076,7 @@
     "methods": [
       "public void <init>(com.fasterxml.jackson.core.JsonGenerator, boolean)",
       "public void <init>(com.fasterxml.jackson.core.JsonGenerator, boolean, boolean)",
+      "public void <init>(com.fasterxml.jackson.core.JsonGenerator, boolean, boolean, boolean)",
       "public void accept(java.lang.String, java.lang.Object)",
       "public void accept(java.lang.String, byte[], int, int)",
       "protected boolean shouldRender(java.lang.String, java.lang.Object)",
diff --git a/container-search/src/main/java/com/yahoo/search/rendering/JsonRenderer.java b/container-search/src/main/java/com/yahoo/search/rendering/JsonRenderer.java
index ee0e7f4fe0e..92cdd10616b 100644
--- a/container-search/src/main/java/com/yahoo/search/rendering/JsonRenderer.java
+++ b/container-search/src/main/java/com/yahoo/search/rendering/JsonRenderer.java
@@ -76,6 +76,7 @@ import static com.fasterxml.jackson.databind.SerializationFeature.FLUSH_AFTER_WR
 // NOTE: The JSON format is a public API. If new elements are added be sure to update the reference doc.
 public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
 
+    private static final CompoundName WRAP_ALL_MAPS = new CompoundName("renderer.json.jsonMaps");
     private static final CompoundName DEBUG_RENDERING_KEY = new CompoundName("renderer.json.debug");
     private static final CompoundName JSON_CALLBACK = new CompoundName("jsoncallback");
     private static final CompoundName TENSOR_FORMAT = new CompoundName("format.tensors");
@@ -125,6 +126,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     private FieldConsumer fieldConsumer;
     private Deque<Integer> renderedChildren;
     private boolean debugRendering;
+    private boolean jsonMaps;
     private LongSupplier timeSource;
     private OutputStream stream;
 
@@ -159,6 +161,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     public void init() {
         super.init();
         debugRendering = false;
+        jsonMaps = false;
         setGenerator(null, debugRendering);
         renderedChildren = null;
         timeSource = System::currentTimeMillis;
@@ -169,6 +172,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     public void beginResponse(OutputStream stream) throws IOException {
         beginJsonCallback(stream);
         debugRendering = getDebugRendering(getResult().getQuery());
+        jsonMaps = getWrapAllMaps(getResult().getQuery());
         tensorShortFormRendering = getTensorShortFormRendering(getResult().getQuery());
         setGenerator(generatorFactory.createGenerator(stream, JsonEncoding.UTF8), debugRendering);
         renderedChildren = new ArrayDeque<>();
@@ -200,6 +204,10 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         generator.writeEndObject();
     }
 
+    private boolean getWrapAllMaps(Query q) {
+        return q != null && q.properties().getBoolean(WRAP_ALL_MAPS, false);
+    }
+
     private boolean getDebugRendering(Query q) {
         return q != null && q.properties().getBoolean(DEBUG_RENDERING_KEY, false);
     }
@@ -514,11 +522,15 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
 
     private void setGenerator(JsonGenerator generator, boolean debugRendering) {
         this.generator = generator;
-        this.fieldConsumer = generator == null ? null : createFieldConsumer(generator, debugRendering);
+        this.fieldConsumer = generator == null ? null : createFieldConsumer(generator, debugRendering, jsonMaps);
     }
 
     protected FieldConsumer createFieldConsumer(JsonGenerator generator, boolean debugRendering) {
-        return new FieldConsumer(generator, debugRendering, tensorShortFormRendering);
+        return createFieldConsumer(generator, debugRendering, this.jsonMaps);
+    }
+
+    private FieldConsumer createFieldConsumer(JsonGenerator generator, boolean debugRendering, boolean jsonMaps) {
+        return new FieldConsumer(generator, debugRendering, tensorShortFormRendering, jsonMaps);
     }
 
     /**
@@ -537,6 +549,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
 
         private final JsonGenerator generator;
         private final boolean debugRendering;
+        private final boolean jsonMaps;
         private final boolean tensorShortForm;
 
         private MutableBoolean hasFieldsField;
@@ -544,11 +557,14 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         public FieldConsumer(JsonGenerator generator, boolean debugRendering) {
             this(generator, debugRendering, false);
         }
-
         public FieldConsumer(JsonGenerator generator, boolean debugRendering, boolean tensorShortForm) {
+            this(generator, debugRendering, tensorShortForm, false);
+        }
+        public FieldConsumer(JsonGenerator generator, boolean debugRendering, boolean tensorShortForm, boolean jsonMaps) {
             this.generator = generator;
             this.debugRendering = debugRendering;
             this.tensorShortForm = tensorShortForm;
+            this.jsonMaps = jsonMaps;
         }
 
         /**
@@ -618,6 +634,43 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
             return true;
         }
 
+        private static Inspector deepWrapAsMap(Inspector data) {
+            if (data.type() == Type.ARRAY) {
+                var map = new Value.ObjectValue();
+                for (int i = 0; i < data.entryCount(); i++) {
+                    Inspector obj = data.entry(i);
+                    if (map != null && obj.type() == Type.OBJECT && obj.fieldCount() == 2) {
+                        Inspector key = obj.field("key");
+                        Inspector value = obj.field("value");
+                        if (key.type() == Type.STRING && value.valid()) {
+                            map.put(key.asString(), deepWrapAsMap(value));
+                        } else {
+                            map = null;
+                        }
+                    } else {
+                        map = null;
+                    }
+                }
+                if (map != null) {
+                    return map;
+                }
+                var array = new Value.ArrayValue();
+                for (int i = 0; i < data.entryCount(); i++) {
+                    Inspector obj = data.entry(i);
+                    array.add(deepWrapAsMap(obj));
+                }
+                return array;
+            }
+            if (data.type() == Type.OBJECT) {
+                var object = new Value.ObjectValue();
+                for (var entry : data.fields()) {
+                    object.put(entry.getKey(), deepWrapAsMap(entry.getValue()));
+                }
+                return object;
+            }
+            return data;
+        }
+
         private static Inspector wrapAsMap(Inspector data) {
             if (data.type() != Type.ARRAY) return null;
             if (data.entryCount() == 0) return null;
@@ -636,11 +689,9 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         }
 
         private void renderInspector(Inspector data) throws IOException {
-            Inspector asMap = wrapAsMap(data);
+            Inspector asMap = jsonMaps ? deepWrapAsMap(data) : wrapAsMap(data);
             if (asMap != null) {
-                StringBuilder intermediate = new StringBuilder();
-                JsonRender.render(asMap, intermediate, true);
-                generator.writeRawValue(intermediate.toString());
+                renderInspectorDirect(asMap);
             } else {
                 renderInspectorDirect(data);
             }
diff --git a/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java b/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java
index 290b7266a3a..2a399ac0e7b 100644
--- a/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java
@@ -1259,6 +1259,54 @@ public class JsonRendererTestCase {
         assertEqualJson(expected, summary);
     }
 
+    private static SlimeAdapter dataFromSimplified(String simplified) {
+        var decoder = new com.yahoo.slime.JsonDecoder();
+        var slime = decoder.decode(new Slime(), Utf8.toBytes(simplified));
+        return new SlimeAdapter(slime.get());
+    }
+
+    @Test
+    public void testMapDeepInFields() throws IOException, InterruptedException, ExecutionException {
+        Result r = new Result(new Query("/?renderer.json.jsonMaps=true"));
+        var expected = dataFromSimplified(
+                "{root: { id:'toplevel', relevance:1.0, fields: { totalCount: 1 }," +
+                "  children: [ { id: 'myHitName', relevance: 1.0," +
+                "    fields: { " +
+                "      f1: [ 'v1', { mykey1: 'myvalue1', mykey2: 'myvalue2' } ]," +
+                "      f2: { i1: 'v2', i2: { mykey3: 'myvalue3' }, i3: 'v3' }," +
+                "      f3: { j1: 42, j2: 17.75, j3: [ 'v4', 'v5' ] }," +
+                "      f4: { mykey4: 'myvalue4', mykey5: 'myvalue5' }" +
+                "    }" +
+                "  } ]" +
+                "}}");
+        Hit h = new Hit("myHitName");
+        h.setField("f1", dataFromSimplified("[ 'v1', [ { key: 'mykey1', value: 'myvalue1' }, { key: 'mykey2', value: 'myvalue2' } ] ]"));
+        h.setField("f2", dataFromSimplified("{ i1: 'v2', i2: [ { key: 'mykey3', value: 'myvalue3' } ], i3: 'v3' }"));
+        h.setField("f3", dataFromSimplified("{ j1: 42, j2: 17.75, j3: [ 'v4', 'v5' ] }"));
+        h.setField("f4", dataFromSimplified("[ { key: 'mykey4', value: 'myvalue4' }, { key: 'mykey5', value: 'myvalue5' } ]"));
+        r.hits().add(h);
+        r.setTotalHitCount(1L);
+        String summary = render(r);
+        assertEqualJson(expected.toString(), summary);
+
+        r = new Result(new Query("/?renderer.json.jsonMaps=false"));
+        expected = dataFromSimplified(
+                "{root:{id:'toplevel',relevance:1.0,fields:{totalCount:1}," +
+                "  children: [ { id: 'myHitName', relevance: 1.0," +
+                "    fields: { " +
+                "      f1: [ 'v1', [ { key: 'mykey1', value: 'myvalue1' }, { key: 'mykey2', value: 'myvalue2' } ] ]," +
+                "      f2: { i1: 'v2', i2: [ { key: 'mykey3', value: 'myvalue3' } ], i3: 'v3' }," +
+                "      f3: { j1: 42, j2: 17.75, j3: [ 'v4', 'v5' ] }," +
+                "      f4: { mykey4: 'myvalue4', mykey5: 'myvalue5' }" +
+                "    }" +
+                "  } ]" +
+                "}}");
+        r.hits().add(h);
+        r.setTotalHitCount(1L);
+        summary = render(r);
+        assertEqualJson(expected.toString(), summary);
+    }
+
     @Test
     public void testThatTheJsonValidatorCanCatchErrors() {
         String json = "{"