feat: add configuration option to read_gbq (#401)

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://togithub.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 https://togithub.com/googleapis/python-bigquery-dataframes/issues/384 🦕

Author: Genesis929

bigframes/pandas/__init__.py

@@ -492,9 +492,10 @@ def read_gbq(
  *,
  index_col: Iterable[str] | str = (),
  columns: Iterable[str] = (),
+ configuration: Optional[Dict] = None,
  max_results: Optional[int] = None,
  filters: vendored_pandas_gbq.FiltersType = (),
- use_cache: bool = True,
+ use_cache: Optional[bool] = None,
  col_order: Iterable[str] = (),
 ) -> bigframes.dataframe.DataFrame:
  _set_default_session_location_if_possible(query_or_table)
@@ -503,6 +504,7 @@ def read_gbq(
  query_or_table,
  index_col=index_col,
  columns=columns,
+ configuration=configuration,
  max_results=max_results,
  filters=filters,
  use_cache=use_cache,
@@ -528,8 +530,9 @@ def read_gbq_query(
  *,
  index_col: Iterable[str] | str = (),
  columns: Iterable[str] = (),
+ configuration: Optional[Dict] = None,
  max_results: Optional[int] = None,
- use_cache: bool = True,
+ use_cache: Optional[bool] = None,
  col_order: Iterable[str] = (),
 ) -> bigframes.dataframe.DataFrame:
  _set_default_session_location_if_possible(query)
@@ -538,6 +541,7 @@ def read_gbq_query(
  query,
  index_col=index_col,
  columns=columns,
+ configuration=configuration,
  max_results=max_results,
  use_cache=use_cache,
  col_order=col_order,

bigframes/session/__init__.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import copy
 import datetime
 import itertools
 import logging
@@ -283,9 +284,10 @@ def read_gbq(
  *,
  index_col: Iterable[str] | str = (),
  columns: Iterable[str] = (),
+ configuration: Optional[Dict] = None,
  max_results: Optional[int] = None,
  filters: third_party_pandas_gbq.FiltersType = (),
- use_cache: bool = True,
+ use_cache: Optional[bool] = None,
  col_order: Iterable[str] = (),
  # Add a verify index argument that fails if the index is not unique.
  ) -> dataframe.DataFrame:
@@ -306,6 +308,7 @@ def read_gbq(
  query_or_table,
  index_col=index_col,
  columns=columns,
+ configuration=configuration,
  max_results=max_results,
  api_name="read_gbq",
  use_cache=use_cache,
@@ -314,13 +317,20 @@ def read_gbq(
  # TODO(swast): Query the snapshot table but mark it as a
  # deterministic query so we can avoid serializing if we have a
  # unique index.
+ if configuration is not None:
+ raise ValueError(
+ "The 'configuration' argument is not allowed when "
+ "directly reading from a table. Please remove "
+ "'configuration' or use a query."
+ )
+
  return self._read_gbq_table(
  query_or_table,
  index_col=index_col,
  columns=columns,
  max_results=max_results,
  api_name="read_gbq",
- use_cache=use_cache,
+ use_cache=use_cache if use_cache is not None else True,
  )
 
  def _to_query(
@@ -405,7 +415,7 @@ def _query_to_destination(
  query: str,
  index_cols: List[str],
  api_name: str,
- use_cache: bool = True,
+ configuration: dict = {"query": {"useQueryCache": True}},
  ) -> Tuple[Optional[bigquery.TableReference], Optional[bigquery.QueryJob]]:
  # If a dry_run indicates this is not a query type job, then don't
  # bother trying to do a CREATE TEMP TABLE ... AS SELECT ... statement.
@@ -427,23 +437,35 @@ def _query_to_destination(
  ][:_MAX_CLUSTER_COLUMNS]
  temp_table = self._create_empty_temp_table(schema, cluster_cols)
 
- job_config = bigquery.QueryJobConfig()
+ timeout_ms = configuration.get("jobTimeoutMs") or configuration["query"].get(
+ "timeoutMs"
+ )
+
+ # Convert timeout_ms to seconds, ensuring a minimum of 0.1 seconds to avoid
+ # the program getting stuck on too-short timeouts.
+ timeout = max(int(timeout_ms) * 1e-3, 0.1) if timeout_ms else None
+
+ job_config = typing.cast(
+ bigquery.QueryJobConfig,
+ bigquery.QueryJobConfig.from_api_repr(configuration),
+ )
  job_config.labels["bigframes-api"] = api_name
  job_config.destination = temp_table
- job_config.use_query_cache = use_cache
 
  try:
  # Write to temp table to workaround BigQuery 10 GB query results
  # limit. See: internal issue 303057336.
  job_config.labels["error_caught"] = "true"
- _, query_job = self._start_query(query, job_config=job_config)
+ _, query_job = self._start_query(
+ query, job_config=job_config, timeout=timeout
+ )
  return query_job.destination, query_job
  except google.api_core.exceptions.BadRequest:
  # Some SELECT statements still aren't compatible with cluster
  # tables as the destination. For example, if the query has a
  # top-level ORDER BY, this conflicts with our ability to cluster
  # the table by the index column(s).
- _, query_job = self._start_query(query)
+ _, query_job = self._start_query(query, timeout=timeout)
  return query_job.destination, query_job
 
  def read_gbq_query(
@@ -452,8 +474,9 @@ def read_gbq_query(
  *,
  index_col: Iterable[str] | str = (),
  columns: Iterable[str] = (),
+ configuration: Optional[Dict] = None,
  max_results: Optional[int] = None,
- use_cache: bool = True,
+ use_cache: Optional[bool] = None,
  col_order: Iterable[str] = (),
  ) -> dataframe.DataFrame:
  """Turn a SQL query into a DataFrame.
@@ -517,6 +540,7 @@ def read_gbq_query(
  query=query,
  index_col=index_col,
  columns=columns,
+ configuration=configuration,
  max_results=max_results,
  api_name="read_gbq_query",
  use_cache=use_cache,
@@ -528,10 +552,34 @@ def _read_gbq_query(
  *,
  index_col: Iterable[str] | str = (),
  columns: Iterable[str] = (),
+ configuration: Optional[Dict] = None,
  max_results: Optional[int] = None,
  api_name: str = "read_gbq_query",
- use_cache: bool = True,
+ use_cache: Optional[bool] = None,
  ) -> dataframe.DataFrame:
+ configuration = _transform_read_gbq_configuration(configuration)
+
+ if "query" not in configuration:
+ configuration["query"] = {}
+
+ if "query" in configuration["query"]:
+ raise ValueError(
+ "The query statement must not be included in the ",
+ "'configuration' because it is already provided as",
+ " a separate parameter.",
+ )
+
+ if "useQueryCache" in configuration["query"]:
+ if use_cache is not None:
+ raise ValueError(
+ "'useQueryCache' in 'configuration' conflicts with"
+ " 'use_cache' parameter. Please specify only one."
+ )
+ else:
+ configuration["query"]["useQueryCache"] = (
+ True if use_cache is None else use_cache
+ )
+
  if isinstance(index_col, str):
  index_cols = [index_col]
  else:
@@ -541,7 +589,7 @@ def _read_gbq_query(
  query,
  index_cols,
  api_name=api_name,
- use_cache=use_cache,
+ configuration=configuration,
  )
 
  # If there was no destination table, that means the query must have
@@ -565,7 +613,7 @@ def _read_gbq_query(
  index_col=index_cols,
  columns=columns,
  max_results=max_results,
- use_cache=use_cache,
+ use_cache=configuration["query"]["useQueryCache"],
  )
 
  def read_gbq_table(
@@ -1656,13 +1704,14 @@ def _start_query(
  sql: str,
  job_config: Optional[bigquery.job.QueryJobConfig] = None,
  max_results: Optional[int] = None,
+ timeout: Optional[float] = None,
  ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]:
  """
  Starts BigQuery query job and waits for results.
  """
  job_config = self._prepare_query_job_config(job_config)
  return bigframes.session._io.bigquery.start_query_with_client(
- self.bqclient, sql, job_config, max_results
+ self.bqclient, sql, job_config, max_results, timeout
  )
 
  def _start_query_ml_ddl(
@@ -1876,3 +1925,25 @@ def _convert_to_nonnull_string(column: ibis_types.Column) -> ibis_types.StringVa
  # Escape backslashes and use backslash as delineator
  escaped = typing.cast(ibis_types.StringColumn, result.fillna("")).replace("\\", "\\\\") # type: ignore
  return typing.cast(ibis_types.StringColumn, ibis.literal("\\")).concat(escaped)
+
+
+def _transform_read_gbq_configuration(configuration: Optional[dict]) -> dict:
+ """
+ For backwards-compatibility, convert any previously client-side only
+ parameters such as timeoutMs to the property name expected by the REST API.
+
+ Makes a copy of configuration if changes are needed.
+ """
+
+ if configuration is None:
+ return {}
+
+ timeout_ms = configuration.get("query", {}).get("timeoutMs")
+ if timeout_ms is not None:
+ # Transform timeoutMs to an actual server-side configuration.
+ # https://github.com/googleapis/python-bigquery-pandas/issues/479
+ configuration = copy.deepcopy(configuration)
+ del configuration["query"]["timeoutMs"]
+ configuration["jobTimeoutMs"] = timeout_ms
+
+ return configuration

bigframes/session/_io/bigquery.py

@@ -220,6 +220,7 @@ def start_query_with_client(
  sql: str,
  job_config: bigquery.job.QueryJobConfig,
  max_results: Optional[int] = None,
+ timeout: Optional[float] = None,
 ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]:
  """
  Starts query job and waits for results.
@@ -230,7 +231,7 @@ def start_query_with_client(
  )
 
  try:
- query_job = bq_client.query(sql, job_config=job_config)
+ query_job = bq_client.query(sql, job_config=job_config, timeout=timeout)
  except google.api_core.exceptions.Forbidden as ex:
  if "Drive credentials" in ex.message:
  ex.message += "\nCheck https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions."

tests/system/small/test_session.py

@@ -20,6 +20,7 @@
 import typing
 from typing import List
 
+import google
 import google.cloud.bigquery as bigquery
 import numpy as np
 import pandas as pd
@@ -363,6 +364,47 @@ def test_read_gbq_table_wildcard_with_filter(session: bigframes.Session):
  assert df.shape == (348485, 32)
 
 
+@pytest.mark.parametrize(
+ ("config"),
+ [
+ {
+ "query": {
+ "useQueryCache": True,
+ "maximumBytesBilled": "1000000000",
+ "timeoutMs": 10000,
+ }
+ },
+ pytest.param(
+ {"query": {"useQueryCache": True, "timeoutMs": 50}},
+ marks=pytest.mark.xfail(
+ raises=google.api_core.exceptions.BadRequest,
+ reason="Expected failure due to timeout being set too short.",
+ ),
+ ),
+ pytest.param(
+ {"query": {"useQueryCache": False, "maximumBytesBilled": "100"}},
+ marks=pytest.mark.xfail(
+ raises=google.api_core.exceptions.InternalServerError,
+ reason="Expected failure when the query exceeds the maximum bytes billed limit.",
+ ),
+ ),
+ ],
+)
+def test_read_gbq_with_configuration(
+ session: bigframes.Session, scalars_table_id: str, config: dict
+):
+ query = f"""SELECT
+ t.float64_col * 2 AS my_floats,
+ CONCAT(t.string_col, "_2") AS my_strings,
+ t.int64_col > 0 AS my_bools,
+ FROM `{scalars_table_id}` AS t
+ """
+
+ df = session.read_gbq(query, configuration=config)
+
+ assert df.shape == (9, 3)
+
+
 def test_read_gbq_model(session, penguins_linear_model_name):
  model = session.read_gbq_model(penguins_linear_model_name)
  assert isinstance(model, bigframes.ml.linear_model.LinearRegression)

third_party/bigframes_vendored/pandas/io/gbq.py

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Iterable, Literal, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, Literal, Optional, Tuple, Union
 
 from bigframes import constants
 
@@ -19,9 +19,10 @@ def read_gbq(
  *,
  index_col: Iterable[str] | str = (),
  columns: Iterable[str] = (),
+ configuration: Optional[Dict] = None,
  max_results: Optional[int] = None,
  filters: FiltersType = (),
- use_cache: bool = True,
+ use_cache: Optional[bool] = None,
  col_order: Iterable[str] = (),
  ):
  """Loads a DataFrame from BigQuery.
@@ -107,6 +108,11 @@ def read_gbq(
  columns (Iterable[str]):
  List of BigQuery column names in the desired order for results
  DataFrame.
+ configuration (dict, optional):
+ Query config parameters for job processing.
+ For example: configuration = {'query': {'useQueryCache': False}}.
+ For more information see `BigQuery REST API Reference
+ <https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.query>`__.
  max_results (Optional[int], default None):
  If set, limit the maximum number of rows to fetch from the
  query results.
@@ -121,8 +127,10 @@ def read_gbq(
  If using wildcard table suffix in query_or_table, can specify
  '_table_suffix' pseudo column to filter the tables to be read
  into the DataFrame.
- use_cache (bool, default True):
- Whether to cache the query inputs. Default to True.
+ use_cache (Optional[bool], default None):
+ Caches query results if set to `True`. When `None`, it behaves
+ as `True`, but should not be combined with `useQueryCache` in
+ `configuration` to avoid conflicts.
  col_order (Iterable[str]):
  Alias for columns, retained for backwards compatibility.