use direct access to tensor when possible

1 files changed, 6 insertions, 0 deletions

diff --git a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp
index 51727846f95..27882a6cc8e 100644
--- a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp
+++ b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp
@@ -2,6 +2,7 @@
 
 #include "tensor_attribute_executor.h"
 #include <vespa/searchlib/tensor/i_tensor_attribute.h>
+#include <vespa/searchlib/tensor/direct_tensor_attribute.h>
 
 namespace search {
 namespace features {
@@ -17,6 +18,11 @@ TensorAttributeExecutor(const search::tensor::ITensorAttribute *attribute)
 void
 TensorAttributeExecutor::execute(uint32_t docId)
 {
+    using DTA = search::tensor::DirectTensorAttribute;
+    if (auto direct = dynamic_cast<const DTA *>(_attribute)) {
+        outputs().set_object(0, direct->get_tensor_ref(docId));
+        return;
+    }
     _tensor = _attribute->getTensor(docId);
     if (_tensor) {
         outputs().set_object(0, *_tensor);