address comments

Author: ashleyxuu

bigframes/bigquery/__init__.py

@@ -20,17 +20,23 @@
 
 from __future__ import annotations
 
+import typing
 from typing import Literal, Optional, Union
 
-import bigframes.bigquery.utils as utils
 import bigframes.constants as constants
 import bigframes.core.groupby as groupby
+import bigframes.core.sql
+import bigframes.ml.utils as utils
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
-import bigframes.pandas as bpd
+import bigframes.series
 
+if typing.TYPE_CHECKING:
+    import bigframes.dataframe as dataframe
+    import bigframes.series as series
 
-def array_length(series: bpd.Series) -> bpd.Series:
+
+def array_length(series: series.Series) -> series.Series:
     """Compute the length of each array element in the Series.
 
     **Examples:**
@@ -67,7 +73,7 @@ def array_length(series: bpd.Series) -> bpd.Series:
 
 def array_agg(
     obj: groupby.SeriesGroupBy | groupby.DataFrameGroupBy,
-) -> bpd.Series | bpd.DataFrame:
+) -> series.Series | dataframe.DataFrame:
     """Group data and create arrays from selected columns, omitting NULLs to avoid
     BigQuery errors (NULLs not allowed in arrays).
 
@@ -118,7 +124,7 @@ def array_agg(
         )
 
 
-def array_to_string(series: bpd.Series, delimiter: str) -> bpd.Series:
+def array_to_string(series: series.Series, delimiter: str) -> series.Series:
     """Converts array elements within a Series into delimited strings.
 
     **Examples:**
@@ -151,14 +157,14 @@ def array_to_string(series: bpd.Series, delimiter: str) -> bpd.Series:
 def vector_search(
     base_table: str,
     column_to_search: str,
-    query: Union[bpd.DataFrame, bpd.Series],
+    query: Union[dataframe.DataFrame, series.Series],
     *,
     query_column_to_search: Optional[str] = None,
     top_k: Optional[int] = 10,
     distance_type: Literal["euclidean", "cosine"] = "euclidean",
     fraction_lists_to_search: Optional[float] = None,
     use_brute_force: bool = False,
-) -> bpd.DataFrame:
+) -> dataframe.DataFrame:
     """
     Conduct vector search to earch embeddings to find semantically similar entities.
 
@@ -258,11 +264,14 @@ def vector_search(
         raise ValueError(
             "You can't specify fraction_lists_to_search when use_brute_force is set to True."
         )
-    if isinstance(query, bpd.Series) and query_column_to_search is not None:
+    if (
+        isinstance(query, bigframes.series.Series)
+        and query_column_to_search is not None
+    ):
         raise ValueError(
             "You can't specify query_column_to_search when query is a Series."
         )
-    ## (TODO: ashleyxu. Support options in vector search.)
+    # TODO(ashleyxu): ashleyxu. Support options in vector search. b/344019989
     if fraction_lists_to_search is not None or use_brute_force is True:
         raise NotImplementedError(
             f"fraction_lists_to_search and use_brute_force is not supported. {constants.FEEDBACK_LINK}"
@@ -277,8 +286,16 @@ def vector_search(
         "use_brute_force": use_brute_force,
     }
 
-    df = utils.apply_sql(
-        query,
-        options,  # type:ignore
+    (query,) = utils.convert_to_dataframe(query)
+    sql_string, index_col_ids, index_labels = query._to_sql_query(include_index=True)
+
+    sql = bigframes.core.sql.create_vector_search_sql(
+        sql_string=sql_string, options=options  # type: ignore
     )
+    if index_col_ids is not None:
+        df = query._session.read_gbq(sql, index_col=index_col_ids)
+    else:
+        df = query._session.read_gbq(sql)
+    df.index.names = index_labels
+
     return df

bigframes/bigquery/utils.py

@@ -20,7 +20,7 @@
 import datetime
 import math
 import textwrap
-from typing import Iterable, TYPE_CHECKING
+from typing import Iterable, Mapping, TYPE_CHECKING, Union
 
 # Literals and identifiers matching this pattern can be unquoted
 unquoted = r"^[A-Za-z_][A-Za-z_0-9]*$"
@@ -169,3 +169,47 @@ def ordering_clause(
         part = f"`{ordering_expr.id}` {asc_desc} {null_clause}"
         parts.append(part)
     return f"ORDER BY {' ,'.join(parts)}"
+
+
+def create_vector_search_sql(
+    sql_string: str,
+    options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
+) -> str:
+    """Encode the VECTOR SEARCH statement for BigQuery Vector Search."""
+
+    base_table = options["base_table"]
+    column_to_search = options["column_to_search"]
+    distance_type = options["distance_type"]
+    top_k = options["top_k"]
+    query_column_to_search = options.get("query_column_to_search", None)
+
+    if query_column_to_search is not None:
+        query_str = f"""
+    SELECT
+        query.*,
+        base.*,
+        distance,
+    FROM VECTOR_SEARCH(
+        TABLE `{base_table}`,
+        "{column_to_search}",
+        ({sql_string}),
+        "{query_column_to_search}",
+        distance_type => "{distance_type}",
+        top_k => {top_k}
+    )
+    """
+    else:
+        query_str = f"""
+    SELECT
+        query.*,
+        base.*,
+        distance,
+    FROM VECTOR_SEARCH(
+        TABLE `{base_table}`,
+        "{column_to_search}",
+        ({sql_string}),
+        distance_type => "{distance_type}",
+        top_k => {top_k}
+    )
+    """
+    return query_str

bigframes/core/sql.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import bigframes.bigquery as bbq
+
+from bigframes.core import sql
 
 
 def test_create_vector_search_sql_simple():
@@ -39,7 +40,7 @@ def test_create_vector_search_sql_simple():
     )
     """
 
-    result_query = bbq.utils.create_vector_search_sql(
+    result_query = sql.create_vector_search_sql(
         sql_string, options  # type:ignore
     )
     assert result_query == expected_query
@@ -71,7 +72,7 @@ def test_create_vector_search_sql_query_column_to_search():
     )
     """
 
-    result_query = bbq.utils.create_vector_search_sql(
+    result_query = sql.create_vector_search_sql(
         sql_string, options  # type:ignore
     )
     assert result_query == expected_query