configurable rendering of map and weightedset in fields

1 files changed, 159 insertions, 76 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 f9a91558254..a2d9e78b37f 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,7 +76,10 @@ 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 WRAP_DEEP_MAPS = new CompoundName("renderer.json.jsonMaps");
+    private static final CompoundName WRAP_MAPS_ALL = new CompoundName("renderer.json.jsonMapsAll");
+    private static final CompoundName WRAP_WSETS = new CompoundName("renderer.json.jsonWsets");
+    private static final CompoundName WRAP_WSETS_ALL = new CompoundName("renderer.json.jsonWsetsAll");
     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,13 +128,19 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     private JsonGenerator generator;
     private FieldConsumer fieldConsumer;
     private Deque<Integer> renderedChildren;
-    private boolean debugRendering;
-    private boolean jsonMaps;
+    static class FieldConsumerSettings {
+        boolean debugRendering = false;
+        boolean jsonDeepMaps = false;
+        boolean jsonWsets = false;
+        boolean jsonMapsAll = false;
+        boolean jsonWsetsAll = false;
+        boolean tensorShortForm = false;
+        boolean convertDeep() { return (jsonDeepMaps || jsonWsets); }
+    }
+    private final FieldConsumerSettings fieldConsumerSettings = new FieldConsumerSettings();
     private LongSupplier timeSource;
     private OutputStream stream;
 
-    private boolean tensorShortFormRendering = false;
-
     public JsonRenderer() {
         this(null);
     }
@@ -160,9 +169,13 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     @Override
     public void init() {
         super.init();
-        debugRendering = false;
-        jsonMaps = false;
-        setGenerator(null, debugRendering);
+        fieldConsumerSettings.debugRendering = false;
+        fieldConsumerSettings.jsonDeepMaps = false;
+        fieldConsumerSettings.jsonWsets = false;
+        fieldConsumerSettings.jsonMapsAll = false;
+        fieldConsumerSettings.jsonWsetsAll = false;
+        fieldConsumerSettings.tensorShortForm = false;
+        setGenerator(null, fieldConsumerSettings);
         renderedChildren = null;
         timeSource = System::currentTimeMillis;
         stream = null;
@@ -171,10 +184,8 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     @Override
     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);
+        getSettings(getResult().getQuery());
+        setGenerator(generatorFactory.createGenerator(stream, JsonEncoding.UTF8), fieldConsumerSettings);
         renderedChildren = new ArrayDeque<>();
         generator.writeStartObject();
         renderTrace(getExecution().trace());
