From 188dc341c1a158271f683d8191affa0ca27c2612 Mon Sep 17 00:00:00 2001
From: Andrey Fedorov <andrey.fedorov@gmail.com>
Date: Thu, 10 Oct 2024 09:24:50 -0400
Subject: [PATCH] ENH: add helper function to load clinical tables

---
 idc_index/index.py | 23 +++++++++++++++++++++++
 tests/idcindex.py  |  3 +++
 2 files changed, 26 insertions(+)

diff --git a/idc_index/index.py b/idc_index/index.py
index f157188..a5e69c3 100644
--- a/idc_index/index.py
+++ b/idc_index/index.py
@@ -362,6 +362,29 @@ def fetch_index(self, index_name) -> None:
                     self.clinical_data_dir,
                 )
 
+    def get_clinical_table(self, table_name):
+        """
+        Returns the requested clinical table as a pandas DataFrame.
+
+        Args:
+            table_name (str): Name of the clinical table to be fetched.
+
+        Returns:
+            pandas.DataFrame: The requested clinical table.
+        """
+        if self.clinical_data_dir is None:
+            logger.error(
+                "Clinical data directory is not available. Please fetch clinical_index first."
+            )
+            return None
+
+        table_path = os.path.join(self.clinical_data_dir, f"{table_name}.parquet")
+        if not os.path.exists(table_path):
+            logger.error(f"Table {table_name} is not available in clinical data.")
+            return None
+
+        return pd.read_parquet(table_path)
+
     def get_collections(self):
         """
         Returns the collections present in IDC
diff --git a/tests/idcindex.py b/tests/idcindex.py
index 7acb70a..0aba6be 100644
--- a/tests/idcindex.py
+++ b/tests/idcindex.py
@@ -563,6 +563,9 @@ def test_clinical_index_install(self):
         assert i.indices_overview["clinical_index"]["installed"] is True
         assert len(os.listdir(i.clinical_data_dir)) > 0
 
+        nlst_clinical = i.get_clinical_table("nlst_clinical")
+        assert nlst_clinical is not None
+
 
 if __name__ == "__main__":
     unittest.main()