diff --git a/thicket/helpers.py b/thicket/helpers.py index 15faf63e..c80e684d 100644 --- 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 index c1380cde..dcb4ba2d 100644 --- 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