1 files changed, 58 insertions, 7 deletions

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);
             }