parameters for kb index

Author: ea-rus

mindsdb_sql_parser/ast/mindsdb/knowledge_base.py

@@ -160,21 +160,34 @@ class CreateKnowledgeBaseIndex(ASTNode):
     """
     Create a new index in the knowledge base
     """
-    def __init__(self, name, *args, **kwargs):
+    def __init__(self, name, params=None, *args, **kwargs):
         """
         Args:
             name: Identifier -- name of the knowledge base
         """
         super().__init__(*args, **kwargs)
         self.name = name
+        if params is None:
+            params = {}
+        self.params = params
 
     def to_tree(self, *args, level=0, **kwargs):
         ind = indent(level)
-        out_str = f"{ind}CreateKnowledgeBaseIndex(name={self.name.to_string()})"
+        params_str = ""
+        if self.params:
+            params_str = f",\n{indent(level+1)}params={repr(self.params)}\n{ind}"
+
+        out_str = f"{ind}CreateKnowledgeBaseIndex(name={self.name.to_string()}{params_str})"
         return out_str
 
     def get_string(self, *args, **kwargs):
         out_str = f'CREATE INDEX ON KNOWLEDGE_BASE {self.name.to_string()}'
+        if self.params:
+            params_ar = [
+                f"{k}={repr(v)}"
+                for k, v in self.params.items()
+            ]
+            out_str += f" WITH ({', '.join(params_ar)})"
         return out_str
 
 

mindsdb_sql_parser/parser.py

@@ -157,8 +157,12 @@ def alter_kb(self, p):
         )
 
     @_('CREATE INDEX ON KNOWLEDGE_BASE identifier')
+    @_('CREATE INDEX ON KNOWLEDGE_BASE identifier WITH LPAREN kw_parameter_list RPAREN')
     def create_index(self, p):
-        return CreateKnowledgeBaseIndex(name=p.identifier)
+        params = None
+        if hasattr(p, 'kw_parameter_list'):
+            params = p.kw_parameter_list
+        return CreateKnowledgeBaseIndex(name=p.identifier, params=params)
 
     @_('DROP INDEX ON KNOWLEDGE_BASE identifier')
     def  drop_index(self, p):

tests/test_mindsdb/test_knowledgebase_index.py

@@ -2,6 +2,7 @@
 from mindsdb_sql_parser.ast.mindsdb.knowledge_base import CreateKnowledgeBaseIndex, DropKnowledgeBaseIndex
 from mindsdb_sql_parser.ast import *
 
+
 class TestKB:
 
     def test_create_knowledge_base_index(self):
@@ -16,6 +17,24 @@ def test_create_knowledge_base_index(self):
         assert str(ast).lower() == sql.lower()
         assert ast.to_tree() == expected_ast.to_tree()
 
+    def test_create_knowledge_base_index_with_params(self):
+        # create without select
+
+        sql = "CREATE INDEX ON KNOWLEDGE_BASE my_kb WITH (type='ivf', nlist=1000)"
+
+        ast = parse_sql(sql)
+        print(ast)
+        expected_ast = CreateKnowledgeBaseIndex(
+            name=Identifier("my_kb"),
+            params={
+                "type": "ivf",
+                "nlist": 1000,
+            }
+        )
+        assert str(ast) == sql
+        assert str(ast) == str(expected_ast)
+        assert ast.to_tree() == expected_ast.to_tree()
+
     def test_drop_knowledge_base_index(self):
         # create without select