restructure trie to store values as a set

Signed-off-by: Ruirui Zhang <mariazrr@amazon.com>

Author: ruai0511

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add temporal routing processors for time-based document routing ([#18920](https://github.com/opensearch-project/OpenSearch/issues/18920))
 - Implement Query Rewriting Infrastructure ([#19060](https://github.com/opensearch-project/OpenSearch/pull/19060))
 - The dynamic mapping parameter supports false_allow_templates ([#19065](https://github.com/opensearch-project/OpenSearch/pull/19065))
+- [Rule-based Auto-tagging] restructure the in-memory trie to store values as a set ([#19344](https://github.com/opensearch-project/OpenSearch/pull/19344))
 - Add a toBuilder method in EngineConfig to support easy modification of configs([#19054](https://github.com/opensearch-project/OpenSearch/pull/19054))
 - Add StoreFactory plugin interface for custom Store implementations([#19091](https://github.com/opensearch-project/OpenSearch/pull/19091))
 - Use S3CrtClient for higher throughput while uploading files to S3 ([#18800](https://github.com/opensearch-project/OpenSearch/pull/18800))

modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/AttributeValueStore.java

@@ -8,7 +8,8 @@
 
 package org.opensearch.rule.storage;
 
-import java.util.Optional;
+import java.util.List;
+import java.util.Set;
 
 /**
  * This interface provides apis to store Rule attribute values
@@ -23,16 +24,17 @@ public interface AttributeValueStore<K, V> {
 
     /**
      * removes the key and associated value from attribute value store
-     * @param key to be removed
+     * @param key key of the value to be removed
+     * @param value to be removed
      */
-    void remove(K key);
+    void remove(K key, V value);
 
     /**
      * Returns the value associated with the key
      * @param key in the data structure
      * @return
      */
-    Optional<V> get(K key);
+    List<Set<V>> get(K key);
 
     /**
      * Clears all the keys and their associated values from the attribute value store

modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/DefaultAttributeValueStore.java

@@ -10,8 +10,10 @@
 
 import org.apache.commons.collections4.trie.PatriciaTrie;
 
-import java.util.Map;
-import java.util.Optional;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 /**
@@ -21,7 +23,7 @@
  * ref: https://commons.apache.org/proper/commons-collections/javadocs/api-4.4/org/apache/commons/collections4/trie/PatriciaTrie.html
  */
 public class DefaultAttributeValueStore<K extends String, V> implements AttributeValueStore<K, V> {
-    private final PatriciaTrie<V> trie;
+    private final PatriciaTrie<Set<V>> trie;
     private static final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
     private static final ReentrantReadWriteLock.ReadLock readLock = lock.readLock();
     private static final ReentrantReadWriteLock.WriteLock writeLock = lock.writeLock();
@@ -37,74 +39,58 @@ public DefaultAttributeValueStore() {
      * Main constructor
      * @param trie A Patricia Trie
      */
-    public DefaultAttributeValueStore(PatriciaTrie<V> trie) {
+    public DefaultAttributeValueStore(PatriciaTrie<Set<V>> trie) {
         this.trie = trie;
     }
 
     @Override
     public void put(K key, V value) {
         writeLock.lock();
         try {
-            trie.put(key, value);
+            trie.computeIfAbsent(key, k -> new HashSet<>()).add(value);
         } finally {
             writeLock.unlock();
         }
     }
 
     @Override
-    public void remove(String key) {
+    public void remove(K key, V value) {
         writeLock.lock();
         try {
-            trie.remove(key);
+            Set<V> values = trie.get(key);
+            if (values != null && values.contains(value)) {
+                values.remove(value);
+                if (values.isEmpty()) {
+                    trie.remove(key);
+                }
+            }
         } finally {
             writeLock.unlock();
         }
     }
 
     @Override
-    public Optional<V> get(String key) {
+    public List<Set<V>> get(String key) {
         readLock.lock();
         try {
-            /**
-             * Since we are inserting prefixes into the trie and searching for larger strings
-             * It is important to find the largest matching prefix key in the trie efficiently
-             * Hence we can do binary search
-             */
-            final String longestMatchingPrefix = findLongestMatchingPrefix(key);
+            List<Set<V>> results = new ArrayList<>();
+            StringBuilder prefixBuilder = new StringBuilder(key);
 
-            /**
-             * Now there are following cases for this prefix
-             * 1. There is a Rule which has this prefix as one of the attribute values. In this case we should return the
-             *     Rule's label otherwise send empty
-             */
-            for (Map.Entry<String, V> possibleMatch : trie.prefixMap(longestMatchingPrefix).entrySet()) {
-                if (key.startsWith(possibleMatch.getKey())) {
-                    return Optional.of(possibleMatch.getValue());
+            for (int i = key.length(); i >= 0; i--) {
+                String prefix = prefixBuilder.toString();
+                Set<V> value = trie.get(prefix);
+                if (value != null && !value.isEmpty()) {
+                    results.add(value);
+                }
+                if (!prefixBuilder.isEmpty()) {
+                    prefixBuilder.deleteCharAt(prefixBuilder.length() - 1);
                 }
             }
+
+            return results;
         } finally {
             readLock.unlock();
         }
-        return Optional.empty();
-    }
-
-    private String findLongestMatchingPrefix(String key) {
-        int low = 0;
-        int high = key.length() - 1;
-
-        while (low < high) {
-            int mid = (high + low + 1) / 2;
-            /**
-             * This operation has O(1) complexity because prefixMap returns only the iterator
-             */
-            if (!trie.prefixMap(key.substring(0, mid)).isEmpty()) {
-                low = mid;
-            } else {
-                high = mid - 1;
-            }
-        }
-
-        return key.substring(0, low);
     }
 
     @Override

modules/autotagging-commons/common/src/test/java/org/opensearch/rule/storage/AttributeValueStoreTests.java

@@ -26,18 +26,24 @@ public void setUp() throws Exception {
 
     public void testPut() {
         subjectUnderTest.put("foo", "bar");
-        assertEquals("bar", subjectUnderTest.get("foo").get());
+        assertEquals("bar", subjectUnderTest.get("foo").getFirst().iterator().next());
+        subjectUnderTest.put("foo", "sing");
+        assertEquals(1, subjectUnderTest.get("foo").size());
+        assertEquals(2, subjectUnderTest.get("foo").get(0).size());
+        assertTrue(subjectUnderTest.get("foo").get(0).contains("sing"));
     }
 
     public void testRemove() {
         subjectUnderTest.put("foo", "bar");
-        subjectUnderTest.remove("foo");
+        subjectUnderTest.remove("foo", "bar");
         assertEquals(0, subjectUnderTest.size());
     }
 
     public void tesGet() {
         subjectUnderTest.put("foo", "bar");
-        assertEquals("bar", subjectUnderTest.get("foo").get());
+        assertEquals("bar", subjectUnderTest.get("foo").getFirst());
+        subjectUnderTest.put("foo", "sing");
+        assertEquals(2, subjectUnderTest.get("foo").size());
     }
 
     public void testGetWhenNoProperPrefixIsPresent() {

modules/autotagging-commons/src/main/java/org/opensearch/rule/InMemoryRuleProcessingService.java

@@ -66,7 +66,7 @@ private void perform(Rule rule, BiConsumer<Map.Entry<Attribute, Set<String>>, Ru
     private void removeOperation(Map.Entry<Attribute, Set<String>> attributeEntry, Rule rule) {
         AttributeValueStore<String, String> valueStore = attributeValueStoreFactory.getAttributeValueStore(attributeEntry.getKey());
         for (String value : attributeEntry.getValue()) {
-            valueStore.remove(value.replace(WILDCARD, ""));
+            valueStore.remove(value.replace(WILDCARD, ""), rule.getFeatureValue());
         }
     }
 
@@ -92,12 +92,13 @@ public Optional<String> evaluateLabel(List<AttributeExtractor<String>> attribute
                 attributeExtractor.getAttribute()
             );
             for (String value : attributeExtractor.extract()) {
-                Optional<String> possibleMatch = valueStore.get(value);
+                List<Set<String>> candidateMatches = valueStore.get(value);
 
-                if (possibleMatch.isEmpty()) {
+                if (candidateMatches == null || candidateMatches.isEmpty()) {
                     return Optional.empty();
                 }
 
+                Optional<String> possibleMatch = candidateMatches.get(0).stream().findAny();
                 if (result.isEmpty()) {
                     result = possibleMatch;
                 } else {