Skip to content
This repository was archived by the owner on Mar 6, 2026. It is now read-only.

Commit 2bd1992

Browse files
google-labs-jules[bot]google-labs-jules[bot]
committed
feat: Add dtype parameters to to_geodataframe
This change adds support for `bool_dtype`, `int_dtype`, `float_dtype`, and `string_dtype` parameters to the `to_geodataframe` method in `RowIterator` and `QueryJob`. These parameters allow you to specify the desired pandas dtypes for boolean, integer, float, and string columns when converting BigQuery results to GeoDataFrames. The changes include: - Updating `RowIterator.to_geodataframe` to accept and pass these dtype parameters to the underlying `to_dataframe` method. - Updating `QueryJob.to_geodataframe` to accept and pass these dtype parameters to the underlying `RowIterator.to_geodataframe` method. - Adding unit tests to verify the correct handling of these parameters.
1 parent 4692747 commit 2bd1992

3 files changed

Lines changed: 183 additions & 0 deletions

File tree

‎google/cloud/bigquery/job/query.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2102,6 +2102,10 @@ def to_geodataframe(
21022102
create_bqstorage_client: bool = True,
21032103
max_results: Optional[int] = None,
21042104
geography_column: Optional[str] = None,
2105+
bool_dtype: Union[Any, None] = DefaultPandasDTypes.BOOL_DTYPE,
2106+
int_dtype: Union[Any, None] = DefaultPandasDTypes.INT_DTYPE,
2107+
float_dtype: Union[Any, None] = None,
2108+
string_dtype: Union[Any, None] = None,
21052109
) -> "geopandas.GeoDataFrame":
21062110
"""Return a GeoPandas GeoDataFrame from a QueryJob
21072111
@@ -2152,6 +2156,34 @@ def to_geodataframe(
21522156
identifies which one to use to construct a GeoPandas
21532157
GeoDataFrame. This option can be ommitted if there's
21542158
only one GEOGRAPHY column.
2159+
bool_dtype (Optional[pandas.Series.dtype, None]):
2160+
If set, indicate a pandas ExtensionDtype (e.g. ``pandas.BooleanDtype()``)
2161+
to convert BigQuery Boolean type, instead of relying on the default
2162+
``pandas.BooleanDtype()``. If you explicitly set the value to ``None``,
2163+
then the data type will be ``numpy.dtype("bool")``. BigQuery Boolean
2164+
type can be found at:
2165+
https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#boolean_type
2166+
int_dtype (Optional[pandas.Series.dtype, None]):
2167+
If set, indicate a pandas ExtensionDtype (e.g. ``pandas.Int64Dtype()``)
2168+
to convert BigQuery Integer types, instead of relying on the default
2169+
``pandas.Int64Dtype()``. If you explicitly set the value to ``None``,
2170+
then the data type will be ``numpy.dtype("int64")``. A list of BigQuery
2171+
Integer types can be found at:
2172+
https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#integer_types
2173+
float_dtype (Optional[pandas.Series.dtype, None]):
2174+
If set, indicate a pandas ExtensionDtype (e.g. ``pandas.Float32Dtype()``)
2175+
to convert BigQuery Float type, instead of relying on the default
2176+
``numpy.dtype("float64")``. If you explicitly set the value to ``None``,
2177+
then the data type will be ``numpy.dtype("float64")``. BigQuery Float
2178+
type can be found at:
2179+
https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#floating_point_types
2180+
string_dtype (Optional[pandas.Series.dtype, None]):
2181+
If set, indicate a pandas ExtensionDtype (e.g. ``pandas.StringDtype()``) to
2182+
convert BigQuery String type, instead of relying on the default
2183+
``numpy.dtype("object")``. If you explicitly set the value to ``None``,
2184+
then the data type will be ``numpy.dtype("object")``. BigQuery String
2185+
type can be found at:
2186+
https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#string_type
21552187
21562188
Returns:
21572189
geopandas.GeoDataFrame:
@@ -2175,6 +2207,10 @@ def to_geodataframe(
21752207
progress_bar_type=progress_bar_type,
21762208
create_bqstorage_client=create_bqstorage_client,
21772209
geography_column=geography_column,
2210+
bool_dtype=bool_dtype,
2211+
int_dtype=int_dtype,
2212+
float_dtype=float_dtype,
2213+
string_dtype=string_dtype,
21782214
)
21792215

21802216
def __iter__(self):

‎google/cloud/bigquery/table.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2727,6 +2727,10 @@ def to_geodataframe(
27272727
progress_bar_type: Optional[str] = None,
27282728
create_bqstorage_client: bool = True,
27292729
geography_column: Optional[str] = None,
2730+
bool_dtype: Union[Any, None] = DefaultPandasDTypes.BOOL_DTYPE,
2731+
int_dtype: Union[Any, None] = DefaultPandasDTypes.INT_DTYPE,
2732+
float_dtype: Union[Any, None] = None,
2733+
string_dtype: Union[Any, None] = None,
27302734
) -> "geopandas.GeoDataFrame":
27312735
"""Create a GeoPandas GeoDataFrame by loading all pages of a query.
27322736
@@ -2778,6 +2782,34 @@ def to_geodataframe(
27782782
identifies which one to use to construct a geopandas
27792783
GeoDataFrame. This option can be ommitted if there's
27802784
only one GEOGRAPHY column.
2785+
bool_dtype (Optional[pandas.Series.dtype, None]):
2786+
If set, indicate a pandas ExtensionDtype (e.g. ``pandas.BooleanDtype()``)
2787+
to convert BigQuery Boolean type, instead of relying on the default
2788+
``pandas.BooleanDtype()``. If you explicitly set the value to ``None``,
2789+
then the data type will be ``numpy.dtype("bool")``. BigQuery Boolean
2790+
type can be found at:
2791+
https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#boolean_type
2792+
int_dtype (Optional[pandas.Series.dtype, None]):
2793+
If set, indicate a pandas ExtensionDtype (e.g. ``pandas.Int64Dtype()``)
2794+
to convert BigQuery Integer types, instead of relying on the default
2795+
``pandas.Int64Dtype()``. If you explicitly set the value to ``None``,
2796+
then the data type will be ``numpy.dtype("int64")``. A list of BigQuery
2797+
Integer types can be found at:
2798+
https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#integer_types
2799+
float_dtype (Optional[pandas.Series.dtype, None]):
2800+
If set, indicate a pandas ExtensionDtype (e.g. ``pandas.Float32Dtype()``)
2801+
to convert BigQuery Float type, instead of relying on the default
2802+
``numpy.dtype("float64")``. If you explicitly set the value to ``None``,
2803+
then the data type will be ``numpy.dtype("float64")``. BigQuery Float
2804+
type can be found at:
2805+
https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#floating_point_types
2806+
string_dtype (Optional[pandas.Series.dtype, None]):
2807+
If set, indicate a pandas ExtensionDtype (e.g. ``pandas.StringDtype()``) to
2808+
convert BigQuery String type, instead of relying on the default
2809+
``numpy.dtype("object")``. If you explicitly set the value to ``None``,
2810+
then the data type will be ``numpy.dtype("object")``. BigQuery String
2811+
type can be found at:
2812+
https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#string_type
27812813
27822814
Returns:
27832815
geopandas.GeoDataFrame:
@@ -2829,6 +2861,10 @@ def to_geodataframe(
28292861
progress_bar_type,
28302862
create_bqstorage_client,
28312863
geography_as_object=True,
2864+
bool_dtype=bool_dtype,
2865+
int_dtype=int_dtype,
2866+
float_dtype=float_dtype,
2867+
string_dtype=string_dtype,
28322868
)
28332869

28342870
return geopandas.GeoDataFrame(
@@ -2932,6 +2968,10 @@ def to_geodataframe(
29322968
progress_bar_type=None,
29332969
create_bqstorage_client=True,
29342970
geography_column: Optional[str] = None,
2971+
bool_dtype: Union[Any, None] = DefaultPandasDTypes.BOOL_DTYPE,
2972+
int_dtype: Union[Any, None] = DefaultPandasDTypes.INT_DTYPE,
2973+
float_dtype: Union[Any, None] = None,
2974+
string_dtype: Union[Any, None] = None,
29352975
) -> "pandas.DataFrame":
29362976
"""Create an empty dataframe.
29372977
@@ -2941,6 +2981,10 @@ def to_geodataframe(
29412981
progress_bar_type (Any): Ignored. Added for compatibility with RowIterator.
29422982
create_bqstorage_client (bool): Ignored. Added for compatibility with RowIterator.
29432983
geography_column (str): Ignored. Added for compatibility with RowIterator.
2984+
bool_dtype (Any): Ignored. Added for compatibility with RowIterator.
2985+
int_dtype (Any): Ignored. Added for compatibility with RowIterator.
2986+
float_dtype (Any): Ignored. Added for compatibility with RowIterator.
2987+
string_dtype (Any): Ignored. Added for compatibility with RowIterator.
29442988
29452989
Returns:
29462990
pandas.DataFrame: An empty :class:`~pandas.DataFrame`.

‎tests/unit/test_table_pandas.py‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -261,3 +261,106 @@ def test_to_dataframe_with_jobs_query_response(class_under_test):
261261
"Tiffani",
262262
]
263263
assert list(df["number"]) == [6, 325, 26, 10, 17, 22, 6, 229, 8]
264+
265+
266+
@mock.patch("google.cloud.bigquery.table.geopandas")
267+
def test_rowiterator_to_geodataframe_with_default_dtypes(
268+
mock_geopandas, monkeypatch, class_under_test
269+
):
270+
mock_geopandas.GeoDataFrame = mock.Mock(spec=True)
271+
mock_client = mock.create_autospec(bigquery.Client)
272+
mock_client.project = "test-proj"
273+
mock_api_request = mock.Mock()
274+
schema = [
275+
bigquery.SchemaField("geo_col", "GEOGRAPHY"),
276+
bigquery.SchemaField("bool_col", "BOOLEAN"),
277+
bigquery.SchemaField("int_col", "INTEGER"),
278+
bigquery.SchemaField("float_col", "FLOAT"),
279+
bigquery.SchemaField("string_col", "STRING"),
280+
]
281+
rows = class_under_test(mock_client, mock_api_request, TEST_PATH, schema)
282+
283+
mock_df = pandas.DataFrame(
284+
{
285+
"geo_col": ["POINT (1 2)"],
286+
"bool_col": [True],
287+
"int_col": [123],
288+
"float_col": [1.23],
289+
"string_col": ["abc"],
290+
}
291+
)
292+
rows.to_dataframe = mock.Mock(return_value=mock_df)
293+
294+
rows.to_geodataframe(geography_column="geo_col")
295+
296+
rows.to_dataframe.assert_called_once_with(
297+
None, # bqstorage_client
298+
None, # dtypes
299+
None, # progress_bar_type
300+
True, # create_bqstorage_client
301+
geography_as_object=True,
302+
bool_dtype=bigquery.enums.DefaultPandasDTypes.BOOL_DTYPE,
303+
int_dtype=bigquery.enums.DefaultPandasDTypes.INT_DTYPE,
304+
float_dtype=None,
305+
string_dtype=None,
306+
)
307+
mock_geopandas.GeoDataFrame.assert_called_once_with(
308+
mock_df, crs="EPSG:4326", geometry="geo_col"
309+
)
310+
311+
312+
@mock.patch("google.cloud.bigquery.table.geopandas")
313+
def test_rowiterator_to_geodataframe_with_custom_dtypes(
314+
mock_geopandas, monkeypatch, class_under_test
315+
):
316+
mock_geopandas.GeoDataFrame = mock.Mock(spec=True)
317+
mock_client = mock.create_autospec(bigquery.Client)
318+
mock_client.project = "test-proj"
319+
mock_api_request = mock.Mock()
320+
schema = [
321+
bigquery.SchemaField("geo_col", "GEOGRAPHY"),
322+
bigquery.SchemaField("bool_col", "BOOLEAN"),
323+
bigquery.SchemaField("int_col", "INTEGER"),
324+
bigquery.SchemaField("float_col", "FLOAT"),
325+
bigquery.SchemaField("string_col", "STRING"),
326+
]
327+
rows = class_under_test(mock_client, mock_api_request, TEST_PATH, schema)
328+
329+
mock_df = pandas.DataFrame(
330+
{
331+
"geo_col": ["POINT (3 4)"],
332+
"bool_col": [False],
333+
"int_col": [456],
334+
"float_col": [4.56],
335+
"string_col": ["def"],
336+
}
337+
)
338+
rows.to_dataframe = mock.Mock(return_value=mock_df)
339+
340+
custom_bool_dtype = "bool"
341+
custom_int_dtype = "int32"
342+
custom_float_dtype = "float32"
343+
custom_string_dtype = "string"
344+
345+
rows.to_geodataframe(
346+
geography_column="geo_col",
347+
bool_dtype=custom_bool_dtype,
348+
int_dtype=custom_int_dtype,
349+
float_dtype=custom_float_dtype,
350+
string_dtype=custom_string_dtype,
351+
)
352+
353+
rows.to_dataframe.assert_called_once_with(
354+
None, # bqstorage_client
355+
None, # dtypes
356+
None, # progress_bar_type
357+
True, # create_bqstorage_client
358+
geography_as_object=True,
359+
bool_dtype=custom_bool_dtype,
360+
int_dtype=custom_int_dtype,
361+
float_dtype=custom_float_dtype,
362+
string_dtype=custom_string_dtype,
363+
)
364+
mock_geopandas.GeoDataFrame.assert_called_once_with(
365+
mock_df, crs="EPSG:4326", geometry="geo_col"
366+
)

0 commit comments

Comments
 (0)