TensorRemoveUpdate support for mixed tensors

1 files changed, 14 insertions, 0 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java
index f44b3ce13b7..22ceed22d3e 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java
@@ -5,6 +5,7 @@ import com.google.common.collect.ImmutableMap;
 
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.DoubleBinaryOperator;
 
 /**
@@ -71,6 +72,19 @@ public class MappedTensor implements Tensor {
     }
 
     @Override
+    public Tensor remove(Set<TensorAddress> addresses) {
+        Tensor.Builder builder = Tensor.Builder.of(type());
+        for (Iterator<Tensor.Cell> i = cellIterator(); i.hasNext(); ) {
+            Tensor.Cell cell = i.next();
+            TensorAddress address = cell.getKey();
+            if ( ! addresses.contains(address)) {
+                builder.cell(address, cell.getValue());
+            }
+        }
+        return builder.build();
+    }
+
+    @Override
     public int hashCode() { return cells.hashCode(); }
 
     @Override