From f471e3202c5e12a211b9101f1cc10970877fff82 Mon Sep 17 00:00:00 2001
From: jzonthemtn <jeff.zemerick@mtnfog.com>
Date: Tue, 26 Nov 2024 09:20:19 -0500
Subject: [PATCH] Adding frequency to query set queries.

---
 .../org/opensearch/eval/samplers/AbstractQuerySampler.java | 2 +-
 .../opensearch/eval/samplers/AllQueriesQuerySampler.java   | 7 +++----
 .../ProbabilityProportionalToSizeAbstractQuerySampler.java | 4 ++--
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AbstractQuerySampler.java b/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AbstractQuerySampler.java
index c82de62..828a88c 100644
--- a/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AbstractQuerySampler.java
+++ b/opensearch-search-quality-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/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AllQueriesQuerySampler.java b/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AllQueriesQuerySampler.java
index e488c7d..4f105c0 100644
--- a/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/AllQueriesQuerySampler.java
+++ b/opensearch-search-quality-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/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/ProbabilityProportionalToSizeAbstractQuerySampler.java b/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/ProbabilityProportionalToSizeAbstractQuerySampler.java
index 1e21e67..1d77e1b 100644
--- a/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/samplers/ProbabilityProportionalToSizeAbstractQuerySampler.java
+++ b/opensearch-search-quality-evaluation-plugin/src/main/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);