fix: do not warn with DefaultIndexWarning in partial ordering mode (#2230)

Thank you for opening a Pull Request! Before submitting your PR, there
are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/python-bigquery-dataframes/issues/new/choose)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
Fixes internal issue b/356872356 🦕

Author: tswast

bigframes/session/_io/bigquery/read_gbq_table.py

@@ -402,6 +402,7 @@ def get_index_cols(
 |bigframes.enums.DefaultIndexKind,
 *,
 rename_to_schema: Optional[Dict[str, str]] =None,
+default_index_type: bigframes.enums.DefaultIndexKind=bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64,
 ) ->List[str]:
 """
  If we can get a total ordering from the table, such as via primary key
@@ -471,7 +472,11 @@ def get_index_cols(
 # find index_cols to use. This is to avoid unexpected performance and
 # resource utilization because of the default sequential index. See
 # internal issue 335727141.
-if_is_table_clustered_or_partitioned(table) andnotprimary_keys:
+if (
+_is_table_clustered_or_partitioned(table)
+andnotprimary_keys
+anddefault_index_type==bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64
+ ):
 msg=bfe.format_message(
 f"Table '{str(table.reference)}' is clustered and/or "
 "partitioned, but BigQuery DataFrames was not able to find a "

bigframes/session/loader.py

@@ -696,6 +696,7 @@ def read_gbq_table(
 table=table,
 index_col=index_col,
 rename_to_schema=rename_to_schema,
+default_index_type=self._default_index_type,
  )
 _check_index_col_param(
 index_cols,

tests/unit/session/test_read_gbq_table.py

@@ -15,10 +15,13 @@
 """Unit tests for read_gbq_table helper functions."""
 
 importunittest.mockasmock
+importwarnings
 
 importgoogle.cloud.bigquery
 importpytest
 
+importbigframes.enums
+importbigframes.exceptions
 importbigframes.session._io.bigquery.read_gbq_tableasbf_read_gbq_table
 frombigframes.testingimportmocks
 
@@ -143,3 +146,43 @@ def test_check_if_index_columns_are_unique(index_cols, values_distinct, expected
  )
 
 assertresult==expected
+
+
+deftest_get_index_cols_warns_if_clustered_but_sequential_index():
+table=google.cloud.bigquery.Table.from_api_repr(
+ {
+"tableReference": {
+"projectId": "my-project",
+"datasetId": "my_dataset",
+"tableId": "my_table",
+ },
+"clustering": {
+"fields": ["col1", "col2"],
+ },
+ },
+ )
+table.schema= (
+google.cloud.bigquery.SchemaField("col1", "INT64"),
+google.cloud.bigquery.SchemaField("col2", "INT64"),
+google.cloud.bigquery.SchemaField("col3", "INT64"),
+google.cloud.bigquery.SchemaField("col4", "INT64"),
+ )
+
+withpytest.warns(bigframes.exceptions.DefaultIndexWarning, match="is clustered"):
+bf_read_gbq_table.get_index_cols(
+table,
+index_col=(),
+default_index_type=bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64,
+ )
+
+# Ensure that we don't raise if using a NULL index by default, such as in
+# partial ordering mode. See: internal issue b/356872356.
+withwarnings.catch_warnings():
+warnings.simplefilter(
+"error", category=bigframes.exceptions.DefaultIndexWarning
+ )
+bf_read_gbq_table.get_index_cols(
+table,
+index_col=(),
+default_index_type=bigframes.enums.DefaultIndexKind.NULL,
+ )