Adding frequency to query set queries.

o19s · Nov 26, 2024 · f471e32 · f471e32
1 parent 7a10c6d
commit f471e32
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 7 deletions.
diff --git a/...ty-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AbstractQuerySampler.java b/...ty-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AbstractQuerySampler.java
@@ -44,7 +44,7 @@ public abstract class AbstractQuerySampler {
     /**
      * Index the query set.
      */
-    protected String indexQuerySet(final NodeClient client, final String name, final String description, final String sampling, Collection<String> queries) throws Exception {
+    protected String indexQuerySet(final NodeClient client, final String name, final String description, final String sampling, Map<String, Long> queries) throws Exception {
 
         LOGGER.info("Indexing {} queries for query set {}", queries.size(), name);
 

diff --git a/...-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AllQueriesQuerySampler.java b/...-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AllQueriesQuerySampler.java
@@ -16,9 +16,8 @@
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.builder.SearchSourceBuilder;
 
-import java.util.HashSet;
+import java.util.HashMap;
 import java.util.Map;
-import java.util.Set;
 
 /**
  * An implementation of {@link AbstractQuerySampler} that uses all UBI queries without any sampling.
@@ -61,10 +60,10 @@ public String sample() throws Exception {
 
         // LOGGER.info("Found {} user queries from the ubi_queries index.", searchResponse.getHits().getTotalHits().toString());
 
-        final Set<String> queries = new HashSet<>();
+        final Map<String, Long> queries = new HashMap<>();
         for(final SearchHit hit : searchResponse.getHits().getHits()) {
             final Map<String, Object> fields = hit.getSourceAsMap();
-            queries.add(fields.get("user_query").toString());
+            queries.merge(fields.get("user_query").toString(), 1L, Long::sum);
         }
 
         // LOGGER.info("Found {} user queries from the ubi_queries index.", queries.size());

diff --git a/.../java/org/opensearch/eval/samplers/ProbabilityProportionalToSizeAbstractQuerySampler.java b/.../java/org/opensearch/eval/samplers/ProbabilityProportionalToSizeAbstractQuerySampler.java
@@ -127,7 +127,7 @@ public String sample() throws Exception {
             LOGGER.info("Summed normalized weights sum to {}", sumOfNormalizedWeights);
         }
 
-        final Set<String> querySet = new HashSet<>();
+        final Map<String, Long> querySet = new HashMap<>();
         final Set<Double> randomNumbers = new HashSet<>();
 
         // Generate random numbers between 0 and 1 for the size of the query set.
@@ -155,7 +155,7 @@ public String sample() throws Exception {
                 }
             }
 
-            querySet.add(closestQuery);
+            querySet.put(closestQuery, weights.get(closestQuery));
             count++;
 
             // LOGGER.info("Generated random value: {}; Smallest delta = {}; Closest query = {}", random, smallestDelta, closestQuery);