Merge pull request Event Theme Visual Example Product Launching from mindsdb/kb-index-params

Parameters for create knowledge base index

Author: Timeline Template For Planning

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