|
23 | 23 | import threading |
24 | 24 | from typing import Any, Iterator, Optional |
25 | 25 | import uuid |
| 26 | +import warnings |
26 | 27 |
|
27 | 28 | import pandas as pd |
28 | 29 |
|
@@ -111,23 +112,32 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame): |
111 | 112 | self.page_size = initial_page_size |
112 | 113 | self.max_columns = initial_max_columns |
113 | 114 |
|
114 | | - # TODO(b/469861913): Nested columns from structs (e.g., 'struct_col.name') are not currently sortable. |
115 | | - # TODO(b/463754889): Support non-string column labels for sorting. |
116 | | - if all(isinstance(col, str) for col in dataframe.columns): |
117 | | - self.orderable_columns = [ |
118 | | - str(col_name) |
119 | | - for col_name, dtype in dataframe.dtypes.items() |
120 | | - if dtypes.is_orderable(dtype) |
121 | | - ] |
122 | | - else: |
123 | | - self.orderable_columns = [] |
| 115 | + self.orderable_columns = self._get_orderable_columns(dataframe) |
124 | 116 |
|
125 | 117 | self._initial_load() |
126 | 118 |
|
127 | 119 | # Signals to the frontend that the initial data load is complete. |
128 | 120 | # Also used as a guard to prevent observers from firing during initialization. |
129 | 121 | self._initial_load_complete = True |
130 | 122 |
|
| 123 | + def _get_orderable_columns( |
| 124 | + self, dataframe: bigframes.dataframe.DataFrame |
| 125 | + ) -> list[str]: |
| 126 | + """Determine which columns can be used for client-side sorting.""" |
| 127 | + # TODO(b/469861913): Nested columns from structs (e.g., 'struct_col.name') are not currently sortable. |
| 128 | + # TODO(b/463754889): Support non-string column labels for sorting. |
| 129 | + if not all(isinstance(col, str) for col in dataframe.columns): |
| 130 | + return [] |
| 131 | + |
| 132 | + with warnings.catch_warnings(): |
| 133 | + warnings.simplefilter("ignore", bigframes.exceptions.JSONDtypeWarning) |
| 134 | + warnings.simplefilter("ignore", category=FutureWarning) |
| 135 | + return [ |
| 136 | + str(col_name) |
| 137 | + for col_name, dtype in dataframe.dtypes.items() |
| 138 | + if dtypes.is_orderable(dtype) |
| 139 | + ] |
| 140 | + |
131 | 141 | def _initial_load(self) -> None: |
132 | 142 | """Get initial data and row count.""" |
133 | 143 | # obtain the row counts |
|
0 commit comments