Add serialization for TensorRemoveUpdate

12 files changed, 105 insertions, 16 deletions

diff --git a/document/abi-spec.json b/document/abi-spec.json
index 7b1bc5d2af2..61390af3523 100644
--- a/document/abi-spec.json
+++ b/document/abi-spec.json
@@ -4495,7 +4495,8 @@
       "public abstract void write(com.yahoo.document.update.RemoveValueUpdate, com.yahoo.document.DataType)",
       "public abstract void write(com.yahoo.document.update.ClearValueUpdate, com.yahoo.document.DataType)",
       "public abstract void write(com.yahoo.document.update.TensorModifyUpdate)",
-      "public abstract void write(com.yahoo.document.update.TensorAddUpdate)"
+      "public abstract void write(com.yahoo.document.update.TensorAddUpdate)",
+      "public abstract void write(com.yahoo.document.update.TensorRemoveUpdate)"
     ],
     "fields": []
   },
@@ -4750,7 +4751,8 @@
       "public void read(com.yahoo.document.annotation.SpanList)",
       "public void read(com.yahoo.document.annotation.AlternateSpanList)",
       "protected com.yahoo.document.update.ValueUpdate readTensorModifyUpdate(com.yahoo.document.DataType)",
-      "protected com.yahoo.document.update.ValueUpdate readTensorAddUpdate(com.yahoo.document.DataType)"
+      "protected com.yahoo.document.update.ValueUpdate readTensorAddUpdate(com.yahoo.document.DataType)",
+      "protected com.yahoo.document.update.ValueUpdate readTensorRemoveUpdate(com.yahoo.document.DataType)"
     ],
     "fields": []
   },
@@ -4763,7 +4765,8 @@
     "methods": [
       "public void <init>(com.yahoo.document.DocumentTypeManager, com.yahoo.io.GrowableByteBuffer)",
       "protected com.yahoo.document.update.ValueUpdate readTensorModifyUpdate(com.yahoo.document.DataType)",
-      "protected com.yahoo.document.update.ValueUpdate readTensorAddUpdate(com.yahoo.document.DataType)"
+      "protected com.yahoo.document.update.ValueUpdate readTensorAddUpdate(com.yahoo.document.DataType)",
+      "protected com.yahoo.document.update.ValueUpdate readTensorRemoveUpdate(com.yahoo.document.DataType)"
     ],
     "fields": []
   },
@@ -4815,6 +4818,7 @@
       "public void write(com.yahoo.document.update.ClearValueUpdate, com.yahoo.document.DataType)",
       "public void write(com.yahoo.document.update.TensorModifyUpdate)",
       "public void write(com.yahoo.document.update.TensorAddUpdate)",
+      "public void write(com.yahoo.document.update.TensorRemoveUpdate)",
       "public static long getSerializedSize(com.yahoo.document.Document)"
     ],
     "fields": []
@@ -4870,6 +4874,7 @@
       "public void write(com.yahoo.document.update.ClearValueUpdate, com.yahoo.document.DataType)",
       "public void write(com.yahoo.document.update.TensorModifyUpdate)",
       "public void write(com.yahoo.document.update.TensorAddUpdate)",
+      "public void write(com.yahoo.document.update.TensorRemoveUpdate)",
       "public static long getSerializedSize(com.yahoo.document.Document)"
     ],
     "fields": []
@@ -4883,7 +4888,8 @@
     "methods": [
       "public void <init>(com.yahoo.io.GrowableByteBuffer)",
       "public void write(com.yahoo.document.update.TensorModifyUpdate)",
-      "public void write(com.yahoo.document.update.TensorAddUpdate)"
+      "public void write(com.yahoo.document.update.TensorAddUpdate)",
+      "public void write(com.yahoo.document.update.TensorRemoveUpdate)"
     ],
     "fields": []
   },
diff --git a/document/src/main/java/com/yahoo/document/json/DocumentUpdateJsonSerializer.java b/document/src/main/java/com/yahoo/document/json/DocumentUpdateJsonSerializer.java
index e9b5a136cf9..8d50f4a14dd 100644
--- a/document/src/main/java/com/yahoo/document/json/DocumentUpdateJsonSerializer.java
+++ b/document/src/main/java/com/yahoo/document/json/DocumentUpdateJsonSerializer.java
@@ -42,6 +42,7 @@ import com.yahoo.document.update.MapValueUpdate;
 import com.yahoo.document.update.RemoveValueUpdate;
 import com.yahoo.document.update.TensorAddUpdate;
 import com.yahoo.document.update.TensorModifyUpdate;
