feat: add push endpoint

requirements.txt

@@ -1,2 +1,3 @@
 jina>=3.19.0
 docarray[hnswlib]>=0.34.0
+more-itertools

vectordb/client/client.py

@@ -54,6 +54,14 @@ def __init__(self, address, reverse_order=False):
  def index(self, *args, **kwargs):
  return self._client.index(*args, **kwargs)
 
+ @unify_input_output
+ @pass_kwargs_as_params
+ def push(self, *args, **kwargs):
+ return self._client.post(on='/push', *args, **kwarg)
+
+ def build(self, *args, **kwargs):
+ return self._client.post(on='/build', *args, **kwarg)
+
  @sort_matches_by_scores
  @unify_input_output
  @pass_kwargs_as_params

vectordb/db/base.py

@@ -85,7 +85,7 @@ def _get_jina_object(cls,
  ServedExecutor = type(f'{executor_cls_name.replace("[", "").replace("]", "")}', (cls._executor_cls,),
  {})
  uses = ServedExecutor
- polling = {'/index': 'ANY', '/search': 'ALL', '/update': 'ALL', '/delete': 'ALL'}
+ polling = {'/index': 'ANY', '/search': 'ALL', '/update': 'ALL', '/delete': 'ALL', '/push': 'ANY', '/build': 'ALL'}
  if to_deploy and replicas > 1:
  import warnings
  warnings.warn(
@@ -232,6 +232,14 @@ async def _deploy():
  def index(self, docs: 'DocList[TSchema]', parameters: Optional[Dict] = None, **kwargs):
  return self._executor.index(docs, parameters)
 
+ @pass_kwargs_as_params
+ @unify_input_output
+ def push(self, docs: 'DocList[TSchema]', parameters: Optional[Dict] = None, **kwargs):
+ return self._executor.push(docs, parameters)
+
+ def build(self, **kwargs):
+ return self._executor.push(**kwargs)
+
  @pass_kwargs_as_params
  @unify_input_output
  def update(self, docs: 'DocList[TSchema]', parameters: Optional[Dict] = None, **kwargs):

vectordb/db/executors/hnsw_indexer.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+from more_itertools import chunked
 from vectordb.db.executors.typed_executor import TypedExecutor
 from jina.serve.executors.decorators import requests, write
 
@@ -24,6 +25,7 @@ def __init__(self,
  num_threads=1,
  *args, **kwargs):
  from docarray.index import HnswDocumentIndex
+ from docarray import DocList
  super().__init__(*args, **kwargs)
  workspace = self.workspace.replace('[', '_').replace(']', '_')
  self.work_dir = f'{workspace}' if self.handle_persistence else f'{workspace}/{"".join(random.choice(string.ascii_lowercase) for _ in range(5))}'
@@ -35,6 +37,7 @@ def __init__(self,
  'M': M,
  'allow_replace_deleted': allow_replace_deleted,
  'num_threads': num_threads})
+ self._pushed_docs = DocList[self._input_schema]()
  db_conf.work_dir = self.work_dir
  self._indexer = HnswDocumentIndex[self._input_schema](db_config=db_conf)
 
@@ -44,6 +47,10 @@ def _index(self, docs, *args, **kwargs):
 
  def index(self, docs, *args, **kwargs):
  self.logger.debug(f'Index {len(docs)}')
+ if len(self._pushed_docs) > 0:
+ self.logger.debug(f'{len(self._pushed_docs)} were waiting to be indexed. Indexing them')
+ self._index(self._pushed_docs)
+ self._pushed_docs.clear()
  return self._index(docs)
 
  @write
@@ -88,6 +95,19 @@ def delete(self, docs, *args, **kwargs):
  self.logger.debug(f'Delete')
  return self._delete(docs, *args, **kwargs)
 
+ @write
+ @requests(on='/push')
+ def push(self, docs, *args, **kwargs):
+ self.logger.debug(f'Push {len(docs)}')
+ self._pushed_docs.extend(docs)
+
+ @write
+ @requests(on='/build')
+ def build(self, *args, **kwargs):
+ self.logger.debug(f'Building index with {len(self._pushed_docs)} pushed docs')
+ self._index(self._pushed_docs, *args, **kwargs)
+ self._pushed_docs.clear()
+
  @write
  @requests(on='/delete')
  async def async_delete(self, docs, *args, **kwargs):

vectordb/db/executors/inmemory_exact_indexer.py

@@ -26,6 +26,16 @@ def index(self, docs, *args, **kwargs):
  self.logger.debug(f'Index {len(docs)}')
  return self._index(docs, *args, **kwargs)
 
+ @write
+ @requests(on='/push')
+ def push(self, docs, *args, **kwargs):
+ self.logger.debug(f'Push {len(docs)}')
+ return self._index(docs, *args, **kwargs)
+
+ @requests(on='/build')
+ def build(self, *args, **kwargs):
+ self.logger.debug(f'Build call has no effect in InMemoryExactNNIndexer.')
+
  def _search(self, docs, parameters, *args, **kwargs):
  from docarray import DocList
  res = DocList[self._output_schema]()

vectordb/db/executors/typed_executor.py

@@ -14,7 +14,7 @@
 InputSchema = TypeVar('InputSchema', bound='BaseDoc')
 OutputSchema = TypeVar('OutputSchema', bound='BaseDoc')
 
-methods = ['/index', '/update', '/delete', '/search']
+methods = ['/index', '/update', '/delete', '/search', '/push', '/build']
 
 
 class TypedExecutor(Executor, Generic[InputSchema, OutputSchema]):