Keep observations from repeated arms in MapData.map_df (#3399)

Summary:
Pull Request resolved: #3399

**Context:** MapData’s df has only the most recent observation for each (arm, metric) pair, so any arm that appears in multiple trials won’t be included. Data.df does allow data from the same arm in different trials. MapData's behavior is surprising, and downstream functions such as `BestPointMixin._get_trace` assume that every trial with data will be present in `map_df`, so something isn't right.

**This diff:** Now keeps the most recent observation for each (trial_index, arm, metric).

Reviewed By: mpolson64

Differential Revision: D69955018

fbshipit-source-id: 608fb2b1332f45fb7b6bd4d085ef6150318adce2

Author: esantorella

ax/core/map_data.py

@@ -95,7 +95,7 @@ class MapData(Data):
     `experiment.attach_data()` (this requires a description to be set.)
     """
 
-    DEDUPLICATE_BY_COLUMNS = ["arm_name", "metric_name"]
+    DEDUPLICATE_BY_COLUMNS = ["trial_index", "arm_name", "metric_name"]
 
     _map_df: pd.DataFrame
     _memo_df: pd.DataFrame | None

ax/core/tests/test_map_data.py

@@ -18,6 +18,15 @@ def setUp(self) -> None:
         super().setUp()
         self.df = pd.DataFrame(
             [
+                {
+                    "arm_name": "0_0",
+                    "epoch": 0,
+                    "mean": 3.0,
+                    "sem": 0.3,
+                    "trial_index": 0,
+                    "metric_name": "a",
+                },
+                # repeated arm 0_0
                 {
                     "arm_name": "0_0",
                     "epoch": 0,
@@ -78,6 +87,10 @@ def setUp(self) -> None:
 
         self.mmd = MapData(df=self.df, map_key_infos=self.map_key_infos)
 
+    def test_df(self) -> None:
+        df = self.mmd.df
+        self.assertEqual(set(df["trial_index"].drop_duplicates()), {0, 1})
+
     def test_map_key_info(self) -> None:
         self.assertEqual(self.map_key_infos, self.mmd.map_key_infos)