IR: How do you match documents based on index values and return the document?

Question

I have a pandas data frame df which has the top 10 documents from a corpus ranked based on their BM25 score, and indexed by their Doc_ID.

Doc_ID	Rank	BM25 Score
1234	1	3.3472
5678	2	3.3238

I also have a list documents containing all of the documents paired up with their Doc_ID, such that the list is in the following form: [['First Doc_ID', 'First doc text], ['Second Doc_ID', 'Second doc text], ...].

I need to take the Doc_ID for the top 3 ranked documents in df, and match each one with the corresponding Doc_ID in documents and print out the document text. I know how to get the Doc_ID for a particular rank from df by doing df.index[df['Rank'] == 1][0], but I'm unsure how to go from there to get the corresponding document text.

Answer 1

You can convert your list to DataFrame and merge:

documents = [[1234, 'First doc text'],
             [5678, 'Second doc text'],
             [5679, 'Third doc text'],
             [5680, 'Fourth doc text']]

(df[df['Rank'].le(3)]
 .merge(pd.DataFrame(documents,
                     columns=['Doc_ID', 'Text']),
        on='Doc_ID')
)

output:

   Doc_ID  Rank  BM25 Score             Text
0    1234     1      3.3472   First doc text
1    5678     2      3.3238  Second doc text
2    5679     3      3.2000   Third doc text

used input:

   Doc_ID  Rank  BM25 Score
0    1234     1      3.3472
1    5678     2      3.3238
2    5679     3      3.2000
3    5680     4      3.1000

Alternatively, if you want a list using python:

top3 = set(df.loc[df['Rank'].le(3), 'Doc_ID'])
out = [text for ID, text in documents if ID in top3]

output: ['First doc text', 'Second doc text', 'Third doc text']