diff --git a/tests/integrations/test_utils.py b/tests/integrations/test_utils.py index dbe1998..502148d 100644 --- 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(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(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" + 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" + 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(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" + 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(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 index 2169dae..1784f72 100644 --- 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(?:\((\d+)\))?", column_type, re.IGNORECASE + r"VECTOR(?:\((\d+)\))?", column_type, re.IGNORECASE ) dimension = ( int(dimension_match.group(1))