Implement applyTo() for TensorRemoveUpdate.

3 files changed, 38 insertions, 6 deletions

diff --git a/document/src/tests/documentupdatetestcase.cpp b/document/src/tests/documentupdatetestcase.cpp
index c74c211756f..017d83893f0 100644
--- a/document/src/tests/documentupdatetestcase.cpp
+++ b/document/src/tests/documentupdatetestcase.cpp
@@ -922,6 +922,17 @@ TEST(DocumentUpdateTest, tensor_add_update_can_be_applied)
                                 .add({{"x", "c"}}, 7));
 }
 
+TEST(DocumentUpdateTest, tensor_remove_update_can_be_applied)
+{
+    TensorUpdateFixture f;
+    f.assertApplyUpdate(f.spec().add({{"x", "a"}}, 2)
+                                .add({{"x", "b"}}, 3),
+
+                        TensorRemoveUpdate(f.makeTensor(f.spec().add({{"x", "b"}}, 1))),
+
+                        f.spec().add({{"x", "a"}}, 2));
+}
+
 TEST(DocumentUpdateTest, tensor_modify_update_can_be_applied)
 {
     TensorUpdateFixture f;
diff --git a/document/src/vespa/document/update/tensor_remove_update.cpp b/document/src/vespa/document/update/tensor_remove_update.cpp
index 3e2bb86c66b..671bf260629 100644
--- a/document/src/vespa/document/update/tensor_remove_update.cpp
+++ b/document/src/vespa/document/update/tensor_remove_update.cpp
@@ -6,6 +6,8 @@
 #include <vespa/document/fieldvalue/document.h>
 #include <vespa/document/fieldvalue/tensorfieldvalue.h>
 #include <vespa/document/serialization/vespadocumentdeserializer.h>
+#include <vespa/eval/tensor/cell_values.h>
+#include <vespa/eval/tensor/sparse/sparse_tensor.h>
 #include <vespa/eval/tensor/tensor.h>
 #include <vespa/vespalib/objects/nbostream.h>
 #include <vespa/vespalib/util/xmlstream.h>
@@ -77,17 +79,35 @@ TensorRemoveUpdate::checkCompatibility(const Field &field) const
 std::unique_ptr<Tensor>
 TensorRemoveUpdate::applyTo(const Tensor &tensor) const
 {
-    // TODO: implement
-    (void) tensor;
+    auto &addressTensor = _tensor->getAsTensorPtr();
+    if (addressTensor) {
+        if (const auto *sparseTensor = dynamic_cast<const vespalib::tensor::SparseTensor *>(addressTensor.get())) {
+            vespalib::tensor::CellValues cellAddresses(*sparseTensor);
+            return tensor.remove(cellAddresses);
+        } else {
+            throw IllegalArgumentException(make_string("Expected address tensor to be sparse, but has type '%s'",
+                                                       addressTensor->type().to_spec().c_str()));
+        }
+    }
     return std::unique_ptr<Tensor>();
 }
 
 bool
 TensorRemoveUpdate::applyTo(FieldValue &value) const
 {
-    // TODO: implement
-    (void) value;
-    return false;
+    if (value.inherits(TensorFieldValue::classId)) {
+        TensorFieldValue &tensorFieldValue = static_cast<TensorFieldValue &>(value);
+        auto &oldTensor = tensorFieldValue.getAsTensorPtr();
+        auto newTensor = applyTo(*oldTensor);
+        if (newTensor) {
+            tensorFieldValue = std::move(newTensor);
+        }
+    } else {
+        std::string err = make_string("Unable to perform a tensor remove update on a '%s' field value.",
+                                      value.getClass().name());
+        throw IllegalStateException(err, VESPA_STRLOC);
+    }
+    return true;
 }
 
 void
diff --git a/document/src/vespa/document/update/valueupdate.h b/document/src/vespa/document/update/valueupdate.h
index 0e15943f8e4..6939d10ce2c 100644
--- a/document/src/vespa/document/update/valueupdate.h
+++ b/document/src/vespa/document/update/valueupdate.h
@@ -55,7 +55,8 @@ public:
         Map        = IDENTIFIABLE_CLASSID(MapValueUpdate),
         Remove     = IDENTIFIABLE_CLASSID(RemoveValueUpdate),
         TensorModifyUpdate = IDENTIFIABLE_CLASSID(TensorModifyUpdate),
-        TensorAddUpdate = IDENTIFIABLE_CLASSID(TensorAddUpdate)
+        TensorAddUpdate = IDENTIFIABLE_CLASSID(TensorAddUpdate),
+        TensorRemoveUpdate = IDENTIFIABLE_CLASSID(TensorRemoveUpdate)
     };
 
     ValueUpdate()