Skip to content

feat: support bigquery.vector_search() #736

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Jun 7, 2024
Merged

feat: support bigquery.vector_search() #736

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
407db5e
feat: supoort bigquery.vector_search()
ashleyxuu May 16, 2024
86a26bd
minor fix
ashleyxuu May 31, 2024
5e6fe30
address comments
ashleyxuu May 31, 2024
b7e3470
docstring fix
ashleyxuu Jun 1, 2024
2a5de8e
address comments
ashleyxuu Jun 6, 2024
2bbbaf4
small fix
ashleyxuu Jun 6, 2024
f2859f3
add docstring clarification
ashleyxuu Jun 6, 2024
65ef138
Merge branch 'main' into ashleyxu-vector-search
ashleyxuu Jun 6, 2024
a2f3b1d
Update bigframes/bigquery/__init__.py
tswast Jun 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
address comments
  • Loading branch information
@ashleyxuu
ashleyxuu committed Jun 6, 2024
commit 2a5de8ed2bd2b0bc7e7d4c028465b200c7959523
16 changes: 8 additions & 8 deletions bigframes/core/sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ def ordering_clause(

def create_vector_search_sql(
sql_string: str,
options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
options: Mapping[str, Union[str | int | bool | float]] = {},
) -> str:
"""Encode the VECTOR SEARCH statement for BigQuery Vector Search."""

Expand All @@ -191,11 +191,11 @@ def create_vector_search_sql(
distance,
FROM VECTOR_SEARCH(
TABLE `{base_table}`,
"{column_to_search}",
{simple_literal(column_to_search)},
({sql_string}),
"{query_column_to_search}",
distance_type => "{distance_type}",
top_k => {top_k}
{simple_literal(query_column_to_search)},
distance_type => {simple_literal(distance_type)},
top_k => {simple_literal(top_k)}
)
"""
else:
Expand All @@ -206,10 +206,10 @@ def create_vector_search_sql(
distance,
FROM VECTOR_SEARCH(
TABLE `{base_table}`,
"{column_to_search}",
{simple_literal(column_to_search)},
({sql_string}),
distance_type => "{distance_type}",
top_k => {top_k}
distance_type => {simple_literal(distance_type)},
top_k => {simple_literal(top_k)}
)
"""
return query_str
12 changes: 6 additions & 6 deletions tests/unit/core/test_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,9 @@ def test_create_vector_search_sql_simple():
distance,
FROM VECTOR_SEARCH(
TABLE `my_base_table`,
"my_embedding_column",
'my_embedding_column',
({sql_string}),
distance_type => "COSINE",
distance_type => 'COSINE',
top_k => 10
)
"""
Expand Down Expand Up @@ -63,11 +63,11 @@ def test_create_vector_search_sql_query_column_to_search():
base.*,
distance,
FROM VECTOR_SEARCH(
TABLE `my_base_table`,
"my_embedding_column",
TABLE my_base_table,
'my_embedding_column',
({sql_string}),
"new_embedding_column",
distance_type => "COSINE",
'new_embedding_column',
distance_type => 'COSINE',
top_k => 10
)
"""
Expand Down