@@ -204,18 +215,24 @@ 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);
-    }
-
-    private boolean getTensorShortFormRendering(Query q) {
-        if (q == null || q.properties().get(TENSOR_FORMAT) == null)
-            return false;
-        return q.properties().getString(TENSOR_FORMAT).equalsIgnoreCase("short");
+    private void getSettings(Query q) {
+        if (q == null) {
+            fieldConsumerSettings.debugRendering = false;
+            fieldConsumerSettings.jsonDeepMaps = false;
+            fieldConsumerSettings.jsonWsets = false;
+            fieldConsumerSettings.jsonMapsAll = false;
+            fieldConsumerSettings.jsonWsetsAll = false;
+            fieldConsumerSettings.tensorShortForm = false;
+            return;
+        }
+        var props = q.properties();
+        fieldConsumerSettings.debugRendering = props.getBoolean(DEBUG_RENDERING_KEY, false);
+        fieldConsumerSettings.jsonDeepMaps = props.getBoolean(WRAP_DEEP_MAPS, false);
+        fieldConsumerSettings.jsonWsets = props.getBoolean(WRAP_WSETS, false);
+        fieldConsumerSettings.jsonMapsAll = props.getBoolean(WRAP_MAPS_ALL, false);
+        fieldConsumerSettings.jsonWsetsAll = props.getBoolean(WRAP_WSETS_ALL, false);
+        fieldConsumerSettings.tensorShortForm = (props.get(TENSOR_FORMAT) != null &&
+                                                 props.getString(TENSOR_FORMAT).equalsIgnoreCase("short"));
     }
 
     protected void renderTrace(Trace trace) throws IOException {
@@ -520,17 +537,18 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         return null;
     }
 
-    private void setGenerator(JsonGenerator generator, boolean debugRendering) {
+    private void setGenerator(JsonGenerator generator, FieldConsumerSettings settings) {
         this.generator = generator;
-        this.fieldConsumer = generator == null ? null : createFieldConsumer(generator, debugRendering, jsonMaps);
+        this.fieldConsumer = generator == null ? null : createFieldConsumer(generator, settings);
     }
 
     protected FieldConsumer createFieldConsumer(JsonGenerator generator, boolean debugRendering) {
-        return createFieldConsumer(generator, debugRendering, this.jsonMaps);
+        fieldConsumerSettings.debugRendering = debugRendering;
+        return createFieldConsumer(generator, fieldConsumerSettings);
     }
 
-    private FieldConsumer createFieldConsumer(JsonGenerator generator, boolean debugRendering, boolean jsonMaps) {
-        return new FieldConsumer(generator, debugRendering, tensorShortFormRendering, jsonMaps);
+    private FieldConsumer createFieldConsumer(JsonGenerator generator, FieldConsumerSettings settings) {
+        return new FieldConsumer(generator, settings);
     }
 
     /**
@@ -548,10 +566,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     public static class FieldConsumer implements Hit.RawUtf8Consumer, TraceRenderer.FieldConsumer {
 
         private final JsonGenerator generator;
-        private final boolean debugRendering;
-        private final boolean jsonMaps;
-        private final boolean tensorShortForm;
-
+        private final FieldConsumerSettings settings;
         private MutableBoolean hasFieldsField;
 
         public FieldConsumer(JsonGenerator generator, boolean debugRendering) {
@@ -560,11 +575,16 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         public FieldConsumer(JsonGenerator generator, boolean debugRendering, boolean tensorShortForm) {
             this(generator, debugRendering, tensorShortForm, false);
         }
-        public FieldConsumer(JsonGenerator generator, boolean debugRendering, boolean tensorShortForm, boolean jsonMaps) {
+        public FieldConsumer(JsonGenerator generator, boolean debugRendering, boolean tensorShortForm, boolean jsonDeepMaps) {
             this.generator = generator;
-            this.debugRendering = debugRendering;
-            this.tensorShortForm = tensorShortForm;
-            this.jsonMaps = jsonMaps;
+            this.settings = new FieldConsumerSettings();
+            this.settings.debugRendering = debugRendering;
+            this.settings.tensorShortForm = tensorShortForm;
+            this.settings.jsonDeepMaps = jsonDeepMaps;
+        }
+        FieldConsumer(JsonGenerator generator, FieldConsumerSettings settings) {
+            this.generator = generator;
+            this.settings = settings;
         }
 
         /**
@@ -618,7 +638,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         }
 
         protected boolean shouldRender(String name, Object value) {
-            if (debugRendering) return true;
+            if (settings.debugRendering) return true;
             if (name.startsWith(VESPA_HIDDEN_FIELD_PREFIX)) return false;
             if (value instanceof CharSequence && ((CharSequence) value).length() == 0) return false;
             // StringFieldValue cannot hold a null, so checking length directly is OK:
@@ -628,47 +648,68 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         }
 
         protected boolean shouldRenderUtf8Value(String name, int length) {
-            if (debugRendering) return true;
+            if (settings.debugRendering) return true;
             if (name.startsWith(VESPA_HIDDEN_FIELD_PREFIX)) return false;
             if (length == 0) return false;
             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;
-                    }
+        private Inspector maybeConvertMap(Inspector data) {
+            var map = new Value.ObjectValue();
+            for (int i = 0; i < data.entryCount(); i++) {
+                Inspector obj = data.entry(i);
+                if (obj.type() != Type.OBJECT || obj.fieldCount() != 2) {
+                    return null;
                 }
-                if (map != null) {
-                    return map;
+                Inspector key = obj.field("key");
+                Inspector value = obj.field("value");
+                if (! key.valid()) return null;
+                if (! value.valid()) return null;
+                if (key.type() != Type.STRING && !settings.jsonMapsAll) {
+                    return null;
+                }
+                if (settings.convertDeep()) {
+                    value = deepMaybeConvert(value);
                 }
-                var array = new Value.ArrayValue();
-                for (int i = 0; i < data.entryCount(); i++) {
-                    Inspector obj = data.entry(i);
-                    array.add(deepWrapAsMap(obj));
+                if (key.type() == Type.STRING) {
+                    map.put(key.asString(), value);
+                } else {
+                    map.put(key.toString(), value);
                 }
-                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 map;
+        }
+
+        private Inspector maybeConvertWset(Inspector data) {
+            var wset = new Value.ObjectValue();
+            for (int i = 0; i < data.entryCount(); i++) {
+                Inspector obj = data.entry(i);
+                if (obj.type() != Type.OBJECT || obj.fieldCount() != 2) {
+                    return null;
+                }
+                Inspector item = obj.field("item");
+                Inspector weight = obj.field("weight");
+                if (! item.valid()) return null;
+                if (! weight.valid()) return null;
+                // TODO support non-integer weights?
+                if (weight.type() != Type.LONG) return null;
+                if (item.type() == Type.STRING) {
+                    wset.put(item.asString(), weight.asLong());
+                } else if (settings.jsonWsetsAll) {
+                    wset.put(item.toString(), weight.asLong());
+                } else {
+                    return null;
                 }
-                return object;
             }
-            return data;
+            return wset;
+        }
+
+        private Inspector convertInsideObject(Inspector data) {
+            var object = new Value.ObjectValue();
+            for (var entry : data.fields()) {
+                object.put(entry.getKey(), deepMaybeConvert(entry.getValue()));
+            }
+            return object;
         }
 
         private static Inspector wrapAsMap(Inspector data) {
@@ -688,13 +729,55 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
             return map;
         }
 
-        private void renderInspector(Inspector data) throws IOException {
-            Inspector asMap = jsonMaps ? deepWrapAsMap(data) : wrapAsMap(data);
-            if (asMap != null) {
-                renderInspectorDirect(asMap);
-            } else {
-                renderInspectorDirect(data);
+        private Inspector deepMaybeConvert(Inspector data) {
+            if (data.type() == Type.ARRAY) {
+                if (settings.jsonDeepMaps) {
+                    var map = maybeConvertMap(data);
+                    if (map != null) return map;
+                }
+                if (settings.jsonWsets) {
+                    var wset = maybeConvertWset(data);
+                    if (wset != null) return wset;
+                }
+            }
+            if (data.type() == Type.OBJECT) {
+                return convertInsideObject(data);
+            }
+            return data;
+        }
+
+        private Inspector convertTopLevelArray(Inspector data) {
+            if (data.entryCount() > 0) {
+                var map = maybeConvertMap(data);
+                if (map != null) return map;
+                if (settings.jsonWsets) {
+                    var wset = maybeConvertWset(data);
+                    if (wset != null) return wset;
+                }
+                if (settings.convertDeep()) {
+                    var array = new Value.ArrayValue();
+                    for (int i = 0; i < data.entryCount(); i++) {
+                        Inspector obj = data.entry(i);
+                        array.add(deepMaybeConvert(obj));
+                    }
+                    return array;
+                }
             }
+            return data;
+        }
+
+        private Inspector maybeConvertData(Inspector data) throws IOException {
+            if (data.type() == Type.ARRAY) {
+                return convertTopLevelArray(data);
+            }
+            if (settings.convertDeep() && data.type() == Type.OBJECT) {
+                return convertInsideObject(data);
+            }
+            return data;
+        }
+
+        private void renderInspector(Inspector data) throws IOException {
+            renderInspectorDirect(maybeConvertData(data));
         }
 
         private void renderInspectorDirect(Inspector data) throws IOException {
@@ -724,7 +807,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
             } else if (field instanceof Tensor) {
                 renderTensor(Optional.of((Tensor)field));
             } else if (field instanceof FeatureData) {
-                generator.writeRawValue(((FeatureData)field).toJson(tensorShortForm));
+                generator.writeRawValue(((FeatureData)field).toJson(settings.tensorShortForm));
             } else if (field instanceof Inspectable) {
                 renderInspectorDirect(((Inspectable)field).inspect());
             } else if (field instanceof JsonProducer) {
@@ -769,7 +852,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
                 generator.writeEndObject();
                 return;
             }
-            if (tensorShortForm) {
+            if (settings.tensorShortForm) {
                 generator.writeRawValue(new String(JsonFormat.encodeShortForm(tensor.get()), StandardCharsets.UTF_8));
             } else {
                 generator.writeRawValue(new String(JsonFormat.encode(tensor.get()), StandardCharsets.UTF_8));