Add flag to db to dataframe method to capture extra data (facebookresearch#720)

Summary:
Pull Request resolved: facebookresearch#720

Experiment to dataframe utility by default ignores extra info. Add an extra argument to capture that data in the dataframe.

Reviewed By: adellari

Differential Revision: D72405413

fbshipit-source-id: 71f755c11694afa2922e2b005ff85fd9c1692c32

Author: JasonKChow

aepsych/database/db.py

@@ -564,9 +564,11 @@ def get_stimuli_per_trial(master_id):
 
         return pd.DataFrame(exp_dict)
 
-    def get_data_frame(self) -> pd.DataFrame:
+    def get_data_frame(self, include_extra_data: bool = False) -> pd.DataFrame:
         """Converts parameter and outcome data in the database into a pandas dataframe.
-
+        Args:
+            include_extra_data (bool): Whether to include columns for extra data from
+                the raw table. Defaults to False.
         Returns:
             pandas.Dataframe: The dataframe containing the parameter and outcome data.
         """
@@ -587,6 +589,15 @@ def get_data_frame(self) -> pd.DataFrame:
                 row.update({par.param_name: par.param_value for par in pars})
                 row.update({out.outcome_name: out.outcome_value for out in outs})
 
+                extra_data = pars[0].parent.extra_data
+                if include_extra_data and extra_data is not None:
+                    if isinstance(extra_data, str):
+                        extra_data_dict = json.loads(extra_data)
+                    elif isinstance(extra_data, dict):
+                        extra_data_dict = extra_data
+
+                    row.update(extra_data_dict)
+
                 rows.append(row)
 
             df = pd.DataFrame(rows)

tests/test_db.py

@@ -17,6 +17,7 @@
 import aepsych.config as configuration
 import aepsych.database.db as db
 import aepsych.database.tables as tables
+import numpy as np
 import pandas as pd
 import sqlalchemy
 
@@ -631,6 +632,31 @@ def test_get_dataframe(self):
             n += len(self.data.get_raw_for(rec.unique_id))
         self.assertEqual(n, len(df))
 
+    def test_get_dataframe_extra_data(self):
+        current_path = Path(os.path.abspath(__file__)).parent
+        db_path = current_path
+        db_path = db_path.joinpath("test_databases/extra_info.db")
+
+        # Make a copy of the database
+        dst_db_path = Path("./{}.db".format(str(uuid.uuid4().hex)))
+        shutil.copy(db_path, dst_db_path)
+
+        time.sleep(0.1)
+        self.assertTrue(dst_db_path.is_file())
+
+        self.data = db.Database(dst_db_path)
+
+        df = self.data.get_data_frame(include_extra_data=True)
+
+        # extra_info db is ragged, so check that we have the right columns
+        self.assertIn("extra", df.columns)
+        self.assertIn("additional", df.columns)
+        self.assertIn("trial_number", df.columns)
+
+        # First row is missing the extra data, so check that it is missing
+        np.testing.assert_equal(df.iloc[0]["trial_number"], np.nan)
+        self.assertEqual(df.iloc[1]["trial_number"], 2)
+
 
 if __name__ == "__main__":
     unittest.main()