enable get_perf_columns to work with columnar joined thicket (LLNL#37)

dyokelson · Jun 3, 2023 · 6cdcf2e · 6cdcf2e
1 parent 2428a54
commit 6cdcf2e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/thicket/helpers.py b/thicket/helpers.py
@@ -190,7 +190,12 @@ def _get_perf_columns(df):
 
     numeric_columns = df.select_dtypes(include=numeric_types).columns.tolist()
 
-    if "nid" in numeric_columns:
-        numeric_columns.remove("nid")
-
-    return numeric_columns
+    # thicket object without columnar index
+    if df.columns.nlevels == 1:
+        if "nid" in numeric_columns:
+            numeric_columns.remove("nid")
+
+        return numeric_columns
+    # columnar joined thicket object
+    else:
+        return [x for x in numeric_columns if "nid" not in x]
diff --git a/thicket/thicket.py b/thicket/thicket.py
@@ -501,6 +501,7 @@ def _create_multiindex_columns(dataframe, upper_idx_name):
                 combined_th.dataframe, multiindex=True
             ),
         )
+        combined_th.performance_cols = helpers._get_perf_columns(combined_th.dataframe)
 
         return combined_th
-Original file line number
+Diff line change
@@ Expand Up / @@ -501,6 +501,7 @@ def _create_multiindex_columns(dataframe, upper_idx_name): @@
                     combined_th.dataframe, multiindex=True
                 ),
             )
+            combined_th.performance_cols = helpers._get_perf_columns(combined_th.dataframe)
             return combined_th
@@ Expand Down @@