Merge branch 'main' into shuowei-time-series-bike

Author: shuoweil

bigframes/display/anywidget.py

@@ -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."""
@@ -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

bigframes/series.py

@@ -2303,7 +2303,7 @@ def to_dict(
         *,
         allow_large_results: Optional[bool] = None,
     ) -> typing.Mapping:
-        return typing.cast(dict, self.to_pandas(allow_large_results=allow_large_results).to_dict(into))  # type: ignore
+        return typing.cast(dict, self.to_pandas(allow_large_results=allow_large_results).to_dict(into=into))  # type: ignore
 
     def to_excel(
         self, excel_writer, sheet_name="Sheet1", *, allow_large_results=None, **kwargs

bigframes/testing/utils.py

@@ -14,6 +14,7 @@
 
 import base64
 import decimal
+import re
 from typing import Iterable, Optional, Sequence, Set, Union
 
 import geopandas as gpd  # type: ignore
@@ -69,6 +70,12 @@
 ]
 
 
+def pandas_major_version() -> int:
+    match = re.search(r"^v?(\d+)", pd.__version__.strip())
+    assert match is not None
+    return int(match.group(1))
+
+
 # Prefer this function for tests that run in both ordered and unordered mode
 def assert_dfs_equivalent(pd_df: pd.DataFrame, bf_df: bpd.DataFrame, **kwargs):
     bf_df_local = bf_df.to_pandas()
@@ -83,7 +90,7 @@ def assert_series_equivalent(pd_series: pd.Series, bf_series: bpd.Series, **kwar
 
 
 def _normalize_all_nulls(col: pd.Series) -> pd.Series:
-    if col.dtype == bigframes.dtypes.FLOAT_DTYPE:
+    if col.dtype in (bigframes.dtypes.FLOAT_DTYPE, bigframes.dtypes.INT_DTYPE):
         col = col.astype("float64")
     if pd_types.is_object_dtype(col):
         col = col.fillna(float("nan"))
@@ -134,6 +141,15 @@ def assert_series_equal(
             left = left.sort_index()
             right = right.sort_index()
 
+    if isinstance(left.index, pd.RangeIndex) or pd_types.is_integer_dtype(
+        left.index.dtype,
+    ):
+        left.index = left.index.astype("Int64")
+    if isinstance(right.index, pd.RangeIndex) or pd_types.is_integer_dtype(
+        right.index.dtype,
+    ):
+        right.index = right.index.astype("Int64")
+
     if nulls_are_nan:
         left = _normalize_all_nulls(left)
         right = _normalize_all_nulls(right)