diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index 954dd270ee..63adc059f3 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -682,6 +682,17 @@ def ibis_type_from_python_type(t: type) -> ibis_dtypes.DataType: def ibis_type_from_type_kind(tk: bigquery.StandardSqlTypeNames) -> ibis_dtypes.DataType: + """Convert bq type to ibis. Only to be used for remote functions, does not handle all types.""" if tk not in SUPPORTED_IO_BIGQUERY_TYPEKINDS: raise UnsupportedTypeError(tk, SUPPORTED_IO_BIGQUERY_TYPEKINDS) return third_party_ibis_bqtypes.BigQueryType.to_ibis(tk) + + +def bf_type_from_type_kind(bf_schema) -> Dict[str, Dtype]: + """Converts bigquery sql type to the default bigframes dtype.""" + ibis_schema: ibis.Schema = third_party_ibis_bqtypes.BigQuerySchema.to_ibis( + bf_schema + ) + return { + name: ibis_dtype_to_bigframes_dtype(type) for name, type in ibis_schema.items() + } diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 74a8325dac..479b3a7bac 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -1889,8 +1889,10 @@ def _get_table_size(self, destination_table): def _rows_to_dataframe( self, row_iterator: bigquery.table.RowIterator, dtypes: Dict ) -> 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) arrow_table = row_iterator.to_arrow() - return bigframes.session._io.pandas.arrow_to_pandas(arrow_table, dtypes) + return bigframes.session._io.pandas.arrow_to_pandas(arrow_table, dtypes_from_bq) def _start_generic_job(self, job: formatting_helpers.GenericJob): if bigframes.options.display.progress_bar is not None: