add num_docs (#62)

* add num_docs

add num_docs

* Update test_inmemory_vectordb.py

* Update test_hnswlib_vectordb.py

tests/unit/test_hnswlib_vectordb.py

@@ -169,4 +169,11 @@ def test_hnswlib_vectordb_restore(docs_to_index, tmpdir):
  assert len(res.matches) == 10
  # assert res.id == res.matches[0].id
  # assert res.text == res.matches[0].text
- # assert res.scores[0] < 0.001 # some precision issues, should be 0
+ # assert res.scores[0] < 0.001 # some precision issues, should be 0
+
+def test_hnswlib_num_dos(tmpdir):
+ db = HNSWVectorDB[MyDoc](workspace=str(tmpdir))
+ doc_list = [MyDoc(text=f'toy doc {i}', embedding=np.random.rand(128)) for i in range(1000)]
+ db.index(inputs=DocList[MyDoc](doc_list))
+ x=db.num_docs()
+ assert x['num_docs']==1000

tests/unit/test_inmemory_vectordb.py

@@ -172,3 +172,10 @@ def test_inmemory_vectordb_restore(docs_to_index, tmpdir):
  assert res.id == res.matches[0].id
  assert res.text == res.matches[0].text
  assert res.scores[0] > 0.99 # some precision issues, should be 1
+
+def test_inmemory_num_dos(tmpdir):
+ db = InMemoryExactNNVectorDB[MyDoc](workspace=str(tmpdir))
+ doc_list = [MyDoc(text=f'toy doc {i}', embedding=np.random.rand(128)) for i in range(1000)]
+ db.index(inputs=DocList[MyDoc](doc_list))
+ x=db.num_docs()
+ assert x['num_docs']==1000

vectordb/db/base.py

@@ -227,6 +227,9 @@ async def _deploy():
  ret = asyncio.run(_deploy())
  return ret
 
+ def num_docs(self, **kwargs):
+ return self._executor.num_docs()
+
  @pass_kwargs_as_params
  @unify_input_output
  def index(self, docs: 'DocList[TSchema]', parameters: Optional[Dict] = None, **kwargs):

vectordb/db/executors/hnsw_indexer.py

@@ -105,7 +105,7 @@ async def async_update(self, docs, *args, **kwargs):
  return self.update(docs, *args, **kwargs)
 
  def num_docs(self, **kwargs):
- return {'num_docs': self._index.num_docs()}
+ return {'num_docs': self._indexer.num_docs()}
 
  def snapshot(self, snapshot_dir):
  # TODO: Maybe copy the work_dir to workspace if `handle` is False

vectordb/db/executors/inmemory_exact_indexer.py

@@ -71,7 +71,7 @@ def update(self, docs, *args, **kwargs):
  return self._index(docs)
 
  def num_docs(self, *args, **kwargs):
- return {'num_docs': self._index.num_docs()}
+ return {'num_docs': self._indexer.num_docs()}
 
  def snapshot(self, snapshot_dir):
  snapshot_file = f'{snapshot_dir}/index.bin'