FIX: keep gt_entity_id as int64 when no match candidates are added

ing-bank · Sep 15, 2024 · 9147749 · 9147749
1 parent 90b9885
commit 9147749
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/emm/indexing/pandas_candidate_selection.py b/emm/indexing/pandas_candidate_selection.py
@@ -212,6 +212,7 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
             if self.with_no_matches:
                 # change gt_uid column to nullable integer
                 candidates["gt_uid"] = candidates["gt_uid"].replace({NO_MATCH_ID: np.nan}).astype("Int64")
+                candidates["gt_entity_id"] = candidates["gt_entity_id"].replace({NO_MATCH_ID: np.nan}).astype("Int64")
 
             timer.log_param("n", len(X))