Update to get_unique_metadata for columnar join (LLNL#35)

* Update to get_unique_metadata for columnar join * small fixes --------- Co-authored-by: Stephanie Brink <[email protected]>
dyokelson · Jun 3, 2023 · e1bb009 · e1bb009
1 parent 2126fbe
commit e1bb009
Showing 1 changed file with 25 additions and 9 deletions.
diff --git a/thicket/thicket.py b/thicket/thicket.py
@@ -1215,19 +1215,35 @@ def get_unique_metadata(self):
 
         Returns:
             (dict): alphabetical ordered dictionary with key's being the column names
-                      and the values being unique values for a column
+                and the values being unique values for a metadata column.
         """
         unique_meta = {}
 
-        for col in self.metadata.columns:
-            # skip columns where the values are a list
-            if isinstance(self.metadata[col].iloc[0], list):
-                continue
-            else:
-                unique_entries = self.metadata[col].unique().tolist()
-                unique_meta[col] = unique_entries
+        # thicket object without columnar index
+        if self.dataframe.columns.nlevels == 1:
+            for col in self.metadata.columns:
+                # skip columns where the values are a list
+                if isinstance(self.metadata[col].iloc[0], list):
+                    continue
+                else:
+                    unique_entries = self.metadata[col].unique().tolist()
+                    unique_meta[col] = unique_entries
+
+            sorted_meta = dict(sorted(unique_meta.items(), key=lambda x: x[0].lower()))
+        # columnar joined thicket object
+        else:
+            sorted_meta = []
+            for idx in list(self.metadata.columns.levels[0]):
+                for col in self.metadata[idx].columns:
+                    if isinstance(self.metadata[idx][col].iloc[0], list):
+                        continue
+                    else:
+                        unique_entries = self.metadata[idx][col].unique().tolist()
+                        unique_meta[col] = unique_entries
 
-        sorted_meta = dict(sorted(unique_meta.items(), key=lambda x: x[0].lower()))
+                sorted_meta.append(
+                    (idx, dict(sorted(unique_meta.items(), key=lambda x: x[0].lower())))
+                )
 
         return sorted_meta