Skip to content

fix: Update to GA syntax #160

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鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
averikitsch wants to merge 2 commits into
base: fix-vul
Choose a base branch
from
Open

fix: Update to GA syntax #160

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
14e47f2
fix: Update to GA syntax
averikitsch Aug 7, 2026
2b63a96
Delete tests/unit/test_vectorstore_sql.py
averikitsch Aug 7, 2026
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
94 changes: 60 additions & 34 deletions src/langchain_google_cloud_sql_mysql/vectorstore.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from __future__ import annotations

import json
import warnings
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union

import numpy as np
Expand All @@ -27,6 +28,7 @@
from .indexes import (
DEFAULT_QUERY_OPTIONS,
DistanceMeasure,
IndexType,
QueryOptions,
SearchType,
VectorIndex,
Expand Down Expand Up @@ -244,42 +246,51 @@ def apply_vector_index(self, vector_index: VectorIndex):
# Construct the default index name
if not vector_index.name:
vector_index.name = f"{self.table_name}_{DEFAULT_INDEX_NAME_SUFFIX}"
query_template = f"CALL mysql.create_vector_index('{vector_index.name}', '{self.db_name}.{self.table_name}', '{self.embedding_column}', '{{}}');"
self.__exec_apply_vector_index(query_template, vector_index)
index_name = (
vector_index.name.split(".")[1]
if "." in vector_index.name
else vector_index.name
)

distance_measure = (
vector_index.distance_measure
or self.query_options.distance_measure
or DistanceMeasure.L2_SQUARED
)

options = ["QUANTIZER=SQ8", f"DISTANCE_MEASURE={distance_measure.value}"]
if vector_index.num_partitions:
options.append(f"NUM_LEAVES={vector_index.num_partitions}")

options_str = " ".join(options)
stmt = f"CREATE VECTOR INDEX `{index_name}` ON `{self.table_name}` (`{self.embedding_column}`) USING SCANN {options_str};"
self.engine._execute_outside_tx(stmt)
# After applying an index to the table, set the query option search type to be ANN
self.query_options.search_type = SearchType.ANN
self.query_options.distance_measure = distance_measure

def alter_vector_index(self, vector_index: VectorIndex):
existing_index_name = self._get_vector_index_name()
if not existing_index_name:
raise ValueError("No existing vector index found.")
existing_short_name = (
existing_index_name.split(".")[1]
if "." in existing_index_name
else existing_index_name
)
if not vector_index.name:
vector_index.name = existing_index_name.split(".")[1]
if existing_index_name.split(".")[1] != vector_index.name:
vector_index.name = existing_short_name
req_short_name = (
vector_index.name.split(".")[1]
if "." in vector_index.name
else vector_index.name
)
if existing_short_name != req_short_name:
raise ValueError(
f"Existing index name {existing_index_name} does not match the new index name {vector_index.name}."
)
query_template = (
f"CALL mysql.alter_vector_index('{existing_index_name}', '{{}}');"
)
self.__exec_apply_vector_index(query_template, vector_index)

def __exec_apply_vector_index(self, query_template: str, vector_index: VectorIndex):
index_options = []
if vector_index.index_type:
index_options.append(f"index_type={vector_index.index_type.value}")
if vector_index.distance_measure:
index_options.append(
f"distance_measure={vector_index.distance_measure.value}"
)
if vector_index.num_partitions:
index_options.append(f"num_partitions={vector_index.num_partitions}")
if vector_index.num_neighbors:
index_options.append(f"num_neighbors={vector_index.num_neighbors}")
index_options_query = ",".join(index_options)

stmt = query_template.format(index_options_query)
self.engine._execute_outside_tx(stmt)
self.drop_vector_index()
self.apply_vector_index(vector_index)

def _get_vector_index_name(self):
query = "SELECT index_name FROM mysql.vector_indexes WHERE table_name=:table_name;"
Expand All @@ -295,8 +306,13 @@ def _get_vector_index_name(self):
def drop_vector_index(self):
existing_index_name = self._get_vector_index_name()
if existing_index_name:
index_name = (
existing_index_name.split(".")[1]
if "." in existing_index_name
else existing_index_name
)
self.engine._execute_outside_tx(
f"CALL mysql.drop_vector_index('{existing_index_name}');"
f"DROP INDEX `{index_name}` ON `{self.table_name}`;"
)
self.query_options.search_type = SearchType.KNN
return existing_index_name
Expand Down Expand Up @@ -662,24 +678,34 @@ def _query_collection(
if query_options.num_partitions and query_options.search_type == SearchType.KNN:
raise ValueError("num_partitions is not supported for the search type KNN")

k = k if k else query_options.num_neighbors
k = k if k is not None else query_options.num_neighbors
distance_function = (
f"{query_options.distance_measure.value}_distance"
if query_options.distance_measure != DistanceMeasure.DOT_PRODUCT
else query_options.distance_measure.value
)
bind_params: Dict[str, Any] = {"embedding": str(embedding)}
if query_options.search_type == SearchType.KNN:
filter = f"WHERE {filter}" if filter else ""
stmt = f"SELECT {column_query}, {distance_function}({self.embedding_column}, string_to_vector(:embedding)) AS distance FROM `{self.table_name}` {filter} ORDER BY distance LIMIT {k};"

search_type = query_options.search_type
if search_type == SearchType.ANN:
if not k or k <= 0:
warnings.warn("A LIMIT is required for ANN search. Falling back to KNN search.")
search_type = SearchType.KNN
elif not query_options.distance_measure:
raise ValueError("Specifying the distance measure is required for ANN search.")

if search_type == SearchType.KNN:
where_clause = f"WHERE {filter}" if filter else ""
stmt = f"SELECT {column_query}, {distance_function}({self.embedding_column}, string_to_vector(:embedding)) AS distance FROM `{self.table_name}` {where_clause} ORDER BY distance LIMIT {k};"
else:
filter = f"AND {filter}" if filter else ""
num_partitions = (
f",num_partitions={query_options.num_partitions}"
where_clause = f"WHERE {filter}" if filter else ""
num_leaves = (
f",num_leaves_to_search={query_options.num_partitions}"
if query_options.num_partitions
else ""
)
stmt = f"SELECT {column_query}, {distance_function}({self.embedding_column}, string_to_vector(:embedding)) AS distance FROM `{self.table_name}` WHERE NEAREST({self.embedding_column}) TO (string_to_vector(:embedding), 'num_neighbors={k}{num_partitions}') {filter} ORDER BY distance;"
ann_options = f"distance_measure={query_options.distance_measure.value}{num_leaves}"
stmt = f"SELECT {column_query}, APPROX_DISTANCE({self.embedding_column}, string_to_vector(:embedding), '{ann_options}') AS distance FROM `{self.table_name}` {where_clause} ORDER BY distance LIMIT {k};"

# return self.engine._fetch(stmt)
if map_results:
Expand Down
Loading