Skip to content

Commit c5cf5ee

Browse files
authored
Merge pull request #380 from lincc-frameworks/nested_columns_speedup
speedup nested_columns logic
2 parents 3d3b6e0 + 8170e86 commit c5cf5ee

File tree

1 file changed

+4
-6
lines changed
  • src/nested_pandas/nestedframe

1 file changed

+4
-6
lines changed

src/nested_pandas/nestedframe/core.py

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -97,16 +97,14 @@ def all_columns(self) -> dict:
9797
@property
9898
def nested_columns(self) -> list:
9999
"""retrieves the base column names for all nested dataframes"""
100-
nest_cols = []
101-
for column in self.columns:
102-
if isinstance(self.dtypes[column], NestedDtype):
103-
nest_cols.append(column)
104-
return nest_cols
100+
nested_mask = self.dtypes.apply(lambda dtype: isinstance(dtype, NestedDtype))
101+
return self.columns[nested_mask].tolist()
105102

106103
@property
107104
def base_columns(self) -> list[str]:
108105
"""Returns the list of base (non-nested) column names"""
109-
return [col for col in self.columns if col not in self.nested_columns]
106+
nested_mask = self.dtypes.apply(lambda dtype: not isinstance(dtype, NestedDtype))
107+
return self.columns[nested_mask].tolist()
110108

111109
def _repr_html_(self) -> str | None:
112110
"""Override html representation"""

0 commit comments

Comments
 (0)