Support accessing tensor refs directly when using an imported DirectTensorAttribute.

Also avoid dynamic casting when instantiating the attribute feature.
Instead extend ITensorAttribute with functions specifying what kind of direct access is supported.

19 files changed, 72 insertions, 55 deletions

diff --git a/searchlib/src/vespa/searchlib/features/attributefeature.cpp b/searchlib/src/vespa/searchlib/features/attributefeature.cpp
index 7c0ed1147f7..eac69f87587 100644
--- a/searchlib/src/vespa/searchlib/features/attributefeature.cpp
+++ b/searchlib/src/vespa/searchlib/features/attributefeature.cpp
@@ -469,13 +469,13 @@ createTensorAttributeExecutor(const IAttributeVector *attribute, const vespalib:
                 tensorType.to_spec().c_str());
         return ConstantTensorExecutor::createEmpty(tensorType, stash);
     }
-    if (tensorType.is_dense()) {
-        return stash.create<DenseTensorAttributeExecutor>(tensorAttribute);
+    if (tensorAttribute->supports_extract_dense_view()) {
+        return stash.create<DenseTensorAttributeExecutor>(*tensorAttribute);
     }
-    if (auto direct = dynamic_cast<const DirectTensorAttribute *>(tensorAttribute)) {
-        return stash.create<DirectTensorAttributeExecutor>(*direct);
+    if (tensorAttribute->supports_get_tensor_ref()) {
+        return stash.create<DirectTensorAttributeExecutor>(*tensorAttribute);
     }
-    return stash.create<TensorAttributeExecutor>(tensorAttribute);
+    return stash.create<TensorAttributeExecutor>(*tensorAttribute);
 }
 
 bool
diff --git a/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.cpp b/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.cpp
index 6eff09b65ab..3dfba3092cd 100644
--- a/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.cpp
+++ b/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.cpp
@@ -10,16 +10,16 @@ using vespalib::tensor::MutableDenseTensorView;
 namespace search::features {
 
 DenseTensorAttributeExecutor::
-DenseTensorAttributeExecutor(const ITensorAttribute *attribute)
+DenseTensorAttributeExecutor(const ITensorAttribute& attribute)
     : _attribute(attribute),
-      _tensorView(_attribute->getTensorType())
+      _tensorView(_attribute.getTensorType())
 {
 }
 
 void
 DenseTensorAttributeExecutor::execute(uint32_t docId)
 {
-    _attribute->getTensor(docId, _tensorView);
+    _attribute.extract_dense_view(docId, _tensorView);
     outputs().set_object(0, _tensorView);
 }
 
diff --git a/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.h b/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.h
index 2bcf356689e..a8a84447c88 100644
--- a/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.h
+++ b/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.h
@@ -16,11 +16,11 @@ namespace search::features {
 class DenseTensorAttributeExecutor : public fef::FeatureExecutor
 {
 private:
-    const search::tensor::ITensorAttribute *_attribute;
+    const search::tensor::ITensorAttribute& _attribute;
     vespalib::tensor::MutableDenseTensorView _tensorView;
 
 public:
-    DenseTensorAttributeExecutor(const search::tensor::ITensorAttribute *attribute);
+    DenseTensorAttributeExecutor(const search::tensor::ITensorAttribute& attribute);
     void execute(uint32_t docId) override;
 };
 
diff --git a/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.cpp b/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.cpp
index cd0f4d4d356..56765bd168f 100644
--- a/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.cpp
+++ b/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.cpp
@@ -1,14 +1,13 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "direct_tensor_attribute_executor.h"
-#include <vespa/searchlib/tensor/direct_tensor_attribute.h>
 #include <vespa/eval/tensor/tensor.h>
+#include <vespa/searchlib/tensor/i_tensor_attribute.h>
 
-namespace search {
-namespace features {
+namespace search::features {
 
 DirectTensorAttributeExecutor::
-DirectTensorAttributeExecutor(const DirectTensorAttribute &attribute)
+DirectTensorAttributeExecutor(const ITensorAttribute &attribute)
     : _attribute(attribute)
 {
 }
@@ -19,5 +18,4 @@ DirectTensorAttributeExecutor::execute(uint32_t docId)
     outputs().set_object(0, _attribute.get_tensor_ref(docId));
 }
 
-} // namespace features
-} // namespace search
+}
diff --git a/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.h b/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.h
index e311fba076e..b9a4df739ff 100644
--- a/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.h
+++ b/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.h
@@ -4,17 +4,17 @@
 
 #include <vespa/searchlib/fef/featureexecutor.h>
 
-namespace search::tensor { class DirectTensorAttribute; }
+namespace search::tensor { class ITensorAttribute; }
 namespace search::features {
 
 class DirectTensorAttributeExecutor : public fef::FeatureExecutor
 {
 public:
-    using DirectTensorAttribute = search::tensor::DirectTensorAttribute;
-    DirectTensorAttributeExecutor(const DirectTensorAttribute &attribute);
+    using ITensorAttribute = search::tensor::ITensorAttribute;
+    DirectTensorAttributeExecutor(const ITensorAttribute &attribute);
     void execute(uint32_t docId) override;
 private:
-    const DirectTensorAttribute &_attribute;
+    const ITensorAttribute &_attribute;
 };
 
 }
diff --git a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp
index 51727846f95..802e98dcd66 100644
--- a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp
+++ b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp
@@ -3,13 +3,12 @@
 #include "tensor_attribute_executor.h"
 #include <vespa/searchlib/tensor/i_tensor_attribute.h>
 
-namespace search {
-namespace features {
+namespace search::features {
 
 TensorAttributeExecutor::
-TensorAttributeExecutor(const search::tensor::ITensorAttribute *attribute)
+TensorAttributeExecutor(const search::tensor::ITensorAttribute& attribute)
     : _attribute(attribute),
-      _emptyTensor(attribute->getEmptyTensor()),
+      _emptyTensor(attribute.getEmptyTensor()),
       _tensor()
 {
 }
@@ -17,7 +16,7 @@ TensorAttributeExecutor(const search::tensor::ITensorAttribute *attribute)
 void
 TensorAttributeExecutor::execute(uint32_t docId)
 {
-    _tensor = _attribute->getTensor(docId);
+    _tensor = _attribute.getTensor(docId);
     if (_tensor) {
         outputs().set_object(0, *_tensor);
     } else {
@@ -25,5 +24,4 @@ TensorAttributeExecutor::execute(uint32_t docId)
     }
 }
 
-} // namespace features
-} // namespace search
+}
diff --git a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.h b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.h
index deba5c05a5d..be900e8a3e7 100644
--- a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.h
+++ b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.h
@@ -15,12 +15,12 @@ namespace search::features {
 class TensorAttributeExecutor : public fef::FeatureExecutor
 {
 private:
-    const search::tensor::ITensorAttribute *_attribute;
+    const search::tensor::ITensorAttribute& _attribute;
     std::unique_ptr<vespalib::eval::Tensor> _emptyTensor;
     std::unique_ptr<vespalib::eval::Tensor> _tensor;
 
 public:
-    TensorAttributeExecutor(const search::tensor::ITensorAttribute *attribute);
+    TensorAttributeExecutor(const search::tensor::ITensorAttribute& attribute);
     void execute(uint32_t docId) override;
 };
 
diff --git a/searchlib/src/vespa/searchlib/queryeval/nearest_neighbor_iterator.cpp b/searchlib/src/vespa/searchlib/queryeval/nearest_neighbor_iterator.cpp
index 07e10271c55..a64168469f0 100644
--- a/searchlib/src/vespa/searchlib/queryeval/nearest_neighbor_iterator.cpp
+++ b/searchlib/src/vespa/searchlib/queryeval/nearest_neighbor_iterator.cpp
@@ -76,7 +76,7 @@ public:
 
 private:
     double computeDistance(uint32_t docId, double limit) {
-        params().tensorAttribute.getTensor(docId, _fieldTensor);
+        params().tensorAttribute.extract_dense_view(docId, _fieldTensor);
         auto rhs = _fieldTensor.cellsRef();
         return params().distanceFunction->calc_with_limit(_lhs, rhs, limit);
     }
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
index 37a042d4e7f..d63249b816f 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
@@ -178,7 +178,7 @@ DenseTensorAttribute::getTensor(DocId docId) const
 }
 
 void
-DenseTensorAttribute::getTensor(DocId docId, MutableDenseTensorView &tensor) const
+DenseTensorAttribute::extract_dense_view(DocId docId, MutableDenseTensorView &tensor) const
 {
     EntryRef ref;
     if (docId < getCommittedDocIdLimit()) {
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
index 7fd06357114..1de74d4238d 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
@@ -37,7 +37,8 @@ public:
     std::unique_ptr<PrepareResult> prepare_set_tensor(DocId docid, const Tensor& tensor) const override;
     void complete_set_tensor(DocId docid, const Tensor& tensor, std::unique_ptr<PrepareResult> prepare_result) override;
     std::unique_ptr<Tensor> getTensor(DocId docId) const override;
-    void getTensor(DocId docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
+    void extract_dense_view(DocId docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
+    bool supports_extract_dense_view() const override { return true; }
     bool onLoad() override;
     std::unique_ptr<AttributeSaver> onInitSave(vespalib::stringref fileName) override;
     void compactWorst() override;
diff --git a/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.cpp
index 4dc6f5d6917..570be586e88 100644
--- a/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.cpp
@@ -108,12 +108,6 @@ DirectTensorAttribute::get_tensor_ref(DocId docId) const
     return *_emptyTensor;
 }
 
-void
-DirectTensorAttribute::getTensor(DocId, vespalib::tensor::MutableDenseTensorView &) const
-{
-    notImplemented();
-}
-
 std::unique_ptr<AttributeSaver>
 DirectTensorAttribute::onInitSave(vespalib::stringref fileName)
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.h
index a02e2678a48..1429a8b4d5d 100644
--- a/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.h
@@ -15,13 +15,13 @@ public:
     virtual ~DirectTensorAttribute();
     virtual void setTensor(DocId docId, const Tensor &tensor) override;
     virtual std::unique_ptr<Tensor> getTensor(DocId docId) const override;
-    virtual void getTensor(DocId docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
     virtual bool onLoad() override;
     virtual std::unique_ptr<AttributeSaver> onInitSave(vespalib::stringref fileName) override;
     virtual void compactWorst() override;
 
     void set_tensor(DocId docId, std::unique_ptr<Tensor> tensor);
-    const Tensor &get_tensor_ref(DocId docId) const;
+    const Tensor &get_tensor_ref(DocId docId) const override;
+    virtual bool supports_get_tensor_ref() const override { return true; }
 };
 
 }  // namespace search::tensor
diff --git a/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h
index 520abb13d06..12b6857d122 100644
--- a/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h
@@ -24,7 +24,11 @@ public:
     virtual ~ITensorAttribute() {}
     virtual std::unique_ptr<Tensor> getTensor(uint32_t docId) const = 0;
     virtual std::unique_ptr<Tensor> getEmptyTensor() const = 0;
-    virtual void getTensor(uint32_t docId, vespalib::tensor::MutableDenseTensorView &tensor) const = 0;
+    virtual void extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const = 0;
+    virtual const Tensor& get_tensor_ref(uint32_t docid) const = 0;
+    virtual bool supports_extract_dense_view() const = 0;
+    virtual bool supports_get_tensor_ref() const = 0;
+
     virtual vespalib::eval::ValueType getTensorType() const = 0;
 
     /**
diff --git a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp
index 593330b7bd4..2e396145138 100644
--- a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp
@@ -51,9 +51,15 @@ ImportedTensorAttributeVectorReadGuard::getEmptyTensor() const
 }
 
 void
-ImportedTensorAttributeVectorReadGuard::getTensor(uint32_t docId, vespalib::tensor::MutableDenseTensorView &tensor) const
+ImportedTensorAttributeVectorReadGuard::extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const
 {
-    _target_tensor_attribute.getTensor(getTargetLid(docId), tensor);
+    _target_tensor_attribute.extract_dense_view(getTargetLid(docid), tensor);
+}
+
+const Tensor&
+ImportedTensorAttributeVectorReadGuard::get_tensor_ref(uint32_t docid) const
+{
+    return _target_tensor_attribute.get_tensor_ref(getTargetLid(docid));
 }
 
 vespalib::eval::ValueType
diff --git a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h
index 74c5cd9a685..e873da510d9 100644
--- a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h
+++ b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h
@@ -30,11 +30,14 @@ public:
 
     const ITensorAttribute *asTensorAttribute() const override;
 
-    virtual std::unique_ptr<Tensor> getTensor(uint32_t docId) const override;
-    virtual std::unique_ptr<Tensor> getEmptyTensor() const override;
-    virtual void getTensor(uint32_t docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
-    virtual vespalib::eval::ValueType getTensorType() const override;
-    virtual void get_state(const vespalib::slime::Inserter& inserter) const override;
+    std::unique_ptr<Tensor> getTensor(uint32_t docId) const override;
+    std::unique_ptr<Tensor> getEmptyTensor() const override;
+    void extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const override;
+    const Tensor& get_tensor_ref(uint32_t docid) const override;
+    bool supports_extract_dense_view() const override { return _target_tensor_attribute.supports_extract_dense_view(); }
+    bool supports_get_tensor_ref() const override { return _target_tensor_attribute.supports_get_tensor_ref(); }
+    vespalib::eval::ValueType getTensorType() const override;
+    void get_state(const vespalib::slime::Inserter& inserter) const override;
 };
 
 }
diff --git a/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.cpp
index d4a20abf2fd..3c0e9b38687 100644
--- a/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.cpp
@@ -55,12 +55,6 @@ SerializedTensorAttribute::getTensor(DocId docId) const
     return _serializedTensorStore.getTensor(ref);
 }
 
-void
-SerializedTensorAttribute::getTensor(DocId, vespalib::tensor::MutableDenseTensorView &) const
-{
-    notImplemented();
-}
-
 bool
 SerializedTensorAttribute::onLoad()
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.h
index 5596341a5b7..7b385359337 100644
--- a/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.h
@@ -17,7 +17,6 @@ public:
     virtual ~SerializedTensorAttribute();
     virtual void setTensor(DocId docId, const Tensor &tensor) override;
     virtual std::unique_ptr<Tensor> getTensor(DocId docId) const override;
-    virtual void getTensor(DocId docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
     virtual bool onLoad() override;
     virtual std::unique_ptr<AttributeSaver> onInitSave(vespalib::stringref fileName) override;
     virtual void compactWorst() override;
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
index 6cf4f6d2689..7abc7bd22b5 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
@@ -202,6 +202,22 @@ TensorAttribute::getEmptyTensor() const
     return _emptyTensor->clone();
 }
 
+void
+TensorAttribute::extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const
+{
+    (void) docid;
+    (void) tensor;
+    notImplemented();
+}
+
+const Tensor&
+TensorAttribute::get_tensor_ref(uint32_t docid) const
+{
+    (void) docid;
+    notImplemented();
+    abort(); // Needed to avoid compile error
+}
+
 vespalib::eval::ValueType
 TensorAttribute::getTensorType() const
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
index 8380e485172..9627ea5b4aa 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
@@ -45,6 +45,10 @@ public:
     void onGenerationChange(generation_t generation) override;
     bool addDoc(DocId &docId) override;
     std::unique_ptr<Tensor> getEmptyTensor() const override;
+    void extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const override;
+    const Tensor& get_tensor_ref(uint32_t docid) const override;
+    bool supports_extract_dense_view() const override { return false; }
+    bool supports_get_tensor_ref() const override { return false; }
     vespalib::eval::ValueType getTensorType() const override;
     void get_state(const vespalib::slime::Inserter& inserter) const override;
     void clearDocs(DocId lidLow, DocId lidLimit) override;