TensorRemoveUpdate support for mixed tensors

3 files changed, 76 insertions, 10 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 e2736dabd2b..01293cb9782 100644
--- a/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java
+++ b/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java
@@ -40,6 +40,7 @@ public class DocumentUpdateJsonSerializerTest {
 
     final static TensorType sparseTensorType = new TensorType.Builder().mapped("x").mapped("y").build();
     final static TensorType denseTensorType = new TensorType.Builder().indexed("x", 2).indexed("y", 3).build();
+    final static TensorType mixedTensorType = new TensorType.Builder().mapped("x").indexed("y", 3).build();
     final static DocumentTypeManager types = new DocumentTypeManager();
     final static JsonFactory parserFactory = new JsonFactory();
     final static DocumentType docType = new DocumentType("doctype");
@@ -60,6 +61,7 @@ public class DocumentUpdateJsonSerializerTest {
         docType.addField(new Field("byte_field", DataType.BYTE));
         docType.addField(new Field("sparse_tensor", new TensorDataType(sparseTensorType)));
         docType.addField(new Field("dense_tensor", new TensorDataType(denseTensorType)));
+        docType.addField(new Field("mixed_tensor", new TensorDataType(mixedTensorType)));
         docType.addField(new Field("reference_field", new ReferenceDataType(refTargetDocType, 777)));
         docType.addField(new Field("predicate_field", DataType.PREDICATE));
         docType.addField(new Field("raw_field", DataType.RAW));
@@ -355,6 +357,25 @@ public class DocumentUpdateJsonSerializerTest {
     }
 
     @Test
+    public void test_tensor_add_update_mixed() {
+        roundtripSerializeJsonAndMatch(inputJson(
+                "{",
+                "  'update': 'DOCUMENT_ID',",
+                "  'fields': {",
+                "    'mixed_tensor': {",
+                "      'add': {",
+                "        'cells': [",
+                "          { 'address': { 'x': '0', 'y': '0' }, 'value': 2.0 },",
+                "          { 'address': { 'x': '1', 'y': '2' }, 'value': 3.0 }",
+                "        ]",
+                "      }",
+                "    }",
+                "  }",
+                "}"
+        ));
+    }
+
+    @Test
     public void test_tensor_remove_update() {
         roundtripSerializeJsonAndMatch(inputJson(
                 "{",
@@ -374,6 +395,25 @@ public class DocumentUpdateJsonSerializerTest {
     }
 
     @Test
+    public void test_tensor_remove_update_mixed() {
+        roundtripSerializeJsonAndMatch(inputJson(
+                "{",
+                "  'update': 'DOCUMENT_ID',",
+                "  'fields': {",
+                "    'mixed_tensor': {",
+                "      'remove': {",
+                "        'addresses': [",
+                "          {'x':'0' }",
+                "        ]",
+                "      }",
+                "    }",
+                "  }",
+                "}"
+        ));
+    }
+
+
+    @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 a20276e5c65..fe24a755d1d 100644
--- a/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
+++ b/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
@@ -1489,6 +1489,7 @@ public class JsonReaderTestCase {
         exception.expect(IllegalArgumentException.class);
         exception.expectMessage("Add update for field 'sparse_tensor' does not contain tensor cells");
         createTensorAddUpdate(inputJson("{}"), "sparse_tensor");
+        createTensorAddUpdate(inputJson("{}"), "mixed_tensor");
     }
 
     @Test
@@ -1500,7 +1501,6 @@ public class JsonReaderTestCase {
                         "    { 'x': 'c', 'y': 'd' } ]}"));
     }
 
-    @Ignore
     @Test
     public void tensor_remove_update_on_mixed_tensor() {
         assertTensorRemoveUpdate("{{x:1}:1.0,{x:2}:1.0}", "mixed_tensor",
@@ -1511,9 +1511,19 @@ public class JsonReaderTestCase {
     }
 
     @Test
+    public void tensor_remove_update_on_mixed_tensor_with_dense_addresses_throws() {
+        exception.expect(IllegalArgumentException.class);
+        exception.expectMessage("Indexed dimension address 'y' should not be specified in remove update");
+        createTensorRemoveUpdate(inputJson("{",
+                        "  'addresses': [",
+                        "    { 'x': '1', 'y': '0' },",
+                        "    { 'x': '2', 'y': '0' } ]}"), "mixed_tensor");
+    }
+
+    @Test
     public void tensor_remove_update_on_dense_tensor_throws() {
         exception.expect(IllegalArgumentException.class);
-        exception.expectMessage("A remove update can only be applied to sparse tensors. Field 'dense_tensor' has unsupported tensor type 'tensor(x[2],y[3])'");
+        exception.expectMessage("A remove update can only be applied to tensors with at least one sparse dimension. Field 'dense_tensor' has unsupported tensor type 'tensor(x[2],y[3])'");
         createTensorRemoveUpdate(inputJson("{",
                 "  'addresses': [] }"), "dense_tensor");
     }
@@ -1532,6 +1542,7 @@ public class JsonReaderTestCase {
         exception.expect(IllegalArgumentException.class);
         exception.expectMessage("Remove update for field 'sparse_tensor' does not contain tensor addresses");
         createTensorRemoveUpdate(inputJson("{'addresses': [] }"), "sparse_tensor");
+        createTensorRemoveUpdate(inputJson("{'addresses': [] }"), "mixed_tensor");
     }
 
     @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 40ab00facdb..52ed6c63356 100644
--- a/document/src/test/java/com/yahoo/document/update/TensorRemoveUpdateTest.java
+++ b/document/src/test/java/com/yahoo/document/update/TensorRemoveUpdateTest.java
@@ -10,17 +10,32 @@ import static org.junit.Assert.assertEquals;
 public class TensorRemoveUpdateTest {
 
     @Test
-    public void apply_remove_update_operations() {
-        assertApplyTo("{{x:0,y:0}:2, {x:0,y:1}:3}", "{{x:0,y:1}:1}", "{{x:0,y:0}:2}");
-        assertApplyTo("{{x:0,y:0}:1, {x:0,y:1}:2}", "{{x:0,y:0}:1,{x:0,y:1}:1}", "{}");
-        assertApplyTo("{}", "{{x:0,y:0}:1}", "{}");
-        assertApplyTo("{{x:0,y:0}:2, {x:0,y:1}:3}", "{}", "{{x:0,y:0}:2, {x:0,y:1}:3}");
+    public void apply_remove_update_operations_sparse() {
+        assertSparseApplyTo("{{x:0,y:0}:2, {x:0,y:1}:3}", "{{x:0,y:1}:1}", "{{x:0,y:0}:2}");
+        assertSparseApplyTo("{{x:0,y:0}:1, {x:0,y:1}:2}", "{{x:0,y:0}:1,{x:0,y:1}:1}", "{}");
+        assertSparseApplyTo("{}", "{{x:0,y:0}:1}", "{}");
+        assertSparseApplyTo("{{x:0,y:0}:2, {x:0,y:1}:3}", "{}", "{{x:0,y:0}:2, {x:0,y:1}:3}");
     }
 
-    private void assertApplyTo(String init, String update, String expected) {
-        String spec = "tensor(x{},y{})";
+    @Test
+    public void apply_remove_update_operations_mixed() {
+        assertMixedApplyTo("{{x:0,y:0}:2, {x:0,y:1}:3}", "{{x:0}:1}", "{}");
+        assertMixedApplyTo("{{x:0,y:0}:1, {x:1,y:0}:2}", "{{x:0}:1}", "{{x:1,y:0}:2,{x:1,y:1}:0,{x:1,y:2}:0}");
+        assertMixedApplyTo("{}", "{{x:0}:1}", "{}");
+        assertMixedApplyTo("{{x:0,y:0}:2, {x:0,y:1}:3}", "{}", "{{x:0,y:0}:2, {x:0,y:1}:3}");
+    }
+
+    private void assertSparseApplyTo(String init, String update, String expected) {
+        assertApplyTo("tensor(x{},y{})", "tensor(x{},y{})", init, update, expected);
+    }
+
+    private void assertMixedApplyTo(String init, String update, String expected) {
+        assertApplyTo("tensor(x{},y[3])", "tensor(x{})", init, update, expected);
+    }
+
+    private void assertApplyTo(String spec, String updateSpec, String init, String update, String expected) {
         TensorFieldValue initialFieldValue = new TensorFieldValue(Tensor.from(spec, init));
-        TensorRemoveUpdate removeUpdate = new TensorRemoveUpdate(new TensorFieldValue(Tensor.from(spec, update)));
+        TensorRemoveUpdate removeUpdate = new TensorRemoveUpdate(new TensorFieldValue(Tensor.from(updateSpec, update)));
         TensorFieldValue updatedFieldValue = (TensorFieldValue) removeUpdate.applyTo(initialFieldValue);
         assertEquals(Tensor.from(spec, expected), updatedFieldValue.getTensor().get());
     }