TensorModifyUpdate support for mixed tensors

5 files changed, 66 insertions, 24 deletions

diff --git a/document/src/main/java/com/yahoo/document/json/readers/TensorModifyUpdateReader.java b/document/src/main/java/com/yahoo/document/json/readers/TensorModifyUpdateReader.java
index 5022185e03f..69b36e7cfa4 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/TensorModifyUpdateReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/TensorModifyUpdateReader.java
@@ -32,7 +32,6 @@ public class TensorModifyUpdateReader {
 
         expectFieldIsOfTypeTensor(field);
         expectTensorTypeHasNoneIndexedUnboundDimensions(field);
-        expectTensorTypeIsNotMixed(field);
         expectObjectStart(buffer.currentToken());
 
         ModifyUpdateResult result = createModifyUpdateResult(buffer, field);
@@ -58,16 +57,6 @@ public class TensorModifyUpdateReader {
         }
     }
 
-    private static void expectTensorTypeIsNotMixed(Field field) {
-        TensorType tensorType = ((TensorDataType)field.getDataType()).getTensorType();
-        long numMappedDimensions = tensorType.dimensions().stream().filter(dim -> dim.type().equals(TensorType.Dimension.Type.mapped)).count();
-        long numIndexedDimensions = tensorType.dimensions().stream().filter(dim -> dim.isIndexed()).count();
-        if (numMappedDimensions > 0 && numIndexedDimensions > 0) {
-            throw new IllegalArgumentException("A modify update cannot be applied to tensor types with mixed dimensions. "
-                    + "Field '" + field.getName() + "' has mixed tensor type '" + tensorType + "'");
-        }
-    }
-
     private static void expectOperationSpecified(TensorModifyUpdate.Operation operation, String fieldName) {
         if (operation == null) {
             throw new IllegalArgumentException("Modify update for field '" + fieldName + "' does not contain an operation");
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 2ab7169fae2..d1855759cb6 100644
--- a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializerHead.java
+++ b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializerHead.java
@@ -36,7 +36,10 @@ public class VespaDocumentDeserializerHead extends VespaDocumentDeserializer6 {
             throw new DeserializationException("Expected tensor data type, got " + type);
         }
         TensorDataType tensorDataType = (TensorDataType)type;
-        TensorFieldValue tensor = new TensorFieldValue(TensorModifyUpdate.convertToCompatibleType(tensorDataType.getTensorType()));
+        TensorType tensorType = tensorDataType.getTensorType();
+        TensorType convertedType = TensorModifyUpdate.convertToCompatibleType(tensorType);
+
+        TensorFieldValue tensor = new TensorFieldValue(convertedType);
         tensor.deserialize(this);
         return new TensorModifyUpdate(operation, tensor);
     }
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 01293cb9782..dd090f7ba32 100644
--- a/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java
+++ b/document/src/test/java/com/yahoo/document/json/DocumentUpdateJsonSerializerTest.java
@@ -338,6 +338,26 @@ public class DocumentUpdateJsonSerializerTest {
     }
 
     @Test
+    public void test_tensor_modify_update_on_mixed_tensor() {
+        roundtripSerializeJsonAndMatch(inputJson(
+                "{",
+                "  'update': 'DOCUMENT_ID',",
+                "  'fields': {",
+                "    'mixed_tensor': {",
+                "      'modify': {",
+                "        'operation': 'multiply',",
+                "        'cells': [",
+                "          { 'address': { 'x': 'a', 'y': '0' }, 'value': 2.0 },",
+                "          { 'address': { 'x': 'c', 'y': '1' }, 'value': 3.0 }",
+                "        ]",
+                "      }",
+                "    }",
+                "  }",
+                "}"
+        ));
+    }
+
+    @Test
     public void test_tensor_add_update() {
         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 fa5479787d7..4bff6f44188 100644
--- a/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
+++ b/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
@@ -1387,12 +1387,30 @@ public class JsonReaderTestCase {
     }
 
     @Test
-    public void tensor_modify_update_on_mixed_tensor_throws() {
-        exception.expect(IllegalArgumentException.class);
-        exception.expectMessage("A modify update cannot be applied to tensor types with mixed dimensions. Field 'mixed_tensor' has mixed tensor type 'tensor(x{},y[3])'");
-        createTensorModifyUpdate(inputJson("{",
-                "  'operation': 'replace',",
-                "  'cells': [] }"), "mixed_tensor");
+    public void tensor_modify_update_with_replace_operation_mixed() {
+        assertTensorModifyUpdate("{{x:a,y:0}:2.0}", TensorModifyUpdate.Operation.REPLACE, "mixed_tensor",
+                inputJson("{",
+                        "  'operation': 'replace',",
+                        "  'cells': [",
+                        "    { 'address': { 'x': 'a', 'y': '0' }, 'value': 2.0 } ]}"));
+    }
+
+    @Test
+    public void tensor_modify_update_with_add_operation_mixed() {
+        assertTensorModifyUpdate("{{x:a,y:0}:2.0}", TensorModifyUpdate.Operation.ADD, "mixed_tensor",
+                inputJson("{",
+                        "  'operation': 'add',",
+                        "  'cells': [",
+                        "    { 'address': { 'x': 'a', 'y': '0' }, 'value': 2.0 } ]}"));
+    }
+
+    @Test
+    public void tensor_modify_update_with_multiply_operation_mixed() {
+        assertTensorModifyUpdate("{{x:a,y:0}:2.0}", TensorModifyUpdate.Operation.MULTIPLY, "mixed_tensor",
+                inputJson("{",
+                        "  'operation': 'multiply',",
+                        "  'cells': [",
+                        "    { 'address': { 'x': 'a', 'y': '0' }, 'value': 2.0 } ]}"));
     }
 
     @Test
@@ -1406,6 +1424,17 @@ public class JsonReaderTestCase {
     }
 
     @Test
+    public void tensor_modify_update_with_out_of_bound_cells_throws_mixed() {
+        exception.expect(IndexOutOfBoundsException.class);
+        exception.expectMessage("Dimension 'y' has label '3' but type is tensor(x{},y[3])");
+        createTensorModifyUpdate(inputJson("{",
+                "  'operation': 'replace',",
+                "  'cells': [",
+                "    { 'address': { 'x': '0', 'y': '3' }, 'value': 2.0 } ]}"), "mixed_tensor");
+    }
+
+
+    @Test
     public void tensor_modify_update_with_unknown_operation_throws() {
         exception.expect(IllegalArgumentException.class);
         exception.expectMessage("Unknown operation 'unknown' in modify update for field 'sparse_tensor'");
diff --git a/document/src/test/java/com/yahoo/document/update/TensorModifyUpdateTest.java b/document/src/test/java/com/yahoo/document/update/TensorModifyUpdateTest.java
index 6e9444de2be..e448f70cf54 100644
--- a/document/src/test/java/com/yahoo/document/update/TensorModifyUpdateTest.java
+++ b/document/src/test/java/com/yahoo/document/update/TensorModifyUpdateTest.java
@@ -1,12 +1,6 @@
 // Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.document.update;
 
-import com.yahoo.document.Document;
-import com.yahoo.document.DocumentId;
-import com.yahoo.document.DocumentType;
-import com.yahoo.document.DocumentTypeManager;
-import com.yahoo.document.Field;
-import com.yahoo.document.TensorDataType;
 import com.yahoo.document.datatypes.TensorFieldValue;
 import com.yahoo.document.update.TensorModifyUpdate.Operation;
 import com.yahoo.tensor.Tensor;
@@ -28,6 +22,7 @@ public class TensorModifyUpdateTest {
         assertConvertToCompatible("tensor(x{})", "tensor(x[10])");
         assertConvertToCompatible("tensor(x{})", "tensor(x{})");
         assertConvertToCompatible("tensor(x{},y{},z{})", "tensor(x[],y[10],z{})");
+        assertConvertToCompatible("tensor(x{},y{})", "tensor(x{},y[3])");
     }
 
     private static void assertConvertToCompatible(String expectedType, String inputType) {
@@ -56,6 +51,12 @@ public class TensorModifyUpdateTest {
                 "{{x:0,y:0}:1, {x:0,y:1}:2}", "{{x:0,y:1}:3}", "{{x:0,y:0}:1,{x:0,y:1}:5}");
         assertApplyTo("tensor(x[1],y[2])", Operation.MULTIPLY,
                 "{{x:0,y:0}:3, {x:0,y:1}:2}", "{{x:0,y:1}:3}", "{{x:0,y:0}:3,{x:0,y:1}:6}");
+        assertApplyTo("tensor(x{},y[2])", Operation.REPLACE,
+                "{{x:0,y:0}:1, {x:0,y:1}:2}", "{{x:0,y:1}:0}", "{{x:0,y:0}:1,{x:0,y:1}:0}");
+        assertApplyTo("tensor(x{},y[2])", Operation.ADD,
+                "{{x:0,y:0}:1, {x:0,y:1}:2}", "{{x:0,y:1}:3}", "{{x:0,y:0}:1,{x:0,y:1}:5}");
+        assertApplyTo("tensor(x{},y[2])", Operation.MULTIPLY,
+                "{{x:0,y:0}:3, {x:0,y:1}:2}", "{{x:0,y:1}:3}", "{{x:0,y:0}:3,{x:0,y:1}:6}");
     }
 
     private void assertApplyTo(String spec, Operation op, String init, String update, String expected) {