+import com.yahoo.document.update.TensorRemoveUpdate;
 import com.yahoo.document.update.ValueUpdate;
 import com.yahoo.vespa.objects.FieldBase;
 import com.yahoo.vespa.objects.Serializer;
@@ -51,7 +52,6 @@ import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.List;
 import java.util.Optional;
 import java.util.regex.Matcher;
 import java.util.stream.Collectors;
@@ -287,6 +287,17 @@ public class DocumentUpdateJsonSerializer
         }
 
         @Override
+        public void write(TensorRemoveUpdate update) {
+            wrapIOException(() -> {
+                generator.writeObjectFieldStart("remove");
+                if (update.getValue().getTensor().isPresent()) {
+                    serializeTensorAddresses(generator, update.getValue().getTensor().get());
+                }
+                generator.writeEndObject();
+            });
+        }
+
+        @Override
         public void write(FieldBase field, FieldValue value) {
             throw new JsonSerializationException(String.format("Serialization of field values of type %s is not supported", value.getClass().getName()));
         }
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 418837ca2a0..5fc9ee0b2f1 100644
--- a/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java
+++ b/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java
@@ -26,6 +26,7 @@ import com.yahoo.document.datatypes.StructuredFieldValue;
 import com.yahoo.document.datatypes.TensorFieldValue;
 import com.yahoo.document.datatypes.WeightedSet;
 import com.yahoo.document.json.readers.TensorReader;
+import com.yahoo.document.json.readers.TensorRemoveUpdateReader;
 import com.yahoo.document.serialization.FieldWriter;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorAddress;
@@ -87,7 +88,6 @@ public class JsonSerializationHelper {
         for (String dimension : dimensions) {
             generator.writeString(dimension);
         }
-
         generator.writeEndArray();
     }
 
@@ -99,7 +99,19 @@ public class JsonSerializationHelper {
             generator.writeNumberField(TensorReader.TENSOR_VALUE, cell.getValue());
             generator.writeEndObject();
         }
+        generator.writeEndArray();
+    }
 
+    static void serializeTensorAddresses(JsonGenerator generator, Tensor tensor) throws IOException {
+        TensorType tensorType = tensor.type();
+        generator.writeArrayFieldStart(TensorRemoveUpdateReader.TENSOR_ADDRESSES);
+        for (Map.Entry<TensorAddress, Double> cell : tensor.cells().entrySet()) {
+            generator.writeStartObject();
+            for (int i = 0; i < tensorType.dimensions().size(); i++) {
+                generator.writeStringField(tensorType.dimensions().get(i).name(), cell.getKey().label(i));
+            }
+            generator.writeEndObject();
+        }
         generator.writeEndArray();
     }
 
@@ -107,7 +119,6 @@ public class JsonSerializationHelper {
         generator.writeObjectFieldStart(TensorReader.TENSOR_ADDRESS);
         for (int i = 0; i < type.dimensions().size(); i++)
             generator.writeStringField(type.dimensions().get(i).name(), address.label(i));
-
         generator.writeEndObject();
     }
 
