Knn python tests (#30328)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: azevaykin

ydb/tests/compatibility/test_vector_index.py

@@ -85,6 +85,30 @@ def _write_data(self, name, vector_type, table_name):
         with ydb.QuerySessionPool(self.driver) as session_pool:
             session_pool.execute_with_retries(sql_upsert)
 
+    def _get_knn_queries(self):
+        """Get KNN search queries without vector index (brute force with pushdown)."""
+        queries = []
+        for prefixed in ['', '_pfx']:
+            for vector_type in self.vector_types.keys():
+                for distance in self.targets.keys():
+                    for distance_func in self.targets[distance].keys():
+                        table_name = f"{vector_type}_{distance}_{distance_func}{prefixed}"
+                        order = "ASC" if distance != "similarity" else "DESC"
+                        vector = self.get_vector(f"{vector_type}Vector", 1)
+                        where = ""
+                        if prefixed:
+                            where = "WHERE user=1"
+                        queries.append([
+                            True, f"""
+                            $Target = {self.vector_types[vector_type]}(Cast([{vector}] AS List<{vector_type}>));
+                            SELECT key, vec, {self.targets[distance][distance_func]}(vec, $Target) as target
+                            FROM `{table_name}`
+                            {where}
+                            ORDER BY {self.targets[distance][distance_func]}(vec, $Target) {order}
+                            LIMIT {self.rows_count};"""
+                        ])
+        return queries
+
     def _get_queries(self):
         queries = []
         for prefixed in ['', '_pfx']:
@@ -159,6 +183,12 @@ def select_from_index_without_roll(self):
         queries = self._get_queries()
         self._do_queries(queries)
 
+    def knn_search(self):
+        """Perform KNN search without vector index during rolling upgrade/downgrade."""
+        queries = self._get_knn_queries()
+        for _ in self.roll():
+            self._do_queries(queries)
+
     def create_table(self, table_name):
         query = f"""
                 CREATE TABLE {table_name} (
@@ -189,5 +219,6 @@ def test_vector_index(self):
                             target=f"{distance}={distance_func}",
                             prefixed=prefixed,
                         )
+        self.knn_search()
         self.wait_index_ready()
         self.select_from_index()

ydb/tests/datashard/knn/test_knn.py

@@ -0,0 +1,171 @@
+from ydb.tests.datashard.lib.vector_base import VectorBase
+
+
+class TestKnn(VectorBase):
+    table_name = "TestTable"
+    target_table_name = "TargetVectors"
+    # Target embedding: 0x67, 0x71 (103, 113)
+    target_embedding_hex = "677102"
+
+    @classmethod
+    def setup_class(cls):
+        super().setup_class()
+        cls._create_tables()
+
+    @classmethod
+    def teardown_class(cls):
+        cls._cleanup_tables()
+        super().teardown_class()
+
+    @classmethod
+    def _cleanup_tables(cls):
+        try:
+            cls.pool.execute_with_retries(f"DROP TABLE `{cls.table_name}`")
+        except Exception:
+            pass
+        try:
+            cls.pool.execute_with_retries(f"DROP TABLE `{cls.target_table_name}`")
+        except Exception:
+            pass
+
+    @classmethod
+    def _create_tables(cls):
+        cls._cleanup_tables()
+
+        cls.pool.execute_with_retries(f"""
+            CREATE TABLE `{cls.table_name}` (
+                pk Int64 NOT NULL,
+                emb String NOT NULL,
+                data String NOT NULL,
+                PRIMARY KEY (pk)
+            );
+        """)
+
+        cls.pool.execute_with_retries(f"""
+            UPSERT INTO `{cls.table_name}` (pk, emb, data) VALUES
+            (0, Unwrap(String::HexDecode("033002")), "0"),
+            (1, Unwrap(String::HexDecode("133102")), "1"),
+            (2, Unwrap(String::HexDecode("233202")), "2"),
+            (3, Unwrap(String::HexDecode("533302")), "3"),
+            (4, Unwrap(String::HexDecode("433402")), "4"),
+            (5, Unwrap(String::HexDecode("506002")), "5"),
+            (6, Unwrap(String::HexDecode("611102")), "6"),
+            (7, Unwrap(String::HexDecode("126202")), "7"),
+            (8, Unwrap(String::HexDecode("757602")), "8"),
+            (9, Unwrap(String::HexDecode("767602")), "9");
+        """)
+
+        cls.pool.execute_with_retries(f"""
+            CREATE TABLE `{cls.target_table_name}` (
+                id Int64 NOT NULL,
+                target_emb String,
+                PRIMARY KEY (id)
+            )
+        """)
+
+        cls.pool.execute_with_retries(f"""
+            UPSERT INTO `{cls.target_table_name}` (id, target_emb) VALUES
+            (1, String::HexDecode("{cls.target_embedding_hex}"))
+        """)
+
+    def test_knn_cosine_distance(self):
+        """Test KNN with CosineDistance."""
+        result = self.query(f"""
+            $TargetEmbedding = String::HexDecode("{self.target_embedding_hex}");
+            SELECT pk FROM `{self.table_name}`
+            ORDER BY Knn::CosineDistance(emb, $TargetEmbedding)
+            LIMIT 3
+        """)
+        assert len(result) == 3
+
+    def test_knn_cosine_similarity(self):
+        """Test KNN with CosineSimilarity (DESC order)."""
+        result = self.query(f"""
+            $TargetEmbedding = String::HexDecode("{self.target_embedding_hex}");
+            SELECT pk FROM `{self.table_name}`
+            ORDER BY Knn::CosineSimilarity(emb, $TargetEmbedding) DESC
+            LIMIT 3
+        """)
+        assert len(result) == 3
+
+    def test_knn_inner_product_similarity(self):
+        """Test KNN with InnerProductSimilarity (DESC order)."""
+        result = self.query(f"""
+            $TargetEmbedding = String::HexDecode("{self.target_embedding_hex}");
+            SELECT pk FROM `{self.table_name}`
+            ORDER BY Knn::InnerProductSimilarity(emb, $TargetEmbedding) DESC
+            LIMIT 3
+        """)
+        assert len(result) == 3
+
+    def test_knn_manhattan_distance(self):
+        """Test KNN with ManhattanDistance."""
+        result = self.query(f"""
+            $TargetEmbedding = String::HexDecode("{self.target_embedding_hex}");
+            SELECT pk FROM `{self.table_name}`
+            ORDER BY Knn::ManhattanDistance(emb, $TargetEmbedding)
+            LIMIT 3
+        """)
+        assert len(result) == 3
+
+    def test_knn_euclidean_distance(self):
+        """Test KNN with EuclideanDistance."""
+        result = self.query(f"""
+            $TargetEmbedding = String::HexDecode("{self.target_embedding_hex}");
+            SELECT pk FROM `{self.table_name}`
+            ORDER BY Knn::EuclideanDistance(emb, $TargetEmbedding)
+            LIMIT 3
+        """)
+        assert len(result) == 3
+
+    def test_knn_verify_results(self):
+        """
+        Verify actual results - check that top 3 PKs are correct.
+        Target vector is 0x67, 0x71 (103, 113)
+        Expected cosine distances:
+            pk=8 (117, 118): 0.000882 - closest
+            pk=5 (80, 96):   0.000985
+            pk=9 (118, 118): 0.001070
+        """
+        result = self.query(f"""
+            $TargetEmbedding = String::HexDecode("{self.target_embedding_hex}");
+            SELECT pk, Knn::CosineDistance(emb, $TargetEmbedding) AS distance
+            FROM `{self.table_name}`
+            ORDER BY distance
+            LIMIT 3
+        """)
+
+        assert len(result) == 3
+        pks = [row['pk'] for row in result]
+        assert pks == [8, 5, 9], f"Expected PKs [8, 5, 9], got {pks}"
+
+    def test_knn_two_stage_query(self):
+        """
+        Test two-stage query: quantized search followed by full precision reranking.
+        """
+        result = self.query(f"""
+            $TargetEmbedding = String::HexDecode("{self.target_embedding_hex}");
+
+            $Pks = SELECT pk
+            FROM `{self.table_name}`
+            ORDER BY Knn::CosineDistance(emb, $TargetEmbedding)
+            LIMIT 3;
+
+            SELECT pk, Knn::CosineDistance(emb, $TargetEmbedding) AS distance
+            FROM `{self.table_name}`
+            WHERE pk IN $Pks
+            ORDER BY distance
+            LIMIT 3;
+        """)
+        assert len(result) == 3
+
+    def test_knn_subquery_target(self):
+        """Test KNN with target vector from subquery."""
+        result = self.query(f"""
+            $TargetEmbedding = (SELECT target_emb FROM `{self.target_table_name}` WHERE id = 1);
+            SELECT pk, Knn::CosineDistance(emb, $TargetEmbedding) AS distance
+            FROM `{self.table_name}`
+            ORDER BY distance
+            LIMIT 3
+        """)
+        assert len(result) == 3

ydb/tests/datashard/knn/ya.make

@@ -0,0 +1,27 @@
+PY3TEST()
+INCLUDE(${ARCADIA_ROOT}/ydb/tests/harness_dep.inc)
+
+FORK_SUBTESTS()
+SPLIT_FACTOR(10)
+
+IF (SANITIZER_TYPE)
+    SIZE(LARGE)
+    INCLUDE(${ARCADIA_ROOT}/ydb/tests/large.inc)
+ELSE()
+    SIZE(MEDIUM)
+ENDIF()
+
+TEST_SRCS(
+    test_knn.py
+)
+
+PEERDIR(
+    ydb/tests/datashard/lib
+)
+
+DEPENDS(
+    ydb/apps/ydb
+)
+
+END()
+

ydb/tests/datashard/ya.make

@@ -4,6 +4,7 @@ RECURSE(
     dml
     dump_restore
     copy_table
+    knn
     lib
     partitioning
     parametrized_queries

ydb/tests/stress/oltp_workload/workload/type/vector_index.py

@@ -143,6 +143,61 @@ def _select(self, index_name, table_path, vector_type, vector_dimension, distanc
         """
         return self.client.query(select_sql, False)
 
+    def _knn_search(self, table_path, vector_type, vector_dimension, distance, similarity, prefixed=False):
+        """KNN search without vector index (brute force with pushdown)."""
+        if distance is not None:
+            target = targets["distance"][distance]
+        else:
+            target = targets["similarity"][similarity]
+        order = "ASC" if distance is not None else "DESC"
+        vector = self._get_random_vector(vector_type, vector_dimension)
+        converter = to_binary_string_converters[vector_type]
+        name = converter.name
+        data_type = converter.data_type
+        if prefixed:
+            where = "WHERE user=1"
+        else:
+            where = ""
+        select_sql = f"""
+            $Target = {name}(Cast([{vector}] AS List<{data_type}>));
+            SELECT pk, embedding, {target}(embedding, $Target) as target
+            FROM `{table_path}`
+            {where}
+            ORDER BY {target}(embedding, $Target) {order}
+            LIMIT {self.limit};
+        """
+        return self.client.query(select_sql, False)
+
+    def _knn_search_check(self, table_path, vector_type, vector_dimension, distance, similarity, prefixed=False):
+        """Perform KNN search and verify results are ordered correctly."""
+        logger.info(f"KNN search: vector_type={vector_type}, distance={distance}, similarity={similarity}")
+        result_set = self._knn_search(
+            table_path=table_path,
+            vector_type=vector_type,
+            vector_dimension=vector_dimension,
+            distance=distance,
+            similarity=similarity,
+            prefixed=prefixed,
+        )
+        if len(result_set) == 0:
+            raise Exception(
+                f"KNN query returned an empty set for vector_type={vector_type}, distance={distance}, similarity={similarity}, prefixed={prefixed}"
+            )
+        rows = result_set[0].rows
+        logger.info(f"KNN search returned {len(rows)} rows")
+
+        if len(rows) > 1:
+            prev = rows[0]['target']
+            for row in rows[1:]:
+                cur = row['target']
+                condition = prev <= cur if distance is not None else prev >= cur
+                if not condition:
+                    raise Exception(
+                        f"KNN results not properly ordered: prev={prev}, cur={cur}"
+                    )
+                prev = cur
+        logger.info("KNN search completed successfully")
+
     def _select_top(self, index_name, table_path, vector_type, vector_dimension, distance, similarity, prefixed=False):
         logger.info("Select values from table")
         result_set = self._select(
@@ -196,6 +251,15 @@ def _wait_index_ready(self, index_name, table_path, vector_type, vector_dimensio
         raise Exception("Error getting index status")
 
     def _check_loop(self, table_path, vector_type, vector_dimension, levels, clusters, distance=None, similarity=None, prefixed=False):
+        self._knn_search_check(
+            table_path=table_path,
+            vector_type=vector_type,
+            vector_dimension=vector_dimension,
+            distance=distance,
+            similarity=similarity,
+            prefixed=prefixed,
+        )
+
         index_name = f"{self.index_name_prefix}_{vector_type}_{vector_dimension}_{levels}_{clusters}_{distance}_{similarity}_{str(prefixed)}"
         self._create_index(
             index_name=index_name,