update tests

pingcap · Jul 10, 2024 · d8288fa · d8288fa
1 parent e8f37d0
commit d8288fa
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/tests/integrations/test_utils.py b/tests/integrations/test_utils.py
@@ -6,7 +6,7 @@
 
 def test_extract_info_from_column_definition():
     # Test case with dimension and distance metric
-    column_type = "VECTOR<FLOAT>(128)"
+    column_type = "VECTOR(128)"
     column_comment = "hnsw(distance=cosine)"
     expected_result = (128, "cosine")
     assert (
@@ -15,7 +15,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with dimension but no distance metric
-    column_type = "VECTOR<FLOAT>(256)"
+    column_type = "VECTOR(256)"
     column_comment = "some comment"
     expected_result = (256, None)
     assert (
@@ -24,7 +24,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with no dimension and no distance metric
-    column_type = "VECTOR<FLOAT>"
+    column_type = "VECTOR"
     column_comment = "another comment"
     expected_result = (None, None)
     assert (
@@ -33,7 +33,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with no dimension and no comment
-    column_type = "VECTOR<FLOAT>"
+    column_type = "VECTOR"
     column_comment = ""
     expected_result = (None, None)
     assert (
@@ -42,7 +42,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with dimension but no comment
-    column_type = "VECTOR<FLOAT>(256)"
+    column_type = "VECTOR(256)"
     column_comment = ""
     expected_result = (256, None)
     assert (
@@ -51,7 +51,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case without index type
-    column_type = "VECTOR<FLOAT>"
+    column_type = "VECTOR"
     column_comment = "distance=l2"
     expected_result = (None, "l2")
     assert (
@@ -60,7 +60,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with addition comment content
-    column_type = "VECTOR<FLOAT>(128)"
+    column_type = "VECTOR(128)"
     column_comment = "test, hnsw(distance=l2)"
     expected_result = (128, "l2")
     assert (

diff --git a/tidb_vector/integrations/utils.py b/tidb_vector/integrations/utils.py
@@ -90,7 +90,7 @@ def extract_info_from_column_definition(column_type, column_comment):
     """
     # Try to extract the dimension, which is optional.
     dimension_match = re.search(
-        r"VECTOR<FLOAT>(?:\((\d+)\))?", column_type, re.IGNORECASE
+        r"VECTOR(?:\((\d+)\))?", column_type, re.IGNORECASE
     )
     dimension = (
         int(dimension_match.group(1))