fix: make check_vector_column to use the proper method for identifying vector columns (#126)

* fix: check_vector_column has an error way to check the vector column

* drop table should check_first by default

* build: bump pytidb version to 0.0.8.post2

---------

Co-authored-by: Mini256 <minianter@foxmail.com>

Author: Icemap

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pytidb"
-version = "0.0.8.post1"
+version = "0.0.8.post2"
 description = "A Python library for TiDB."
 readme = "README.md"
 requires-python = ">=3.10"

pytidb/client.py

@@ -148,7 +148,7 @@ def drop_table(self, table_name: str):
         table = sqlalchemy.Table(
             table_name, Base.metadata, autoload_with=self._db_engine
         )
-        return table.drop(self._db_engine)
+        return table.drop(self._db_engine, checkfirst=True)
 
     # Raw SQL API
 

pytidb/utils.py

@@ -84,12 +84,15 @@ def filter_vector_columns(columns: Dict) -> List[Column]:
 
 
 def check_vector_column(columns: Dict, column_name: str) -> Optional[str]:
+    if not isinstance(column_name, str):
+        raise ValueError(f"Invalid vector column name: {column_name}")
+
     if column_name not in columns:
         raise ValueError(f"Non-exists vector column: {column_name}")
 
     vector_column = columns[column_name]
-    if vector_column.type != VectorType:
-        raise ValueError(f"Invalid vector column: {vector_column}")
+    if not isinstance(vector_column.type, VectorType):
+        raise ValueError(f"Invalid vector column: {column_name}")
 
     return vector_column
 
@@ -140,4 +143,8 @@ def get_row_id_from_row(row: Row, table: Table) -> Optional[RowKeyType]:
 
 
 def get_index_type(index: Index) -> str:
-    return index.dialect_kwargs.get("mysql_prefix", "").lower()
+    dialect_kwargs = getattr(index, "dialect_kwargs", None)
+    if dialect_kwargs is None:
+        return ""
+    mysql_prefix = dialect_kwargs.get("mysql_prefix", "")
+    return mysql_prefix.lower() if mysql_prefix else ""

tests/test_search_vector.py

@@ -212,3 +212,66 @@ def test_rerank(vector_table: Table, reranker: BaseReranker):
     assert len(reranked_results) > 0
     assert reranked_results[0]["text"] == "bar"
     assert reranked_results[0]["_score"] > 0
+
+
+def test_with_multi_vector_fields(client: TiDBClient):
+    class ChunkWithMultiVec(TableModel):
+        __tablename__ = "test_vector_search_multi_vec"
+        id: int = Field(None, primary_key=True)
+        title: str = Field(None)
+        title_vec: list[float] = Field(sa_column=Column(Vector(3)))
+        body: str = Field(None)
+        body_vec: list[float] = Field(sa_column=Column(Vector(3)))
+
+    tbl = client.create_table(schema=ChunkWithMultiVec, mode="overwrite")
+    tbl.bulk_insert(
+        [
+            ChunkWithMultiVec(
+                id=1, title="tree", title_vec=[4, 5, 6], body="cat", body_vec=[1, 2, 3]
+            ),
+            ChunkWithMultiVec(
+                id=2, title="grass", title_vec=[1, 2, 3], body="dog", body_vec=[7, 8, 9]
+            ),
+            ChunkWithMultiVec(
+                id=3, title="leaf", title_vec=[7, 8, 9], body="bird", body_vec=[4, 5, 6]
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="more than two vector columns"):
+        tbl.search([1, 2, 3], search_type="vector").limit(3).to_list()
+
+    with pytest.raises(ValueError, match="Invalid vector column"):
+        tbl.search([1, 2, 3], search_type="vector").vector_column("title").limit(3)
+
+    with pytest.raises(ValueError, match="Non-exists vector column"):
+        tbl.search([1, 2, 3], search_type="vector").vector_column(
+            "non_exist_column"
+        ).limit(3)
+
+    with pytest.raises(ValueError, match="Invalid vector column name"):
+        tbl.search([1, 2, 3], search_type="vector").vector_column(None).limit(3)
+
+    results = (
+        tbl.search([1, 2, 3], search_type="vector")
+        .vector_column("title_vec")
+        .limit(3)
+        .to_list()
+    )
+    assert len(results) == 3
+    assert results[0]["id"] == 2
+    assert results[0]["title"] == "grass"
+    assert results[0]["_distance"] == 0
+    assert results[0]["_score"] == 1
+
+    results = (
+        tbl.search([1, 2, 3], search_type="vector")
+        .vector_column("body_vec")
+        .limit(3)
+        .to_list()
+    )
+    assert len(results) == 3
+    assert results[0]["id"] == 1
+    assert results[0]["body"] == "cat"
+    assert results[0]["_distance"] == 0
+    assert results[0]["_score"] == 1