diff --git a/document/src/main/java/com/yahoo/document/json/readers/TensorRemoveUpdateReader.java b/document/src/main/java/com/yahoo/document/json/readers/TensorRemoveUpdateReader.java
index 6638320699c..210a6a80ee5 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/TensorRemoveUpdateReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/TensorRemoveUpdateReader.java
@@ -20,7 +20,7 @@ import static com.yahoo.document.json.readers.JsonParserHelpers.expectObjectStar
  */
 public class TensorRemoveUpdateReader {
 
-    private static final String TENSOR_ADDRESSES = "addresses";
+    public static final String TENSOR_ADDRESSES = "addresses";
 
     static TensorRemoveUpdate createTensorRemoveUpdate(TokenBuffer buffer, Field field) {
         expectObjectStart(buffer.currentToken());
diff --git a/document/src/main/java/com/yahoo/document/serialization/DocumentUpdateWriter.java b/document/src/main/java/com/yahoo/document/serialization/DocumentUpdateWriter.java
index 22b13d57aea..4b87edeeded 100644
--- a/document/src/main/java/com/yahoo/document/serialization/DocumentUpdateWriter.java
+++ b/document/src/main/java/com/yahoo/document/serialization/DocumentUpdateWriter.java
@@ -12,6 +12,7 @@ import com.yahoo.document.update.MapValueUpdate;
 import com.yahoo.document.update.RemoveValueUpdate;
 import com.yahoo.document.update.TensorAddUpdate;
 import com.yahoo.document.update.TensorModifyUpdate;
+import com.yahoo.document.update.TensorRemoveUpdate;
 
 /**
  * Interface for writing document updates in custom serializers.
@@ -30,4 +31,5 @@ public interface DocumentUpdateWriter {
     public void write(ClearValueUpdate clearValueUpdate, DataType superType);
     public void write(TensorModifyUpdate update);
     public void write(TensorAddUpdate update);
+    public void write(TensorRemoveUpdate update);
 }
diff --git a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializer6.java b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializer6.java
index 2943f6a8b22..ef33fb6b804 100644
--- a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializer6.java
+++ b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializer6.java
@@ -657,6 +657,8 @@ public class VespaDocumentDeserializer6 extends BufferSerializer implements Docu
                 return readTensorModifyUpdate(superType);
             case TENSORADD:
                 return readTensorAddUpdate(superType);
+            case TENSORREMOVE:
+                return readTensorRemoveUpdate(superType);
             default:
                 throw new DeserializationException(
                         "Could not deserialize ValueUpdate, unknown valueUpdateClassID type " + vuTypeId);
@@ -889,4 +891,8 @@ public class VespaDocumentDeserializer6 extends BufferSerializer implements Docu
         throw new DeserializationException("Cannot deserialize TensorAddUpdate, not implemented for Vespa 6");
     }
 
+    protected ValueUpdate readTensorRemoveUpdate(DataType type) {
+        throw new DeserializationException("Cannot deserialize TensorRemoveUpdate, not implemented for Vespa 6");
+    }
+
 }
diff --git a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializerHead.java b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializerHead.java
index 0f19937db41..2f22def9aa1 100644
--- a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializerHead.java
+++ b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializerHead.java
@@ -4,12 +4,13 @@ package com.yahoo.document.serialization;
 import com.yahoo.document.DataType;
 import com.yahoo.document.DocumentTypeManager;
 import com.yahoo.document.TensorDataType;
-import com.yahoo.document.datatypes.FieldValue;
 import com.yahoo.document.datatypes.TensorFieldValue;
 import com.yahoo.document.update.TensorAddUpdate;
 import com.yahoo.document.update.TensorModifyUpdate;
+import com.yahoo.document.update.TensorRemoveUpdate;
 import com.yahoo.document.update.ValueUpdate;
 import com.yahoo.io.GrowableByteBuffer;
+import com.yahoo.tensor.TensorType;
 
 /**
  * Class used for de-serializing documents on the current head document format.
@@ -49,4 +50,19 @@ public class VespaDocumentDeserializerHead extends VespaDocumentDeserializer6 {
         tensor.deserialize(this);
         return new TensorAddUpdate(tensor);
     }
+
+    @Override
+    protected ValueUpdate readTensorRemoveUpdate(DataType type) {
+        if (!(type instanceof TensorDataType)) {
+            throw new DeserializationException("Expected tensor data type, got " + type);
+        }
+        TensorDataType tensorDataType = (TensorDataType)type;
+        TensorType tensorType = tensorDataType.getTensorType();
+
+        // TODO: for mixed case extract a new tensor type based only on mapped dimensions
+
+        TensorFieldValue tensor = new TensorFieldValue(tensorType);
+        tensor.deserialize(this);
+        return new TensorRemoveUpdate(tensor);
+    }
 }
diff --git a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializer42.java b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializer42.java
index 28ae1c3db90..746b3c9d68d 100644
--- a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializer42.java
+++ b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializer42.java
@@ -47,6 +47,7 @@ import com.yahoo.document.update.MapValueUpdate;
 import com.yahoo.document.update.RemoveValueUpdate;
 import com.yahoo.document.update.TensorAddUpdate;
 import com.yahoo.document.update.TensorModifyUpdate;
+import com.yahoo.document.update.TensorRemoveUpdate;
 import com.yahoo.document.update.ValueUpdate;
 import com.yahoo.io.GrowableByteBuffer;
 import com.yahoo.tensor.serialization.TypedBinaryFormat;
@@ -648,12 +649,17 @@ public class VespaDocumentSerializer42 extends BufferSerializer implements Docum
 
     @Override
     public void write(TensorModifyUpdate update) {
-        throw new IllegalArgumentException("Write of TensorModifyUpdate not implemeneted for Vespa 4.2");
+        throw new IllegalArgumentException("Write of TensorModifyUpdate not implemented for Vespa 4.2");
     }
 
     @Override
     public void write(TensorAddUpdate update) {
-        throw new IllegalArgumentException("Write of TensorAddUpdate not implemeneted for Vespa 4.2");
+        throw new IllegalArgumentException("Write of TensorAddUpdate not implemented for Vespa 4.2");
+    }
+
+    @Override
+    public void write(TensorRemoveUpdate update) {
+        throw new IllegalArgumentException("Write of TensorRemoveUpdate not implemented for Vespa 4.2");
     }
 
     /**
diff --git a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializer6.java b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializer6.java
index 22f71bb3f65..8f033c84635 100644
--- a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializer6.java
+++ b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializer6.java
@@ -50,6 +50,7 @@ import com.yahoo.document.update.MapValueUpdate;
 import com.yahoo.document.update.RemoveValueUpdate;
 import com.yahoo.document.update.TensorAddUpdate;
 import com.yahoo.document.update.TensorModifyUpdate;
+import com.yahoo.document.update.TensorRemoveUpdate;
 import com.yahoo.document.update.ValueUpdate;
 import com.yahoo.io.GrowableByteBuffer;
 import com.yahoo.tensor.serialization.TypedBinaryFormat;
@@ -63,7 +64,6 @@ import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
-import java.util.logging.Logger;
 
 import static com.yahoo.text.Utf8.calculateBytePositions;
 
@@ -686,14 +686,20 @@ public class VespaDocumentSerializer6 extends BufferSerializer implements Docume
 
     @Override
     public void write(TensorModifyUpdate update) {
-        throw new IllegalArgumentException("Write of TensorModifyUpdate not implemeneted for Vespa 6");
+        throw new IllegalArgumentException("Write of TensorModifyUpdate not implemented for Vespa 6");
     }
 
     @Override
     public void write(TensorAddUpdate update) {
-        throw new IllegalArgumentException("Write of TensorAddUpdate not implemeneted for Vespa 6");
+        throw new IllegalArgumentException("Write of TensorAddUpdate not implemented for Vespa 6");
     }
 
+    @Override
+    public void write(TensorRemoveUpdate update) {
+        throw new IllegalArgumentException("Write of TensorRemoveUpdate not implemented for Vespa 6");
+    }
+
+
     /**
      * Returns the serialized size of the given {@link Document}. Please note that this method performs actual
      * serialization of the document, but simply return the size of the final {@link GrowableByteBuffer}. If you need
diff --git a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializerHead.java b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializerHead.java
index 144d3f768ab..4be4d6fcd93 100644
--- a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializerHead.java
+++ b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentSerializerHead.java
@@ -3,6 +3,7 @@ package com.yahoo.document.serialization;
 
 import com.yahoo.document.update.TensorAddUpdate;
 import com.yahoo.document.update.TensorModifyUpdate;
+import com.yahoo.document.update.TensorRemoveUpdate;
 import com.yahoo.io.GrowableByteBuffer;
 
 /**
@@ -27,4 +28,10 @@ public class VespaDocumentSerializerHead extends VespaDocumentSerializer6 {
         update.getValue().serialize(this);
     }
 
+    @Override
+    public void write(TensorRemoveUpdate update) {
+        update.getValue().serialize(this);
+    }
+
+
 }
diff --git a/document/src/main/java/com/yahoo/document/update/TensorRemoveUpdate.java b/document/src/main/java/com/yahoo/document/update/TensorRemoveUpdate.java
index 7af00abad1e..e9fb1e3efd5 100644
--- a/document/src/main/java/com/yahoo/document/update/TensorRemoveUpdate.java
+++ b/document/src/main/java/com/yahoo/document/update/TensorRemoveUpdate.java
@@ -36,8 +36,7 @@ public class TensorRemoveUpdate extends ValueUpdate<TensorFieldValue> {
 
     @Override
     public void serialize(DocumentUpdateWriter data, DataType superType) {
-        // TODO: implement
-        throw new UnsupportedOperationException("Not implemented yet");
+        data.write(this);
     }
 
     @Override
diff --git a/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java b/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java
index 65b51ad34b8..e2736dabd2b 100644
--- a/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java
+++ b/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java
@@ -355,6 +355,25 @@ public class DocumentUpdateJsonSerializerTest {
     }
 
     @Test
+    public void test_tensor_remove_update() {
+        roundtripSerializeJsonAndMatch(inputJson(
+                "{",
+                "  'update': 'DOCUMENT_ID',",
+                "  'fields': {",
+                "    'sparse_tensor': {",
+                "      'remove': {",
+                "        'addresses': [",
+                "          {'x':'0','y':'0'},",
+                "          {'x':'1','y':'2'}",
+                "        ]",
+                "      }",
+                "    }",
+                "  }",
+                "}"
+        ));
+    }
+
+    @Test
     public void reference_field_id_can_be_update_assigned_non_empty_id() {
         roundtripSerializeJsonAndMatch(inputJson(
                 "{",