Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
15 changes: 7 additions & 8 deletions ydb/tests/datashard/vector_index/large/test_vector_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ def _create_table(self, table_path):
PRIMARY KEY(pk)
);
"""
self.query(create_table_sql, True)
self.query(create_table_sql)

def _create_index(
self, table_path, vector_type, vector_dimension, levels, clusters, distance=None, similarity=None
Expand Down Expand Up @@ -73,7 +73,7 @@ def _create_index(
);
"""
logger.info(create_index_sql)
self.query(create_index_sql, True)
self.query(create_index_sql)

def _upsert_values(self, table_path, vector_type, vector_dimension):
logger.info("Upsert values")
Expand All @@ -90,7 +90,7 @@ def _upsert_values(self, table_path, vector_type, vector_dimension):
UPSERT INTO `{table_path}` (pk, embedding)
VALUES {",".join(values)};
"""
self.query(upsert_sql, False)
self.query(upsert_sql)

def _select(self, table_path, vector_type, vector_dimension, distance, similarity):
if distance is not None:
Expand All @@ -110,20 +110,19 @@ def _select(self, table_path, vector_type, vector_dimension, distance, similarit
ORDER BY {target}(embedding, $Target) {order}
LIMIT {self.limit};
"""
return self.query(select_sql, False)
return self.query(select_sql)

def _select_top(self, table_path, vector_type, vector_dimension, distance, similarity):
logger.info("Select values from table")
result_set = self._select(
rows = self._select(
table_path=table_path,
vector_type=vector_type,
vector_dimension=vector_dimension,
distance=distance,
similarity=similarity,
)
assert len(result_set) != 0, "Query returned an empty set"
assert len(rows) != 0, "Query returned an empty set"

rows = result_set[0].rows
logger.info(f"Rows count {len(rows)}")

prev = 0.0 if distance is not None else 1.0
Expand Down Expand Up @@ -183,7 +182,7 @@ def test_vector_index(self):
similarity_data = ["cosine"] # "inner_product", "cosine"
vector_type_data = ["float", "int8"]
levels_data = [1, 3]
clusters_data = [1, 17]
clusters_data = [2, 17]
vector_dimension_data = [5]

for vector_type in vector_type_data:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ def setup_method(self):
self.rows_count = 1000
self.count_prefix = 5

def test_vecot_index_large_levels_and_clusters(self):
def test_vector_index_large_levels_and_clusters(self):
prefix_data = {"String": lambda i: f"{i}"}
vector = {"String": lambda i: f"{i}"}
all_types = {
Expand Down Expand Up @@ -52,8 +52,8 @@ def test_vecot_index_large_levels_and_clusters(self):
create_table_sql = create_table_sql_request(
table_name=self.table_name,
columns=columns,
pk_colums=pk_columns,
index_colums={},
pk_columns=pk_columns,
index_columns={},
unique="",
sync="",
)
Expand Down