Skip to content

Commit ededfa9

Browse files
committed
feat(bigquery): add bigquery_vector_search_brute_force sample to run exact brute-force vector search
Problem: The BigQuery documentation (https://cloud.google.com/bigquery/docs/vector-search) needs a standalone, tested Python code sample demonstrating how to run exact brute-force vector search (region tag `bigquery_vector_search_brute_force`). Solution: - Add `vector_search_brute_force(project_id: str) -> None` in `bigquery/vector_search/bigquery_vector_search_brute_force.py` (Compute exact nearest-neighbor distances over embeddings using VECTOR_SEARCH with use_brute_force => TRUE.) - Add unit tests (`bigquery/vector_search/bigquery_vector_search_brute_force_test.py`) covering both the success path and API error handling (`GoogleAPICallError`)
1 parent 186be98 commit ededfa9

5 files changed

Lines changed: 178 additions & 0 deletions

File tree

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START bigquery_vector_search_brute_force]
16+
import sys
17+
18+
from google.api_core import exceptions
19+
from google.cloud import bigquery
20+
21+
22+
def vector_search_brute_force(
23+
project_id: str,
24+
) -> None:
25+
"""Runs an exact brute-force vector similarity search over BBC news embeddings.
26+
27+
Args:
28+
project_id: Google Cloud project ID used to bill the query job
29+
(for example, 'example-project').
30+
"""
31+
client = bigquery.Client(project=project_id)
32+
# Query the public BBC News dataset in bigquery-public-data
33+
# (https://cloud.google.com/bigquery/public-data).
34+
table_ref = "bigquery-public-data.bbc_news.fulltext"
35+
# Pass a 768-dimensional query vector and enable use_brute_force for exact distance calculation.
36+
category_filter = "tech"
37+
query_embedding = [0.01] * 768
38+
job_config = bigquery.QueryJobConfig(
39+
query_parameters=[
40+
bigquery.ScalarQueryParameter("category_filter", "STRING", category_filter),
41+
bigquery.ArrayQueryParameter("query_embedding", "FLOAT64", query_embedding),
42+
]
43+
)
44+
sql = f"""
45+
SELECT
46+
base.title AS matched_title,
47+
base.category AS matched_category,
48+
distance
49+
FROM VECTOR_SEARCH(
50+
(
51+
SELECT
52+
title,
53+
category,
54+
body_embedded
55+
FROM `{table_ref}`
56+
WHERE category = @category_filter
57+
AND ARRAY_LENGTH(body_embedded) = 768
58+
),
59+
'body_embedded',
60+
(SELECT @query_embedding AS query_embedding),
61+
query_column_to_search => 'query_embedding',
62+
top_k => 5,
63+
distance_type => 'COSINE',
64+
options => '{{"use_brute_force": true}}'
65+
)
66+
ORDER BY distance ASC, matched_title ASC;
67+
"""
68+
69+
try:
70+
print(
71+
f"Running brute-force vector search on {table_ref} (category={category_filter})"
72+
)
73+
query_job = client.query(sql, job_config=job_config)
74+
rows = list(query_job.result())
75+
except exceptions.NotFound as err:
76+
print(
77+
f"error: public embeddings table '{table_ref}' not found when querying "
78+
f"from project '{project_id}': {err.message}",
79+
file=sys.stderr,
80+
)
81+
return
82+
except exceptions.GoogleAPICallError as err:
83+
print(
84+
f"error: failed to run brute-force vector search on '{table_ref}' in "
85+
f"project '{project_id}': {err.message}",
86+
file=sys.stderr,
87+
)
88+
return
89+
90+
print(f"Resource ID: {table_ref}")
91+
print("Distance Type: COSINE")
92+
print(f"Embedding Dimensions: {len(query_embedding)}")
93+
print(f"Job ID: {query_job.job_id}")
94+
print(f"Neighbors Returned: {len(rows)}")
95+
96+
97+
# [END bigquery_vector_search_brute_force]
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from unittest import mock
16+
17+
import bigquery_vector_search_brute_force
18+
from google.api_core import exceptions
19+
import pytest
20+
21+
22+
def test_vector_search_brute_force_success(
23+
capsys: pytest.CaptureFixture[str],
24+
) -> None:
25+
with mock.patch("google.cloud.bigquery.Client") as mock_client_cls:
26+
mock_client = mock_client_cls.return_value
27+
mock_job = mock.MagicMock()
28+
mock_job.job_id = "example-job-123"
29+
mock_job.result.return_value = [mock.MagicMock() for _ in range(5)]
30+
mock_client.query.return_value = mock_job
31+
32+
bigquery_vector_search_brute_force.vector_search_brute_force(
33+
project_id="example-project",
34+
)
35+
36+
mock_client.query.assert_called_once()
37+
captured = capsys.readouterr()
38+
assert "Resource ID: bigquery-public-data.bbc_news.fulltext" in captured.out
39+
assert "Embedding Dimensions: 768" in captured.out
40+
assert "Neighbors Returned: 5" in captured.out
41+
42+
43+
def test_vector_search_brute_force_error(
44+
capsys: pytest.CaptureFixture[str],
45+
) -> None:
46+
with mock.patch("google.cloud.bigquery.Client") as mock_client_cls:
47+
mock_client = mock_client_cls.return_value
48+
mock_client.query.side_effect = exceptions.NotFound("Table not found")
49+
50+
bigquery_vector_search_brute_force.vector_search_brute_force(
51+
project_id="example-project",
52+
)
53+
54+
captured = capsys.readouterr()
55+
assert "error:" in captured.err
56+
assert "bbc_news.fulltext" in captured.err
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
TEST_CONFIG_OVERRIDE = {
16+
# You can opt out from the test for specific Python versions.
17+
"ignored_versions": ["2.7", "3.7", "3.8", "3.9", "3.10", "3.11"],
18+
# Enforce type hints on new samples.
19+
"enforce_type_hints": True,
20+
"gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
21+
"pip_version_override": None,
22+
"envs": {},
23+
}
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
pytest==9.0.3; python_version >= "3.10"
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
google-cloud-bigquery==3.47.0

0 commit comments

Comments
 (0)