TensorRemoveUpdate support for mixed tensors

1 files changed, 18 insertions, 0 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java
index 3630a016691..00229c56171 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java
@@ -6,9 +6,11 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.DoubleBinaryOperator;
 import java.util.stream.Collectors;
 
@@ -127,6 +129,22 @@ public class MixedTensor implements Tensor {
     }
 
     @Override
+    public Tensor remove(Set<TensorAddress> addresses) {
+        Tensor.Builder builder = Tensor.Builder.of(type());
+        for (Map.Entry<TensorAddress, Long> entry : index.sparseMap.entrySet()) {
+            TensorAddress sparsePartialAddress = entry.getKey();
+            if ( ! addresses.contains(sparsePartialAddress)) {
+                long offset = entry.getValue();
+                for (int i = 0; i < index.denseSubspaceSize; ++i) {
+                    Cell cell = cells.get((int)offset + i);
+                    builder.cell(cell.getKey(), cell.getValue());
+                }
+            }
+        }
+        return builder.build();
+    }
+
+    @Override
     public int hashCode() { return cells.hashCode(); }
 
     @Override