feat(bigquery): add QueryJobConfig properties to bigquery backend specified at query time (#11255)

Author: Wopple

ibis/backends/bigquery/__init__.py

@@ -4,6 +4,7 @@
 
 import concurrent.futures
 import contextlib
+import copy
 import glob
 import os
 import re
@@ -765,17 +766,18 @@ def _get_schema_using_query(self, query: str) -> sch.Schema:
         )
         return BigQuerySchema.to_ibis(job.schema)
 
-    def raw_sql(self, query: str, params=None) -> RowIterator:
-        query_parameters = [
-            bigquery_param(param.type(), value, param.get_name())
-            for param, value in (params or {}).items()
-        ]
+    def raw_sql(
+        self,
+        query: str,
+        params: Mapping[ir.Scalar, Any] | None = None,
+        query_job_config: bq.QueryJobConfig | None = None,
+    ) -> RowIterator:
+        job_config = _merge_params_into_config(query_job_config, params)
+
         with contextlib.suppress(AttributeError):
             query = query.sql(self.dialect)
 
-        job_config = bq.job.QueryJobConfig(query_parameters=query_parameters or [])
-
-        return self._client_query(
+        return self.client.query_and_wait(
             query, job_config=job_config, project=self.billing_project
         )
 
@@ -867,13 +869,18 @@ def _to_query(
         *,
         params: Mapping[ir.Scalar, Any] | None = None,
         limit: int | str | None = None,
+        query_job_config: bq.QueryJobConfig | None = None,
         **kwargs: Any,
     ) -> RowIterator:
         self._run_pre_execute_hooks(table_expr)
         sql = self.compile(table_expr, limit=limit, params=params, **kwargs)
         self._log(sql)
 
-        return self.raw_sql(sql, params=params)
+        return self.raw_sql(
+            sql,
+            params=params,
+            query_job_config=query_job_config,
+        )
 
     def to_pyarrow(
         self,
@@ -882,14 +889,21 @@ def to_pyarrow(
         *,
         params: Mapping[ir.Scalar, Any] | None = None,
         limit: int | str | None = None,
+        query_job_config: bq.QueryJobConfig | None = None,
         **kwargs: Any,
     ) -> pa.Table:
         self._import_pyarrow()
 
         table_expr = expr.as_table()
         schema = table_expr.schema() - ibis.schema({"_TABLE_SUFFIX": "string"})
 
-        query = self._to_query(table_expr, params=params, limit=limit, **kwargs)
+        query = self._to_query(
+            table_expr,
+            params=params,
+            limit=limit,
+            query_job_config=query_job_config,
+            **kwargs,
+        )
         table = query.to_arrow(
             progress_bar_type=None, bqstorage_client=self.storage_client
         )
@@ -904,6 +918,7 @@ def to_pyarrow_batches(
         params: Mapping[ir.Scalar, Any] | None = None,
         limit: int | str | None = None,
         chunk_size: int = 1_000_000,
+        query_job_config: bq.QueryJobConfig | None = None,
         **kwargs: Any,
     ):
         pa = self._import_pyarrow()
@@ -912,7 +927,13 @@ def to_pyarrow_batches(
         schema = table_expr.schema() - ibis.schema({"_TABLE_SUFFIX": "string"})
         colnames = list(schema.names)
 
-        query = self._to_query(table_expr, params=params, limit=limit, **kwargs)
+        query = self._to_query(
+            table_expr,
+            params=params,
+            limit=limit,
+            query_job_config=query_job_config,
+            **kwargs,
+        )
         batch_iter = query.to_arrow_iterable(bqstorage_client=self.storage_client)
         return pa.ipc.RecordBatchReader.from_batches(
             schema.to_pyarrow(),
@@ -926,6 +947,7 @@ def execute(
         *,
         params: Mapping[ir.Scalar, Any] | None = None,
         limit: int | str | None = None,
+        query_job_config: bq.QueryJobConfig | None = None,
         **kwargs: Any,
     ) -> pd.DataFrame | pd.Series | Any:
         """Compile and execute the given Ibis expression.
@@ -942,6 +964,8 @@ def execute(
             already set on the expression.
         params
             Query parameters
+        query_job_config
+            QueryJobConfig, the values in the `params` argument take precedence over the ones in this object
         kwargs
             Extra arguments specific to the backend
 
@@ -955,7 +979,13 @@ def execute(
 
         table_expr = expr.as_table()
         schema = table_expr.schema() - ibis.schema({"_TABLE_SUFFIX": "string"})
-        query = self._to_query(table_expr, params=params, limit=limit, **kwargs)
+        query = self._to_query(
+            table_expr,
+            params=params,
+            limit=limit,
+            query_job_config=query_job_config,
+            **kwargs,
+        )
         df = query.to_arrow(
             progress_bar_type=None, bqstorage_client=self.storage_client
         ).to_pandas(timestamp_as_object=True)
@@ -1245,6 +1275,32 @@ def _safe_raw_sql(self, *args, **kwargs):
         yield self.raw_sql(*args, **kwargs)
 
 
+def _merge_params_into_config(
+    query_job_config: bq.QueryJobConfig | None = None,
+    params: Mapping[ir.Scalar, Any] | None = None,
+) -> bq.QueryJobConfig:
+    """Merge parameters into a QueryJobConfig.
+
+    Returns a copy of `query_job_config` with the `params` merged into the `query_parameters`
+    field. `params` will override values with a key naming conflict in `query_job_config`.
+    """
+
+    if query_job_config is not None:
+        query_job_config = copy.deepcopy(query_job_config)  # do not modify the input
+    else:
+        query_job_config = bq.QueryJobConfig()
+
+    config_params = {param.name: param for param in query_job_config.query_parameters}
+
+    params_as_bq = {
+        param.get_name(): bigquery_param(param.type(), value, param.get_name())
+        for param, value in (params or {}).items()
+    }
+
+    query_job_config.query_parameters = list({**config_params, **params_as_bq}.values())
+    return query_job_config
+
+
 def compile(expr, params=None, **kwargs):
     """Compile an expression for BigQuery."""
     backend = Backend()

ibis/backends/bigquery/tests/unit/test___init__.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from datetime import date
+
+import pytest
+from google.cloud.bigquery import QueryJobConfig
+from google.cloud.bigquery.query import ScalarQueryParameter
+from google.cloud.bigquery.table import TableReference
+
+import ibis
+from ibis.backends.bigquery import _merge_params_into_config
+
+
+@pytest.mark.parametrize(
+    "query_job_config, params, expected",
+    [
+        (None, None, []),
+        (QueryJobConfig(), None, []),
+        (None, {}, []),
+        (QueryJobConfig(), {}, []),
+        (
+            QueryJobConfig(
+                query_parameters=[
+                    ScalarQueryParameter("param1", "INT64", 1),
+                    ScalarQueryParameter("param2", "INT64", 2),
+                ],
+            ),
+            None,
+            [
+                ScalarQueryParameter("param1", "INT64", 1),
+                ScalarQueryParameter("param2", "INT64", 2),
+            ],
+        ),
+        (
+            None,
+            {
+                ibis.literal(0).name("param1"): 1,
+                ibis.literal(0).name("param2"): 2,
+            },
+            [
+                ScalarQueryParameter("param1", "INT64", 1),
+                ScalarQueryParameter("param2", "INT64", 2),
+            ],
+        ),
+        (
+            QueryJobConfig(
+                query_parameters=[
+                    ScalarQueryParameter("param1", "INT64", 1),
+                    ScalarQueryParameter("param2", "INT64", 2),
+                ],
+            ),
+            {
+                ibis.literal(0).name("param2"): 3,
+                ibis.literal(0).name("param3"): 4,
+            },
+            [
+                ScalarQueryParameter("param1", "INT64", 1),
+                ScalarQueryParameter("param2", "INT64", 3),
+                ScalarQueryParameter("param3", "INT64", 4),
+            ],
+        ),
+        (
+            QueryJobConfig(
+                query_parameters=[
+                    ScalarQueryParameter("config1", "BOOL", True),
+                    ScalarQueryParameter("config2", "INT64", 1),
+                    ScalarQueryParameter("config3", "FLOAT64", 2.3),
+                    ScalarQueryParameter("config4", "STRING", "abc"),
+                    ScalarQueryParameter("config5", "DATE", "2025-01-01"),
+                ],
+                # ensure this is preserved
+                destination=TableReference.from_string(
+                    "test_project.test_dataset.test_table",
+                ),
+            ),
+            {
+                ibis.literal(False).name("param1"): False,
+                ibis.literal(0).name("param2"): 4,
+                ibis.literal(0.0).name("param3"): 5.6,
+                ibis.literal("").name("param4"): "def",
+                ibis.literal(date.today()).name("param5"): date(2025, 1, 2),
+            },
+            [
+                ScalarQueryParameter("config1", "BOOL", True),
+                ScalarQueryParameter("config2", "INT64", 1),
+                ScalarQueryParameter("config3", "FLOAT64", 2.3),
+                ScalarQueryParameter("config4", "STRING", "abc"),
+                ScalarQueryParameter("config5", "DATE", date(2025, 1, 1)),
+                ScalarQueryParameter("param1", "BOOL", False),
+                ScalarQueryParameter("param2", "INT64", 4),
+                ScalarQueryParameter("param3", "FLOAT64", 5.6),
+                ScalarQueryParameter("param4", "STRING", "def"),
+                ScalarQueryParameter("param5", "DATE", date(2025, 1, 2)),
+            ],
+        ),
+    ],
+)
+def test__merge_params_into_config(query_job_config, params, expected):
+    # check the merge is correct
+    result = _merge_params_into_config(query_job_config, params)
+    assert result is not query_job_config
+    assert result.query_parameters == expected
+
+    # check all the other fields are preserved
+    if query_job_config is not None:
+        expected_repr = query_job_config.to_api_repr()
+        result_repr = result.to_api_repr()
+
+        if "queryParameters" in expected_repr["query"]:
+            del expected_repr["query"]["queryParameters"]
+
+        if "queryParameters" in result_repr["query"]:
+            del result_repr["query"]["queryParameters"]
+
+        assert result_repr == expected_repr