diff --git a/ax/core/data.py b/ax/core/data.py
index 7ec309442a4..2939143ff79 100644
--- a/ax/core/data.py
+++ b/ax/core/data.py
@@ -357,7 +357,7 @@ def df(self) -> pd.DataFrame:
             # In the case where all MAP_KEY values are NaN for a group we return an
             # arbitrary row from that group.
             .fillna(
-                self.full_df.groupby(self.DEDUPLICATE_BY_COLUMNS).apply(
+                self.full_df.groupby(self.DEDUPLICATE_BY_COLUMNS)[MAP_KEY].apply(
                     lambda group: group.index[0]
                 )
             )