Cleanup of tensor updates - Java

6 files changed, 25 insertions, 14 deletions

diff --git a/document/src/main/java/com/yahoo/document/json/readers/TensorAddUpdateReader.java b/document/src/main/java/com/yahoo/document/json/readers/TensorAddUpdateReader.java
index 5ed4455435a..6310fa62d15 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/TensorAddUpdateReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/TensorAddUpdateReader.java
@@ -27,11 +27,10 @@ public class TensorAddUpdateReader {
 
         TensorDataType tensorDataType = (TensorDataType)field.getDataType();
         TensorType tensorType = tensorDataType.getTensorType();
-
         TensorFieldValue tensorFieldValue = new TensorFieldValue(tensorType);
         fillTensor(buffer, tensorFieldValue);
-        expectTensorIsNonEmpty(field, tensorFieldValue.getTensor().get());
 
+        expectTensorIsNonEmpty(field, tensorFieldValue.getTensor().get());
         return new TensorAddUpdate(tensorFieldValue);
     }
 
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 aa5fed78bfe..66588debbca 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
@@ -29,7 +29,6 @@ public class TensorModifyUpdateReader {
     private static final String MODIFY_MULTIPLY = "multiply";
 
     public static TensorModifyUpdate createModifyUpdate(TokenBuffer buffer, Field field) {
-
         expectFieldIsOfTypeTensor(field);
         expectTensorTypeHasNoneIndexedUnboundDimensions(field);
         expectObjectStart(buffer.currentToken());
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 0d12e7c074b..3bb4b2e262f 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
@@ -28,9 +28,9 @@ public class TensorRemoveUpdateReader {
 
         TensorDataType tensorDataType = (TensorDataType)field.getDataType();
         TensorType originalType = tensorDataType.getTensorType();
-        TensorType convertedType = extractSparseDimensions(originalType);
-
+        TensorType convertedType = TensorRemoveUpdate.extractSparseDimensions(originalType);
         Tensor tensor = readRemoveUpdateTensor(buffer, convertedType, originalType);
+
         expectAddressesAreNonEmpty(field, tensor);
         return new TensorRemoveUpdate(new TensorFieldValue(tensor));
     }
@@ -87,9 +87,4 @@ public class TensorRemoveUpdateReader {
         return builder.build();
     }
 
-    public static TensorType extractSparseDimensions(TensorType type) {
-        TensorType.Builder builder = new TensorType.Builder();
-        type.dimensions().stream().filter(dim -> ! dim.isIndexed()).forEach(dim -> builder.mapped(dim.name()));
-        return builder.build();
-    }
 }
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 fb252b1a30a..a763db33e7a 100644
--- a/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializerHead.java
+++ b/document/src/main/java/com/yahoo/document/serialization/VespaDocumentDeserializerHead.java
@@ -63,7 +63,7 @@ public class VespaDocumentDeserializerHead extends VespaDocumentDeserializer6 {
         }
         TensorDataType tensorDataType = (TensorDataType)type;
         TensorType tensorType = tensorDataType.getTensorType();
-        TensorType convertedType = TensorRemoveUpdateReader.extractSparseDimensions(tensorType);
+        TensorType convertedType = TensorRemoveUpdate.extractSparseDimensions(tensorType);
 
         TensorFieldValue tensor = new TensorFieldValue(convertedType);
         tensor.deserialize(this);
diff --git a/document/src/main/java/com/yahoo/document/update/TensorAddUpdate.java b/document/src/main/java/com/yahoo/document/update/TensorAddUpdate.java
index d0d0cfc2480..f8d2464deb7 100644
--- a/document/src/main/java/com/yahoo/document/update/TensorAddUpdate.java
+++ b/document/src/main/java/com/yahoo/document/update/TensorAddUpdate.java
@@ -12,8 +12,6 @@ import java.util.Objects;
 
 /**
  *  An update used to add cells to a sparse or mixed tensor (has at least one mapped dimension).
- *
- *  The cells to add are contained in a sparse tensor.
  */
 public class TensorAddUpdate extends ValueUpdate<TensorFieldValue> {
 
@@ -50,7 +48,7 @@ public class TensorAddUpdate extends ValueUpdate<TensorFieldValue> {
 
         Tensor old = ((TensorFieldValue) oldValue).getTensor().get();
         Tensor update = tensor.getTensor().get();
-        Tensor result = old.merge((left, right) -> right, update.cells());
+        Tensor result = old.merge((left, right) -> right, update.cells());  // note this might be slow for large mixed tensor updates
         return new TensorFieldValue(result);
     }
 
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 0c5345c7a9f..335cda8e133 100644
--- a/document/src/main/java/com/yahoo/document/update/TensorRemoveUpdate.java
+++ b/document/src/main/java/com/yahoo/document/update/TensorRemoveUpdate.java
@@ -7,6 +7,7 @@ import com.yahoo.document.datatypes.FieldValue;
 import com.yahoo.document.datatypes.TensorFieldValue;
 import com.yahoo.document.serialization.DocumentUpdateWriter;
 import com.yahoo.tensor.Tensor;
+import com.yahoo.tensor.TensorType;
 
 import java.util.Objects;
 
@@ -22,6 +23,18 @@ public class TensorRemoveUpdate extends ValueUpdate<TensorFieldValue> {
     public TensorRemoveUpdate(TensorFieldValue value) {
         super(ValueUpdateClassID.TENSORREMOVE);
         this.tensor = value;
+        verifyCompatibleType();
+    }
+
+    private void verifyCompatibleType() {
+        if ( ! tensor.getTensor().isPresent()) {
+            throw new IllegalArgumentException("Tensor must be present in remove update");
+        }
+        TensorType tensorType = tensor.getTensor().get().type();
+        TensorType expectedType = extractSparseDimensions(tensor.getDataType().getTensorType());
+        if ( ! tensorType.equals(expectedType)) {
+            throw new IllegalArgumentException("Unexpected type '" + tensorType + "' in remove update. Expected is '" + expectedType + "'");
+        }
     }
 
     @Override
@@ -83,4 +96,11 @@ public class TensorRemoveUpdate extends ValueUpdate<TensorFieldValue> {
         return super.toString() + " " + tensor;
     }
 
+    public static TensorType extractSparseDimensions(TensorType type) {
+        TensorType.Builder builder = new TensorType.Builder();
+        type.dimensions().stream().filter(dim -> ! dim.isIndexed()).forEach(dim -> builder.mapped(dim.name()));
+        return builder.build();
+    }
+
+
 }