17
17
import os
18
18
import re
19
19
from unittest import mock
20
+ import warnings
20
21
21
22
import google .api_core .exceptions
22
23
import google .cloud .bigquery
@@ -186,7 +187,7 @@ def get_table_mock(table_ref):
186
187
187
188
188
189
@pytest .mark .parametrize ("table" , CLUSTERED_OR_PARTITIONED_TABLES )
189
- def test_no_default_index_error_raised_by_read_gbq (table ):
190
+ def test_default_index_warning_raised_by_read_gbq (table ):
190
191
"""Because of the windowing operation to create a default index, row
191
192
filters can't push down to the clustering column.
192
193
@@ -202,12 +203,12 @@ def test_no_default_index_error_raised_by_read_gbq(table):
202
203
session = resources .create_bigquery_session (bqclient = bqclient )
203
204
table ._properties ["location" ] = session ._location
204
205
205
- with pytest .raises (bigframes .exceptions .NoDefaultIndexError ):
206
+ with pytest .warns (bigframes .exceptions .DefaultIndexWarning ):
206
207
session .read_gbq ("my-project.my_dataset.my_table" )
207
208
208
209
209
210
@pytest .mark .parametrize ("table" , CLUSTERED_OR_PARTITIONED_TABLES )
210
- def test_no_default_index_error_not_raised_by_read_gbq_index_col_sequential_int64 (
211
+ def test_default_index_warning_not_raised_by_read_gbq_index_col_sequential_int64 (
211
212
table ,
212
213
):
213
214
"""Because of the windowing operation to create a default index, row
@@ -224,11 +225,13 @@ def test_no_default_index_error_not_raised_by_read_gbq_index_col_sequential_int6
224
225
session = resources .create_bigquery_session (bqclient = bqclient )
225
226
table ._properties ["location" ] = session ._location
226
227
227
- # No exception raised because we set the option allowing the default indexes.
228
- df = session .read_gbq (
229
- "my-project.my_dataset.my_table" ,
230
- index_col = bigframes .enums .DefaultIndexKind .SEQUENTIAL_INT64 ,
231
- )
228
+ # No warnings raised because we set the option allowing the default indexes.
229
+ with warnings .catch_warnings ():
230
+ warnings .simplefilter ("error" , bigframes .exceptions .DefaultIndexWarning )
231
+ df = session .read_gbq (
232
+ "my-project.my_dataset.my_table" ,
233
+ index_col = bigframes .enums .DefaultIndexKind .SEQUENTIAL_INT64 ,
234
+ )
232
235
233
236
# We expect a window operation because we specificaly requested a sequential index.
234
237
generated_sql = df .sql .casefold ()
@@ -246,7 +249,7 @@ def test_no_default_index_error_not_raised_by_read_gbq_index_col_sequential_int6
246
249
),
247
250
)
248
251
@pytest .mark .parametrize ("table" , CLUSTERED_OR_PARTITIONED_TABLES )
249
- def test_no_default_index_error_not_raised_by_read_gbq_index_col_columns (
252
+ def test_default_index_warning_not_raised_by_read_gbq_index_col_columns (
250
253
total_count ,
251
254
distinct_count ,
252
255
table ,
@@ -270,18 +273,20 @@ def test_no_default_index_error_not_raised_by_read_gbq_index_col_columns(
270
273
)
271
274
table ._properties ["location" ] = session ._location
272
275
273
- # No exception raised because there are columns to use as the index.
274
- df = session .read_gbq (
275
- "my-project.my_dataset.my_table" , index_col = ("idx_1" , "idx_2" )
276
- )
276
+ # No warning raised because there are columns to use as the index.
277
+ with warnings .catch_warnings ():
278
+ warnings .simplefilter ("error" , bigframes .exceptions .DefaultIndexWarning )
279
+ df = session .read_gbq (
280
+ "my-project.my_dataset.my_table" , index_col = ("idx_1" , "idx_2" )
281
+ )
277
282
278
283
# There should be no analytic operators to prevent row filtering pushdown.
279
284
assert "OVER" not in df .sql
280
285
assert tuple (df .index .names ) == ("idx_1" , "idx_2" )
281
286
282
287
283
288
@pytest .mark .parametrize ("table" , CLUSTERED_OR_PARTITIONED_TABLES )
284
- def test_no_default_index_error_not_raised_by_read_gbq_primary_key (table ):
289
+ def test_default_index_warning_not_raised_by_read_gbq_primary_key (table ):
285
290
"""If a primary key is set on the table, we use that as the index column
286
291
by default, no error should be raised in this case.
287
292
@@ -310,8 +315,10 @@ def test_no_default_index_error_not_raised_by_read_gbq_primary_key(table):
310
315
)
311
316
table ._properties ["location" ] = session ._location
312
317
313
- # No exception raised because there is a primary key to use as the index.
314
- df = session .read_gbq ("my-project.my_dataset.my_table" )
318
+ # No warning raised because there is a primary key to use as the index.
319
+ with warnings .catch_warnings ():
320
+ warnings .simplefilter ("error" , bigframes .exceptions .DefaultIndexWarning )
321
+ df = session .read_gbq ("my-project.my_dataset.my_table" )
315
322
316
323
# There should be no analytic operators to prevent row filtering pushdown.
317
324
assert "OVER" not in df .sql
0 commit comments