Extend tensor remove update to support not fully specified addresses and update JSON parser.

Previously, all the sparse dimensions of the sparse or mixed tensor type (to remove from)
had to be specified in the addresses to remove.

3 files changed, 77 insertions, 6 deletions

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 4e8fa427e7d..1772a410a36 100644
--- a/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java
+++ b/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java
@@ -436,6 +436,25 @@ public class DocumentUpdateJsonSerializerTest {
     }
 
     @Test
+    public void test_tensor_remove_update_with_not_fully_specified_address() {
+        roundtripSerializeJsonAndMatch(inputJson(
+                "{",
+                "  'update': 'DOCUMENT_ID',",
+                "  'fields': {",
+                "    'sparse_tensor': {",
+                "      'remove': {",
+                "        'addresses': [",
+                "          {'y':'0'},",
+                "          {'y':'2'}",
+                "        ]",
+                "      }",
+                "    }",
+                "  }",
+                "}"
+        ));
+    }
+
+    @Test
     public void reference_field_id_can_be_update_assigned_non_empty_id() {
         roundtripSerializeJsonAndMatch(inputJson(
                 "{",
diff --git a/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java b/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
index 7fc43656d55..da9ab4ea7bf 100644
--- a/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
+++ b/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
@@ -168,6 +168,8 @@ public class JsonReaderTestCase {
                     new TensorDataType(new TensorType.Builder().indexed("x").indexed("y").build())));
             x.addField(new Field("mixed_tensor",
                     new TensorDataType(new TensorType.Builder().mapped("x").indexed("y", 3).build())));
+            x.addField(new Field("mixed_tensor_adv",
+                    new TensorDataType(new TensorType.Builder().mapped("x").mapped("y").mapped("z").indexed("a", 3).build())));
             types.registerDocumentType(x);
         }
         {
@@ -1685,6 +1687,24 @@ public class JsonReaderTestCase {
     }
 
     @Test
+    public void tensor_remove_update_on_sparse_tensor_with_not_fully_specified_address() {
+        assertTensorRemoveUpdate("{{y:b}:1.0,{y:d}:1.0}", "sparse_tensor",
+                inputJson("{",
+                        "  'addresses': [",
+                        "    { 'y': 'b' },",
+                        "    { 'y': 'd' } ]}"));
+    }
+
+    @Test
+    public void tensor_remove_update_on_mixed_tensor_with_not_fully_specified_address() {
+        assertTensorRemoveUpdate("{{x:1,z:a}:1.0,{x:2,z:b}:1.0}", "mixed_tensor_adv",
+                inputJson("{",
+                        "  'addresses': [",
+                        "    { 'x': '1', 'z': 'a' },",
+                        "    { 'x': '2', 'z': 'b' } ]}"));
+    }
+
+    @Test
     public void tensor_remove_update_on_mixed_tensor_with_dense_addresses_throws() {
         illegalTensorRemoveUpdate("Error in 'mixed_tensor': Indexed dimension address 'y' should not be specified in remove update",
                                   "mixed_tensor",
@@ -1703,12 +1723,19 @@ public class JsonReaderTestCase {
     }
 
     @Test
-    public void tensor_remove_update_on_not_fully_specified_cell_throws() {
-        illegalTensorRemoveUpdate("Error in 'sparse_tensor': Missing a label for dimension y for tensor(x{},y{})",
-                                  "sparse_tensor",
-                                  "{",
-                                  "  'addresses': [",
-                                  "    { 'x': 'a' } ]}");
+    public void tensor_remove_update_with_stray_dimension_throws() {
+        illegalTensorRemoveUpdate("Error in 'sparse_tensor': tensor(x{},y{}) does not contain dimension 'foo'",
+                "sparse_tensor",
+                "{",
+                "  'addresses': [",
+                "    { 'x': 'a', 'foo': 'b' } ]}");
+
+        illegalTensorRemoveUpdate("Error in 'sparse_tensor': tensor(x{}) does not contain dimension 'foo'",
+                "sparse_tensor",
+                "{",
+                "  'addresses': [",
+                "    { 'x': 'c' },",
+                "    { 'x': 'a', 'foo': 'b' } ]}");
     }
 
     @Test
diff --git a/document/src/test/java/com/yahoo/document/update/TensorRemoveUpdateTest.java b/document/src/test/java/com/yahoo/document/update/TensorRemoveUpdateTest.java
index 3a005e858c8..86f07db1b2d 100644
--- a/document/src/test/java/com/yahoo/document/update/TensorRemoveUpdateTest.java
+++ b/document/src/test/java/com/yahoo/document/update/TensorRemoveUpdateTest.java
@@ -3,9 +3,12 @@ package com.yahoo.document.update;
 
 import com.yahoo.document.datatypes.TensorFieldValue;
 import com.yahoo.tensor.Tensor;
+import com.yahoo.tensor.TensorType;
+import com.yahoo.yolean.Exceptions;
 import org.junit.Test;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 
 public class TensorRemoveUpdateTest {
 
@@ -22,4 +25,26 @@ public class TensorRemoveUpdateTest {
         assertEquals(Tensor.from(spec, expected), updatedFieldValue.getTensor().get());
     }
 
+    @Test
+    public void verify_compatible_type_throws_on_mismatch() {
+        // Contains an indexed dimension, which is not allowed.
+        illegalTensorRemoveUpdate("tensor(x{},y[1])", "{{x:a,y:0}:1}", "tensor(x{},y[1])",
+                "Unexpected type 'tensor(x{},y[1])' in remove update. Expected dimensions to be a subset of 'tensor(x{})'");
+
+        // Sparse dimension is not found in the original type.
+        illegalTensorRemoveUpdate("tensor(y{})", "{{y:a}:1}", "tensor(x{},z[2])",
+                "Unexpected type 'tensor(y{})' in remove update. Expected dimensions to be a subset of 'tensor(x{})'");
+    }
+
+    private void illegalTensorRemoveUpdate(String updateType, String updateTensor, String originalType, String expectedMessage) {
+        try {
+            var value = new TensorFieldValue(Tensor.from(updateType, updateTensor));
+            new TensorRemoveUpdate(value).verifyCompatibleType(TensorType.fromSpec(originalType));
+            fail("Expected exception");
+        }
+        catch (IllegalArgumentException expected) {
+            assertEquals(expectedMessage, Exceptions.toMessageString(expected));
+        }
+    }
+
 }