Use provided id_field in get_many (#252)

* Use provided id_field * Fix 2
allenai · Dec 14, 2023 · c7e61da · c7e61da
1 parent 81d7dd2
commit c7e61da
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ir_datasets/indices/base.py b/ir_datasets/indices/base.py
@@ -17,9 +17,9 @@ def get_many(self, doc_ids, field=None):
         field_idx = self._doc_cls._fields.index(field) if field is not None else None
         for doc in self.get_many_iter(doc_ids):
             if field is not None:
-                result[doc.doc_id] = doc[field_idx]
+                result[getattr(doc, self._id_field)] = doc[field_idx]
             else:
-                result[doc.doc_id] = doc
+                result[getattr(doc, self._id_field)] = doc
         return result
 
     def get_many_iter(self, doc_ids):