MSMARCO 10m Recall Operations (#1092)

* added a 10m recall test for msmarco

Author: john-wagster

msmarco-v2-vector/operations/default.json

@@ -69,7 +69,12 @@
     "visit-percentage": {{p_search_ops[i][3]}},
   {%- endif -%}
   "oversample-rescore": {{p_search_ops[i][2]}},
-  "include-in-reporting": false
+  "include-in-reporting": false,
+  {%- if initial_indexing_ingest_doc_count is defined and initial_indexing_ingest_doc_count == 10000000 -%}
+    "recall-doc-set": "10m"
+  {%- else -%}
+    "recall-doc-set": "full"
+  {%- endif -%}
 }
 {%- endfor %}
 {%- set p_hybrid_knn_ops = (hybrid_knn_ops | default([(10, 0), (10, 50), (100, 200), (100, 300)]) ) -%}

msmarco-v2-vector/queries-recall-10m.json.bz2

@@ -16,6 +16,7 @@
 
 QUERIES_FILENAME: str = "queries.json.bz2"
 QUERIES_RECALL_FILENAME: str = "queries-recall.json.bz2"
+QUERIES_RECALL_10M_FILENAME: str = "queries-recall-10m.json.bz2"
 
 
 def extract_vector_operations_count(knn_result):
@@ -147,6 +148,7 @@ def params(self):
             "num_candidates": self._params.get("num-candidates", 100),
             "visit_percentage": self._params.get("visit-percentage", -1),
             "oversample_rescore": self._params.get("oversample-rescore", -1),
+            "recall_doc_set": self._params.get("recall-doc-set", -1),
         }
 
 
@@ -157,6 +159,7 @@ async def __call__(self, es, params):
         visit_percentage = params["visit_percentage"]
         index = params["index"]
         request_cache = params["cache"]
+        recall_doc_set = params["recall_doc_set"]
 
         cwd = os.path.dirname(__file__)
         qrels = read_qrels(os.path.join(cwd, "qrels.tsv"))
@@ -166,7 +169,13 @@ async def __call__(self, es, params):
         exact_total = 0
         min_recall = top_k
         nodes_visited = []
-        with bz2.open(os.path.join(cwd, QUERIES_RECALL_FILENAME), "r") as queries_file:
+
+        if recall_doc_set == "10m":
+            queries_recall = QUERIES_RECALL_10M_FILENAME
+        else:
+            queries_recall = QUERIES_RECALL_FILENAME
+
+        with bz2.open(os.path.join(cwd, queries_recall), "r") as queries_file:
             for line in queries_file:
                 query = json.loads(line)
                 query_id = query["query_id"]