diff options
4 files changed, 84 insertions, 8 deletions
diff --git a/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java b/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java index 340bd542885..a9d34fc78d8 100644 --- a/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java +++ b/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java @@ -32,10 +32,12 @@ import com.yahoo.document.serialization.FieldWriter; import com.yahoo.tensor.Tensor; import com.yahoo.tensor.TensorAddress; import com.yahoo.tensor.TensorType; +import com.yahoo.tensor.serialization.JsonFormat; import com.yahoo.vespa.objects.FieldBase; import java.io.IOException; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.util.Base64; import java.util.Iterator; import java.util.Map; @@ -71,6 +73,19 @@ public class JsonSerializationHelper { } } + public static void serializeTensorFieldShortForm(JsonGenerator generator, FieldBase field, TensorFieldValue value) { + wrapIOException(() -> { + fieldNameIfNotNull(generator, field); + if (value.getTensor().isPresent()) { + Tensor tensor = value.getTensor().get(); + generator.writeRawValue(new String(JsonFormat.encodeShortForm(tensor), StandardCharsets.UTF_8)); + } else { + generator.writeStartObject(); + generator.writeEndObject(); + } + }); + } + public static void serializeTensorField(JsonGenerator generator, FieldBase field, TensorFieldValue value) { wrapIOException(() -> { fieldNameIfNotNull(generator, field); diff --git a/document/src/main/java/com/yahoo/document/json/JsonWriter.java b/document/src/main/java/com/yahoo/document/json/JsonWriter.java index 95faf753877..05147d2c0ab 100644 --- a/document/src/main/java/com/yahoo/document/json/JsonWriter.java +++ b/document/src/main/java/com/yahoo/document/json/JsonWriter.java @@ -63,6 +63,7 @@ import static com.yahoo.document.json.JsonSerializationHelper.serializeStringFie import static com.yahoo.document.json.JsonSerializationHelper.serializeStructField; import static com.yahoo.document.json.JsonSerializationHelper.serializeStructuredField; import static com.yahoo.document.json.JsonSerializationHelper.serializeTensorField; +import static com.yahoo.document.json.JsonSerializationHelper.serializeTensorFieldShortForm; import static com.yahoo.document.json.JsonSerializationHelper.serializeWeightedSet; import static com.yahoo.document.json.document.DocumentParser.FIELDS; import static com.yahoo.document.json.document.DocumentParser.REMOVE; @@ -77,6 +78,8 @@ public class JsonWriter implements DocumentWriter { private static final JsonFactory jsonFactory = new JsonFactory(); private final JsonGenerator generator; + private final boolean tensorShortForm; + // I really hate exception unsafe constructors, but the alternative // requires generator to not be a final /** @@ -108,7 +111,12 @@ public class JsonWriter implements DocumentWriter { * the output JSON generator */ public JsonWriter(JsonGenerator generator) { + this(generator, false); + } + + public JsonWriter(JsonGenerator generator, boolean tensorShortForm) { this.generator = generator; + this.tensorShortForm = tensorShortForm; } /** @@ -206,7 +214,11 @@ public class JsonWriter implements DocumentWriter { @Override public void write(FieldBase field, TensorFieldValue value) { - serializeTensorField(generator, field, value); + if (tensorShortForm) { + serializeTensorFieldShortForm(generator, field, value); + } else { + serializeTensorField(generator, field, value); + } } @Override diff --git a/document/src/test/java/com/yahoo/document/json/JsonWriterTestCase.java b/document/src/test/java/com/yahoo/document/json/JsonWriterTestCase.java index 7573aba519f..b7f24368ad7 100644 --- a/document/src/test/java/com/yahoo/document/json/JsonWriterTestCase.java +++ b/document/src/test/java/com/yahoo/document/json/JsonWriterTestCase.java @@ -2,6 +2,7 @@ package com.yahoo.document.json; import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.core.JsonGenerator; import com.fasterxml.jackson.core.io.JsonStringEncoder; import com.fasterxml.jackson.databind.ObjectMapper; import com.yahoo.document.ArrayDataType; @@ -25,6 +26,7 @@ import com.yahoo.document.datatypes.TensorFieldValue; import com.yahoo.document.internal.GeoPosType; import com.yahoo.document.json.readers.DocumentParseInfo; import com.yahoo.document.json.readers.VespaJsonDocumentReader; +import com.yahoo.tensor.Tensor; import com.yahoo.tensor.TensorType; import com.yahoo.text.Utf8; import org.junit.After; @@ -32,6 +34,7 @@ import org.junit.Before; import org.junit.Test; import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.util.HashMap; @@ -73,6 +76,7 @@ public class JsonWriterTestCase { registerSinglePositionDocumentType(); registerMultiPositionDocumentType(); registerTensorDocumentType(); + registerIndexedTensorDocumentType(); registerReferenceDocumentType(); } @@ -90,6 +94,13 @@ public class JsonWriterTestCase { types.registerDocumentType(x); } + private void registerIndexedTensorDocumentType() { + DocumentType x = new DocumentType("testindexedtensor"); + TensorType tensorType = new TensorType.Builder().indexed("x", 3).build(); + x.addField(new Field("tensorfield", new TensorDataType(tensorType))); + types.registerDocumentType(x); + } + private void registerMultiPositionDocumentType() { DocumentType x = new DocumentType("testmultipos"); DataType d = new ArrayDataType(PositionDataType.INSTANCE); @@ -336,7 +347,6 @@ public class JsonWriterTestCase { private Document readDocumentFromJson(String docId, String fields) throws IOException { InputStream rawDoc = new ByteArrayInputStream(asFeed(docId, fields)); - JsonReader r = new JsonReader(types, rawDoc, parserFactory); DocumentParseInfo raw = r.parseDocument().get(); DocumentType docType = r.readDocumentType(raw.documentId); @@ -433,6 +443,30 @@ public class JsonWriterTestCase { } @Test + public void testTensorShortForm() throws IOException { + DocumentType documentTypeWithTensor = types.getDocumentType("testindexedtensor"); + String docId = "id:unittest:testindexedtensor::0"; + Document doc = new Document(documentTypeWithTensor, docId); + Field tensorField = documentTypeWithTensor.getField("tensorfield"); + Tensor tensor = Tensor.from("tensor(x[3]):[1,2,3]"); + doc.setFieldValue(tensorField, new TensorFieldValue(tensor)); + + assertEqualJson(asDocument(docId, "{ \"tensorfield\": {\"cells\":[{\"address\":{\"x\":\"0\"},\"value\":1.0},{\"address\":{\"x\":\"1\"},\"value\":2.0},{\"address\":{\"x\":\"2\"},\"value\":3.0}]} }"), + writeDocument(doc, false)); + assertEqualJson(asDocument(docId, "{ \"tensorfield\": {\"type\":\"tensor(x[3])\", \"values\":[1.0, 2.0, 3.0] } }"), + writeDocument(doc, true)); + } + + private byte[] writeDocument(Document doc, boolean tensorShortForm) throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + JsonFactory factory = new JsonFactory(); + JsonGenerator generator = factory.createGenerator(out); + JsonWriter writer = new JsonWriter(generator, tensorShortForm); + writer.write(doc); + return out.toByteArray(); + } + + @Test public void non_empty_reference_field_is_roundtrip_json_serialized() throws IOException { roundTripEquality("id:unittest:testrefs::helloworld", "{ \"ref_field\": \"id:unittest:smoke::and_mirrors_too\" }"); diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java index 73bea9f7a00..d3f8bb40b04 100644 --- a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java +++ b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java @@ -426,7 +426,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler { rawParameters = rawParameters.withFieldSet(path.documentType().orElseThrow() + ":[document]"); DocumentOperationParameters parameters = rawParameters.withResponseHandler(response -> { outstanding.decrementAndGet(); - handle(path, handler, response, (document, jsonResponse) -> { + handle(path, request, handler, response, (document, jsonResponse) -> { if (document != null) { jsonResponse.writeSingleDocument(document); jsonResponse.commit(Response.Status.OK); @@ -591,6 +591,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler { private final OutputStream out = new ContentChannelOutputStream(buffer); private final JsonGenerator json; private final ResponseHandler handler; + private final HttpRequest request; private final Queue<CompletionHandler> acks = new ConcurrentLinkedQueue<>(); private final Queue<ByteArrayOutputStream> docs = new ConcurrentLinkedQueue<>(); private final AtomicLong documentsWritten = new AtomicLong(); @@ -601,14 +602,24 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler { private ContentChannel channel; private JsonResponse(ResponseHandler handler) throws IOException { + this(handler, null); + } + + private JsonResponse(ResponseHandler handler, HttpRequest request) throws IOException { this.handler = handler; + this.request = request; json = jsonFactory.createGenerator(out); json.writeStartObject(); } /** Creates a new JsonResponse with path and id fields written. */ static JsonResponse create(DocumentPath path, ResponseHandler handler) throws IOException { - JsonResponse response = new JsonResponse(handler); + return create(path, handler, null); + } + + /** Creates a new JsonResponse with path and id fields written. */ + static JsonResponse create(DocumentPath path, ResponseHandler handler, HttpRequest request) throws IOException { + JsonResponse response = new JsonResponse(handler, request); response.writePathId(path.rawPath()); response.writeDocId(path.id()); return response; @@ -703,7 +714,11 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler { } synchronized void writeSingleDocument(Document document) throws IOException { - new JsonWriter(json).writeFields(document); + boolean tensorShortForm = false; + if (request != null && request.parameters().containsKey("format.tensors")) { + tensorShortForm = request.parameters().get("format.tensors").contains("short"); + } + new JsonWriter(json, tensorShortForm).writeFields(document); } synchronized void writeDocumentsArrayStart() throws IOException { @@ -1001,8 +1016,8 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler { void onSuccess(Document document, JsonResponse response) throws IOException; } - private static void handle(DocumentPath path, ResponseHandler handler, com.yahoo.documentapi.Response response, SuccessCallback callback) { - try (JsonResponse jsonResponse = JsonResponse.create(path, handler)) { + private static void handle(DocumentPath path, HttpRequest request, ResponseHandler handler, com.yahoo.documentapi.Response response, SuccessCallback callback) { + try (JsonResponse jsonResponse = JsonResponse.create(path, handler, request)) { jsonResponse.writeTrace(response.getTrace()); if (response.isSuccess()) callback.onSuccess((response instanceof DocumentResponse) ? ((DocumentResponse) response).getDocument() : null, jsonResponse); @@ -1037,7 +1052,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler { } private static void handleFeedOperation(DocumentPath path, ResponseHandler handler, com.yahoo.documentapi.Response response) { - handle(path, handler, response, (document, jsonResponse) -> jsonResponse.commit(Response.Status.OK)); + handle(path, null, handler, response, (document, jsonResponse) -> jsonResponse.commit(Response.Status.OK)); } private void updatePutMetrics(Outcome outcome) { |