Skip to content
16 changes: 13 additions & 3 deletions bigframes/display/anywidget.py
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ def _cached_data(self) -> pd.DataFrame:
"""Combine all cached batches into a single DataFrame."""
if not self._cached_batches:
return pd.DataFrame(columns=self._dataframe.columns)
return pd.concat(self._cached_batches, ignore_index=True)
return pd.concat(self._cached_batches)

def _reset_batch_cache(self) -> None:
"""Resets batch caching attributes."""
Expand Down Expand Up @@ -294,8 +294,18 @@ def _set_table_html(self) -> None:
break

# Get the data for the current page
page_data = cached_data.iloc[start:end]

page_data = cached_data.iloc[start:end].copy()

# Handle index display
# TODO(b/438181139): Add tests for custom multiindex
if self._dataframe._block.has_index:
index_name = page_data.index.name
page_data.insert(
0, index_name if index_name is not None else "", page_data.index
)
else:
# Default index - include as "Row" column
page_data.insert(0, "Row", range(start + 1, start + len(page_data) + 1))
# Handle case where user navigated beyond available data with unknown row count
is_unknown_count = self.row_count is None
is_beyond_data = self._all_data_loaded and len(page_data) == 0 and self.page > 0
Expand Down
Loading