optimize views

3 files changed, 104 insertions, 10 deletions

diff --git a/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value.cpp b/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value.cpp
index 6ff1e130468..ba9c78aeb7b 100644
--- a/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value.cpp
+++ b/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value.cpp
@@ -1,20 +1,21 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "sparse_tensor_value.h"
+#include "sparse_tensor_address_builder.h"
 #include "sparse_tensor_address_decoder.h"
 
 #include <vespa/vespalib/stllike/hash_map.hpp>
 #include <vespa/vespalib/stllike/hash_map_equal.hpp>
-//#include <vespa/vespalib/util/array_equal.hpp>
 
 #include <vespa/log/log.h>
 LOG_SETUP(".eval.tensor.sparse.sparse_tensor_value");
 
 namespace vespalib::tensor {
 
-namespace {
-
 using SubspaceMap = SparseTensorValueIndex::SubspaceMap;
+using View = vespalib::eval::Value::Index::View;
+
+namespace {
 
 void copyMap(SubspaceMap &map, const SubspaceMap &map_in, Stash &stash) {
     // copy the exact hashtable structure:
@@ -38,9 +39,9 @@ size_t needed_memory_for(const SubspaceMap &map, ConstArrayRef<T> cells) {
     return needs;
 }
 
-} // namespace <unnamed>
+//-----------------------------------------------------------------------------
 
-class SparseTensorValueView : public vespalib::eval::Value::Index::View
+class SparseTensorValueView : public View
 {
 private:
     const SubspaceMap &map;
@@ -110,7 +111,89 @@ SparseTensorValueView::next_result(const std::vector<vespalib::stringref*> &addr
     return false;
 }
 
-SparseTensorValueIndex::SparseTensorValueIndex() : map() {}
+//-----------------------------------------------------------------------------
+
+class SparseTensorValueLookup : public View
+{
+private:
+    const SubspaceMap &map;
+    SubspaceMap::const_iterator iter;
+public:
+    SparseTensorValueLookup(const SubspaceMap & map_in) : map(map_in), iter(map.end()) {}
+    ~SparseTensorValueLookup();
+    void lookup(const std::vector<const vespalib::stringref*> &addr) override;
+    bool next_result(const std::vector<vespalib::stringref*> &addr_out, size_t &idx_out) override;
+};
+
+SparseTensorValueLookup::~SparseTensorValueLookup() = default;
+
+void
+SparseTensorValueLookup::lookup(const std::vector<const vespalib::stringref*> &addr)
+{
+    SparseTensorAddressBuilder builder;
+    for (const auto & label : addr) {
+        builder.add(*label);
+    }
+    auto ref = builder.getAddressRef();
+    iter = map.find(ref);
+}
+
+bool
+SparseTensorValueLookup::next_result(const std::vector<vespalib::stringref*> &, size_t &idx_out)
+{
+    if (iter != map.end()) {
+        idx_out = iter->second;
+        iter = map.end();
+        return true;
+    }
+    return false;
+}
+
+//-----------------------------------------------------------------------------
+
+class SparseTensorValueAllMappings : public View
+{
+private:
+    const SubspaceMap &map;
+    SubspaceMap::const_iterator iter;
+public:
+    SparseTensorValueAllMappings(const SubspaceMap & map_in) : map(map_in), iter(map.end()) {}
+    ~SparseTensorValueAllMappings();
+    void lookup(const std::vector<const vespalib::stringref*> &addr) override;
+    bool next_result(const std::vector<vespalib::stringref*> &addr_out, size_t &idx_out) override;
+};
+
+SparseTensorValueAllMappings::~SparseTensorValueAllMappings() = default;
+
+void
+SparseTensorValueAllMappings::lookup(const std::vector<const vespalib::stringref*> &)
+{
+    iter = map.begin();
+}
+
+bool
+SparseTensorValueAllMappings::next_result(const std::vector<vespalib::stringref*> &addr_out, size_t &idx_out)
+{
+    if (iter != map.end()) {
+        const auto & ref = iter->first;
+        idx_out = iter->second;
+        ++iter;
+        SparseTensorAddressDecoder decoder(ref);
+        for (const auto ptr : addr_out) {
+            const auto label = decoder.decodeLabel();
+            *ptr = label;
+        }
+        return true;
+    }
+    return false;
+}
+
+} // namespace <unnamed>
+
+//-----------------------------------------------------------------------------
+
+SparseTensorValueIndex::SparseTensorValueIndex(size_t num_mapped_in)
+    : map(), num_mapped_dims(num_mapped_in) {}
 
 SparseTensorValueIndex::~SparseTensorValueIndex() = default;
 
@@ -118,16 +201,24 @@ size_t SparseTensorValueIndex::size() const {
     return map.size();
 }
 
-std::unique_ptr<vespalib::eval::Value::Index::View>
+std::unique_ptr<View>
 SparseTensorValueIndex::create_view(const std::vector<size_t> &dims) const
 {
+    if (dims.size() == num_mapped_dims) {
+        return std::make_unique<SparseTensorValueLookup>(map);
+    }
+    if (dims.size() == 0) {
+        return std::make_unique<SparseTensorValueAllMappings>(map);
+    }
     return std::make_unique<SparseTensorValueView>(map, dims);
 }
 
+//-----------------------------------------------------------------------------
+
 template<typename T>
 SparseTensorValue<T>::SparseTensorValue(const eval::ValueType &type_in, const SparseTensorValueIndex &index_in, ConstArrayRef<T> cells_in)
     : _type(type_in),
-      _index(),
+      _index(index_in.num_mapped_dims),
       _cells(),
       _stash(needed_memory_for(index_in.map, cells_in))
 {
@@ -149,6 +240,8 @@ template<typename T> SparseTensorValue<T>::~SparseTensorValue() = default;
 template class SparseTensorValue<float>;
 template class SparseTensorValue<double>;
 
+//-----------------------------------------------------------------------------
+
 } // namespace
 
 VESPALIB_HASH_MAP_INSTANTIATE(vespalib::tensor::SparseTensorAddressRef, uint32_t);
diff --git a/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value.h b/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value.h
index 1ee73f2dcbf..ad916021bbc 100644
--- a/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value.h
+++ b/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value.h
@@ -17,7 +17,8 @@ struct SparseTensorValueIndex : public vespalib::eval::Value::Index
     using SubspaceMap = hash_map<SparseTensorAddressRef, uint32_t, hash<SparseTensorAddressRef>,
                                  std::equal_to<>, hashtable_base::and_modulator>;
     SubspaceMap map;
-    SparseTensorValueIndex();
+    size_t num_mapped_dims;
+    explicit SparseTensorValueIndex(size_t num_mapped_dims_in);
     ~SparseTensorValueIndex();
     size_t size() const override;
     std::unique_ptr<View> create_view(const std::vector<size_t> &dims) const override;
diff --git a/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value_builder.h b/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value_builder.h
index 1f79c8b2243..37de7bacdb2 100644
--- a/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value_builder.h
+++ b/eval/src/vespa/eval/tensor/sparse/sparse_tensor_value_builder.h
@@ -26,7 +26,7 @@ public:
                              size_t subspace_size_in,
                              size_t expected_subspaces)
       : _type(type),
-        _index(),
+        _index(num_mapped_in),
         _cells(),
         _stash()
     {