diff --git a/distance_batsim_output.py b/distance_batsim_output.py
index b5e324e4b05be9eee72eeff1e4d7a8900e517e80..3f32433b047682d76281387040779b8c1f41bc94 100755
--- a/distance_batsim_output.py
+++ b/distance_batsim_output.py
@@ -21,8 +21,9 @@ def clean_and_select(df):
     select.job_id = select.job_id.astype(str)
     select["job_id"] = select["job_id"].str.split(':', expand=True)[0]
     select.job_id = select.job_id.astype(int)
+    select.sort_values(by="job_id")
 
-    return select.sort_values(by="job_id")
+    return select.set_index("job_id")
 
 
 def open_and_compare(file1, file2):
@@ -35,7 +36,7 @@ def open_and_compare(file1, file2):
     out1 = clean_and_select(out1)
     out2 = clean_and_select(out2)
 
-    if not out1.job_id.equals(out2.job_id):
+    if not out1.index.equals(out2.index):
         raise KeyError(
             f"{file1} and {file2} cannot be compared: they don't have the same job_ids")