diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 25acfe2f42..fef91f88dc 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -438,9 +438,7 @@ def reorder_levels(self, ids: typing.Sequence[str]): def _to_dataframe(self, result) -> pd.DataFrame: """Convert BigQuery data to pandas DataFrame with specific dtypes.""" - dtypes = dict(zip(self.index_columns, self.index.dtypes)) - dtypes.update(zip(self.value_columns, self.dtypes)) - result_dataframe = self.session._rows_to_dataframe(result, dtypes) + result_dataframe = self.session._rows_to_dataframe(result) # Runs strict validations to ensure internal type predictions and ibis are completely in sync # Do not execute these validations outside of testing suite. if "PYTEST_CURRENT_TEST" in os.environ: @@ -2582,7 +2580,6 @@ def to_pandas(self, *, ordered: Optional[bool] = None) -> pd.Index: ) # Project down to only the index column. So the query can be cached to visualize other data. index_columns = list(self._block.index_columns) - dtypes = dict(zip(index_columns, self.dtypes)) expr = self._expr.select_columns(index_columns) results, _ = self.session._execute( expr, @@ -2590,7 +2587,7 @@ def to_pandas(self, *, ordered: Optional[bool] = None) -> pd.Index: if (ordered is not None) else self.session._strictly_ordered, ) - df = expr.session._rows_to_dataframe(results, dtypes) + df = expr.session._rows_to_dataframe(results) df = df.set_index(index_columns) index = df.index index.names = list(self._block._index_labels) # type:ignore diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 02d30cdbd2..0f7953d3d4 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -2038,7 +2038,7 @@ def _get_table_size(self, destination_table): return table.num_bytes def _rows_to_dataframe( - self, row_iterator: bigquery.table.RowIterator, dtypes: Dict + self, row_iterator: bigquery.table.RowIterator ) -> pandas.DataFrame: # Can ignore inferred datatype until dtype emulation breaks 1:1 mapping between BQ types and bigframes types dtypes_from_bq = bigframes.dtypes.bf_type_from_type_kind(row_iterator.schema)