@@ -81,6 +81,41 @@ def test_pandas_gbq_query():
81
81
assert len (df ) > 0
82
82
83
83
84
+ def test_client_library_query_bqstorage ():
85
+ # [START bigquery_migration_client_library_query_bqstorage]
86
+ import google .auth
87
+ from google .cloud import bigquery
88
+ from google .cloud import bigquery_storage_v1beta1
89
+
90
+ # Create a BigQuery client and a BigQuery Storage API client with the same
91
+ # credentials to avoid authenticating twice.
92
+ credentials , project_id = google .auth .default (
93
+ scopes = ["https://www.googleapis.com/auth/cloud-platform" ]
94
+ )
95
+ client = bigquery .Client (credentials = credentials , project = project_id )
96
+ bqstorage_client = bigquery_storage_v1beta1 .BigQueryStorageClient (
97
+ credentials = credentials
98
+ )
99
+ sql = "SELECT * FROM `bigquery-public-data.irs_990.irs_990_2012`"
100
+
101
+ # Use a BigQuery Storage API client to download results more quickly.
102
+ df = client .query (sql ).to_dataframe (bqstorage_client = bqstorage_client )
103
+ # [END bigquery_migration_client_library_query_bqstorage]
104
+ assert len (df ) > 0
105
+
106
+
107
+ def test_pandas_gbq_query_bqstorage ():
108
+ # [START bigquery_migration_pandas_gbq_query_bqstorage]
109
+ import pandas
110
+
111
+ sql = "SELECT * FROM `bigquery-public-data.irs_990.irs_990_2012`"
112
+
113
+ # Use the BigQuery Storage API to download results more quickly.
114
+ df = pandas .read_gbq (sql , dialect = 'standard' , use_bqstorage_api = True )
115
+ # [END bigquery_migration_pandas_gbq_query_bqstorage]
116
+ assert len (df ) > 0
117
+
118
+
84
119
def test_client_library_legacy_query ():
85
120
# [START bigquery_migration_client_library_query_legacy]
86
121
from google .cloud import bigquery
@@ -184,16 +219,28 @@ def test_client_library_upload_from_dataframe(temp_dataset):
184
219
}
185
220
)
186
221
client = bigquery .Client ()
187
- dataset_ref = client . dataset ( 'my_dataset' )
222
+ table_id = 'my_dataset.new_table'
188
223
# [END bigquery_migration_client_library_upload_from_dataframe]
189
- dataset_ref = client .dataset (temp_dataset .dataset_id )
224
+ table_id = (
225
+ temp_dataset .dataset_id
226
+ + ".test_client_library_upload_from_dataframe"
227
+ )
190
228
# [START bigquery_migration_client_library_upload_from_dataframe]
191
- table_ref = dataset_ref .table ('new_table' )
229
+ # Since string columns use the "object" dtype, pass in a (partial) schema
230
+ # to ensure the correct BigQuery data type.
231
+ job_config = bigquery .LoadJobConfig (schema = [
232
+ bigquery .SchemaField ("my_string" , "STRING" ),
233
+ ])
234
+
235
+ job = client .load_table_from_dataframe (
236
+ df , table_id , job_config = job_config
237
+ )
192
238
193
- client .load_table_from_dataframe (df , table_ref ).result ()
239
+ # Wait for the load job to complete.
240
+ job .result ()
194
241
# [END bigquery_migration_client_library_upload_from_dataframe]
195
242
client = bigquery .Client ()
196
- table = client .get_table (table_ref )
243
+ table = client .get_table (table_id )
197
244
assert table .num_rows == 3
198
245
199
246
@@ -209,16 +256,16 @@ def test_pandas_gbq_upload_from_dataframe(temp_dataset):
209
256
'my_float64' : [4.0 , 5.0 , 6.0 ],
210
257
}
211
258
)
212
- full_table_id = 'my_dataset.new_table'
213
- project_id = 'my-project-id'
259
+ table_id = 'my_dataset.new_table'
214
260
# [END bigquery_migration_pandas_gbq_upload_from_dataframe]
215
- table_id = 'new_table'
216
- full_table_id = '{}.{}' .format (temp_dataset .dataset_id , table_id )
217
- project_id = os .environ ['GCLOUD_PROJECT' ]
261
+ table_id = (
262
+ temp_dataset .dataset_id
263
+ + ".test_pandas_gbq_upload_from_dataframe"
264
+ )
218
265
# [START bigquery_migration_pandas_gbq_upload_from_dataframe]
219
266
220
- df .to_gbq (full_table_id , project_id = project_id )
267
+ df .to_gbq (table_id )
221
268
# [END bigquery_migration_pandas_gbq_upload_from_dataframe]
222
269
client = bigquery .Client ()
223
- table = client .get_table (temp_dataset . table ( table_id ) )
270
+ table = client .get_table (table_id )
224
271
assert table .num_rows == 3
0 commit comments