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