yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2025-11-10 06:43:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add document count query

Browse Source
This commit is contained in:
Tanay Soni 2020-01-22 16:05:35 +01:00
parent ec7ad29d51
commit ecfd19d2d8
3 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
haystack/database/base.py
View File

@ -18,4 +18,8 @@ class BaseDocumentStore:
def get_document_ids_by_tag(self, tag): def get_document_ids_by_tag(self, tag):
pass pass
@abstractmethod
def get_document_count(self):
pass

4
haystack/database/sql.py
View File

@ -98,3 +98,7 @@ class SQLDocumentStore(BaseDocumentStore):
row = Document(name=doc["name"], text=doc["text"]) row = Document(name=doc["name"], text=doc["text"])
self.session.add(row) self.session.add(row)
self.session.commit() self.session.commit()
def get_document_count(self):
return self.session.query(Document).count()

3
haystack/indexing/io.py
View File

@ -19,11 +19,10 @@ def write_documents_to_db(datastore, document_dir, clean_func=None, only_empty_d
:return: None :return: None
""" """
file_paths = Path(document_dir).glob("**/*.txt") file_paths = Path(document_dir).glob("**/*.txt")
n_docs = 0
# check if db has already docs # check if db has already docs
if only_empty_db: if only_empty_db:
n_docs = len(datastore.get_all_documents()) n_docs = datastore.get_document_count()
if n_docs > 0: if n_docs > 0:
logger.info(f"Skip writing documents since DB already contains {n_docs} docs ... " logger.info(f"Skip writing documents since DB already contains {n_docs} docs ... "
"(Disable `only_empty_db`, if you want to add docs anyway.)") "(Disable `only_empty_db`, if you want to add docs anyway.)")