Merge pull request #32998 from vespa-engine/toregge/implement-get-enu…

…m-for-reference-attribute Implement getEnum for reference attribute.
vespa-engine · Dec 5, 2024 · 5dbc3c6 · 5dbc3c6
2 parents de61821 + 8f2829c
commit 5dbc3c6
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 1 deletion.
diff --git a/searchlib/src/tests/attribute/reference_attribute/reference_attribute_test.cpp b/searchlib/src/tests/attribute/reference_attribute/reference_attribute_test.cpp
@@ -123,6 +123,8 @@ struct ReferenceAttributeTest : public ::testing::Test {
         return _attr->getReference(doc);
     }
 
+    auto get_enum(uint32_t doc) { return _attr->getEnum(doc); }
+
     void set(uint32_t doc, const GlobalId &gid) {
         _attr->update(doc, gid);
     }
@@ -467,6 +469,23 @@ TEST_F(ReferenceAttributeTest, unique_gids_are_tracked)
     EXPECT_EQ(0u, getUniqueGids());
 }
 
+TEST_F(ReferenceAttributeTest, getEnum_returns_same_value_for_same_reference)
+{
+    ensureDocIdLimit(7);
+    set(1, toGid(doc1));
+    set(2, toGid(doc2));
+    set(4, toGid(doc1));
+    set(5, toGid(doc2));
+    commit();
+    EXPECT_EQ(get_enum(1), get_enum(4));
+    EXPECT_EQ(get_enum(2), get_enum(5));
+    EXPECT_EQ(get_enum(3), get_enum(6));
+    EXPECT_EQ(get_enum(3), get_enum(7));
+    EXPECT_NE(get_enum(1), get_enum(2));
+    EXPECT_NE(get_enum(1), get_enum(3));
+    EXPECT_NE(get_enum(2), get_enum(3));
+}
+
 struct ReferenceAttributeSearchTest : public ReferenceAttributeTest {
 
     constexpr static uint32_t doc_id_limit = 6;

diff --git a/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp b/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
@@ -52,7 +52,6 @@ ReferenceAttribute::ReferenceAttribute(const std::string_view baseFileName, cons
       _gidToLidMapperFactory(),
       _referenceMappings(getGenerationHolder(), getCommittedDocIdLimitRef(), get_initial_alloc())
 {
-    //TODO this is not safe without implementing getEnum
     setEnum(true);
 }
 
@@ -298,6 +297,15 @@ ReferenceAttribute::getReference(DocId doc) const
     }
 }
 
+IAttributeVector::EnumHandle
+ReferenceAttribute::getEnum(DocId doc) const
+{
+    if (doc >= getCommittedDocIdLimit()) [[unlikely]] {
+        return EntryRef().ref();
+    }
+    return _indices.acquire_elem_ref(doc).load_acquire().ref();
+}
+
 bool
 ReferenceAttribute::consider_compact_values(const CompactionStrategy &compactionStrategy)
 {

diff --git a/searchlib/src/vespa/searchlib/attribute/reference_attribute.h b/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
@@ -75,6 +75,7 @@ class ReferenceAttribute : public NotImplementedAttribute
     uint32_t clearDoc(DocId doc) override;
     void update(DocId doc, const GlobalId &gid);
     const Reference *getReference(DocId doc) const;
+    EnumHandle getEnum(DocId doc) const override;
     void setGidToLidMapperFactory(std::shared_ptr<IGidToLidMapperFactory> gidToLidMapperFactory);
     std::shared_ptr<IGidToLidMapperFactory> getGidToLidMapperFactory() const { return _gidToLidMapperFactory; }
     TargetLids getTargetLids() const { return _referenceMappings.getTargetLids(); }