support pgvector operators

Author: ea-rus

mindsdb_sql_parser/__about__.py

@@ -1,6 +1,6 @@
 __title__ = 'mindsdb_sql_parser'
 __package_name__ = 'mindsdb_sql_parser'
-__version__ = '0.3.1'
+__version__ = '0.4.0'
 __description__ = "Mindsdb SQL parser"
 __email__ = "[email protected]"
 __author__ = 'MindsDB Inc'

mindsdb_sql_parser/lexer.py

@@ -73,6 +73,7 @@ class MindsDBLexer(Lexer):
         EQUALS, NEQUALS, GREATER, GEQ, LESS, LEQ,
         AND, OR, NOT, IS, IS_NOT, TYPECAST,
         IN, NOT_IN, LIKE, NOT_LIKE, CONCAT, BETWEEN, WINDOW, OVER, PARTITION_BY,
+        VECT_L2, VECT_L1, VECT_INNER, VECT_COS, VECT_HAMM, VECT_JACC,
         JSON_GET, JSON_GET_STR, INTERVAL,
 
         # Data types
@@ -272,6 +273,13 @@ class MindsDBLexer(Lexer):
     SEMICOLON = r'\;'
 
     # Operators
+    VECT_L2 = r'<->'
+    VECT_L1 = r'<\+>'
+    VECT_INNER = r'<#>'
+    VECT_COS = r'<=>'
+    VECT_HAMM = r'<~>'
+    VECT_JACC = r'<%>'
+
     JSON_GET = r'->'
     JSON_GET_STR = r'->>'
     PLUS = r'\+'

mindsdb_sql_parser/parser.py

@@ -1572,6 +1572,12 @@ def star(self, p):
        'expr LIKE expr',
        'expr NOT_LIKE expr',
        'expr CONCAT expr',
+       'expr VECT_L2 expr',
+       'expr VECT_L1 expr',
+       'expr VECT_INNER expr',
+       'expr VECT_COS expr',
+       'expr VECT_HAMM expr',
+       'expr VECT_JACC expr',
        'expr JSON_GET constant',
        'expr JSON_GET_STR constant',
        'expr NOT_IN expr',

tests/test_mindsdb/test_selects.py

@@ -1,3 +1,4 @@
+import pytest
 from mindsdb_sql_parser import parse_sql
 from mindsdb_sql_parser.ast import *
 from mindsdb_sql_parser.utils import JoinType
@@ -188,6 +189,26 @@ def test_json(self):
         assert ast.to_tree() == expected_ast.to_tree()
         assert str(ast) == str(expected_ast)
 
+    @pytest.mark.parametrize('op', ['<->', '<+>', '<#>', '<=>', '<~>', '<%>'])
+    def test_vector(self, op):
+        sql = f"SELECT * from TAB1 WHERE embedding {op} '[3,1,2]'"
+
+        ast = parse_sql(sql)
+        expected_ast = Select(
+            targets=[Star()],
+            where=BinaryOperation(
+                op=op,
+                args=[
+                    Identifier('embedding'),
+                    Constant('[3,1,2]')
+                ]
+            ),
+            from_table=Identifier(parts=['TAB1']),
+        )
+
+        assert ast.to_tree() == expected_ast.to_tree()
+        assert str(ast) == str(expected_ast)
+
     def test_match(self):
         sql = "SELECT a~b, a!~c from TAB1"