pythonpandasdataframeinformation-retrieval

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


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.


Solution

  • 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']