Skip to content

Commit ebfd0a8

Browse files
chalmerlowegoogle-labs-jules[bot]gcf-owl-bot[bot]tswast
authored
feat: Add dtype parameters to to_geodataframe functions (#2176)
* 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. * updates to several tests re geopandas as well as imports * updates to enum import * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Update pyproject.toml Co-authored-by: Tim Sweña (Swast) <[email protected]> * Update testing/constraints-3.9.txt Co-authored-by: Tim Sweña (Swast) <[email protected]> --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com> Co-authored-by: Tim Sweña (Swast) <[email protected]>
1 parent 0217637 commit ebfd0a8

File tree

9 files changed

+205
-10
lines changed

9 files changed

+205
-10
lines changed

‎docs/conf.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@
6161

6262
# autodoc/autosummary flags
6363
autoclass_content = "both"
64-
autodoc_default_options = {"members": True, "inherited-members": True}
64+
autodoc_default_options = {"members": True}
6565
autosummary_generate = True
6666

6767

@@ -109,7 +109,6 @@
109109
# List of patterns, relative to source directory, that match files and
110110
# directories to ignore when looking for source files.
111111
exclude_patterns = [
112-
"google/cloud/bigquery_v2/**", # Legacy proto-based types.
113112
"_build",
114113
"**/.nox/**/*",
115114
"samples/AUTHORING_GUIDE.md",

‎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`.

‎noxfile.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -109,9 +109,7 @@ def default(session, install_extras=True):
109109
# that logic (and the associated tests) we avoid installing the [ipython] extra
110110
# which has a downstream effect of then avoiding installing bigquery_magics.
111111
if install_extras and session.python == UNIT_TEST_PYTHON_VERSIONS[0]:
112-
install_target = (
113-
".[bqstorage,pandas,ipywidgets,geopandas,tqdm,opentelemetry,bigquery_v2]"
114-
)
112+
install_target = ".[bqstorage,pandas,ipywidgets,geopandas,matplotlib,tqdm,opentelemetry,bigquery_v2]"
115113
elif install_extras: # run against all other UNIT_TEST_PYTHON_VERSIONS
116114
install_target = ".[all]"
117115
else:

‎pyproject.toml‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,11 @@ pandas = [
8282
ipywidgets = ["ipywidgets >= 7.7.1", "ipykernel >= 6.2.0"]
8383
geopandas = ["geopandas >= 0.9.0, < 2.0.0", "Shapely >= 1.8.4, < 3.0.0"]
8484
ipython = ["ipython >= 7.23.1", "bigquery-magics >= 0.6.0"]
85-
tqdm = ["tqdm >= 4.7.4, < 5.0.0"]
85+
matplotlib = [
86+
"matplotlib >= 3.7.1, <= 3.9.2; python_version == '3.9'",
87+
"matplotlib >= 3.10.3; python_version >= '3.10'",
88+
]
89+
tqdm = ["tqdm >= 4.23.4, < 5.0.0"]
8690
opentelemetry = [
8791
"opentelemetry-api >= 1.1.0",
8892
"opentelemetry-sdk >= 1.1.0",
@@ -93,7 +97,7 @@ bigquery_v2 = [
9397
"protobuf >= 3.20.2, < 7.0.0, != 4.21.0, != 4.21.1, != 4.21.2, != 4.21.3, != 4.21.4, != 4.21.5", # For the legacy proto-based types.
9498
]
9599
all = [
96-
"google-cloud-bigquery[bqstorage,pandas,ipywidgets,geopandas,ipython,tqdm,opentelemetry,bigquery_v2]",
100+
"google-cloud-bigquery[bqstorage,pandas,ipywidgets,geopandas,ipython,matplotlib,tqdm,opentelemetry,bigquery_v2]",
97101
]
98102

99103
[tool.setuptools.dynamic]

‎testing/constraints-3.9.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,4 +29,4 @@ pyarrow==4.0.0
2929
python-dateutil==2.8.2
3030
requests==2.21.0
3131
Shapely==1.8.4
32-
tqdm==4.7.4
32+
matplotlib==3.7.1

‎tests/unit/job/test_query_pandas.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from ..helpers import make_connection
2323
from .helpers import _make_client
2424
from .helpers import _make_job_resource
25+
from google.cloud.bigquery.enums import DefaultPandasDTypes
2526

2627
try:
2728
from google.cloud import bigquery_storage
@@ -30,6 +31,7 @@
3031
except (ImportError, AttributeError):
3132
bigquery_storage = None
3233

34+
3335
try:
3436
import shapely
3537
except (ImportError, AttributeError):
@@ -1019,5 +1021,9 @@ def test_query_job_to_geodataframe_delegation(wait_for_query):
10191021
progress_bar_type=progress_bar_type,
10201022
create_bqstorage_client=create_bqstorage_client,
10211023
geography_column=geography_column,
1024+
bool_dtype=DefaultPandasDTypes.BOOL_DTYPE,
1025+
int_dtype=DefaultPandasDTypes.INT_DTYPE,
1026+
float_dtype=None,
1027+
string_dtype=None,
10221028
)
10231029
assert df is row_iterator.to_geodataframe.return_value

‎tests/unit/test_table.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
from google.cloud.bigquery import exceptions
3232
from google.cloud.bigquery import external_config
3333
from google.cloud.bigquery import schema
34+
from google.cloud.bigquery.enums import DefaultPandasDTypes
3435
from google.cloud.bigquery.table import TableReference
3536
from google.cloud.bigquery.dataset import DatasetReference
3637

@@ -4065,7 +4066,7 @@ def test_to_dataframe_no_tqdm(self):
40654066

40664067
def test_to_dataframe_tqdm_error(self):
40674068
pytest.importorskip("pandas")
4068-
pytest.importorskip("tqdm")
4069+
tqdm = pytest.importorskip("tqdm")
40694070
mock.patch("tqdm.tqdm_gui", new=None)
40704071
mock.patch("tqdm.notebook.tqdm", new=None)
40714072
mock.patch("tqdm.tqdm", new=None)
@@ -4100,7 +4101,7 @@ def test_to_dataframe_tqdm_error(self):
41004101
for warning in warned: # pragma: NO COVER
41014102
self.assertIn(
41024103
warning.category,
4103-
[UserWarning, DeprecationWarning],
4104+
[UserWarning, DeprecationWarning, tqdm.TqdmExperimentalWarning],
41044105
)
41054106

41064107
def test_to_dataframe_w_empty_results(self):
@@ -5639,6 +5640,10 @@ def test_rowiterator_to_geodataframe_delegation(self, to_dataframe):
56395640
progress_bar_type,
56405641
create_bqstorage_client,
56415642
geography_as_object=True,
5643+
bool_dtype=DefaultPandasDTypes.BOOL_DTYPE,
5644+
int_dtype=DefaultPandasDTypes.INT_DTYPE,
5645+
float_dtype=None,
5646+
string_dtype=None,
56425647
)
56435648

56445649
self.assertIsInstance(df, geopandas.GeoDataFrame)

‎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)