Extend meta data support for SQLDocumentStore

haystack/document_stores/filter_utils.py

@@ -1,6 +1,12 @@
 from typing import Union, List, Dict
 from abc import ABC, abstractmethod
 from collections import defaultdict
+from functools import reduce
+
+from haystack.utils.import_utils import safe_import
+from haystack.document_stores.sql import MetaDocumentORM
+from sqlalchemy.sql import select
+from sqlalchemy import and_, or_
 
 
 def nested_defaultdict():
@@ -122,6 +128,13 @@ def convert_to_elasticsearch(self):
         """
         pass
 
+    @abstractmethod
+    def convert_to_sql(self):
+        """
+        Converts the LogicalFilterClause instance to an SQL filter.
+        """
+        pass
+
     def _merge_es_range_queries(self, conditions: List[Dict]) -> List[Dict]:
         """
         Merges Elasticsearch range queries that perform on the same metadata field.
@@ -187,6 +200,13 @@ def convert_to_elasticsearch(self):
         """
         pass
 
+    @abstractmethod
+    def convert_to_sql(self):
+        """
+        Converts the ComparisonOperation instance to an SQL filter.
+        """
+        pass
+
 
 class NotOperation(LogicalFilterClause):
     """
@@ -197,7 +217,11 @@ def convert_to_elasticsearch(self):
         conditions = [condition.convert_to_elasticsearch() for condition in self.conditions]
         conditions = self._merge_es_range_queries(conditions)
         return {"bool": {"must_not": conditions}}
-
+
+    def convert_to_sql(self):
+        conditions = [MetaDocumentORM.document_id.in_(condition.convert_to_sql()) for condition in self.conditions]
+        return select(MetaDocumentORM.document_id).filter(~or_(*conditions))
+
 
 class AndOperation(LogicalFilterClause):
     """
@@ -208,7 +232,11 @@ def convert_to_elasticsearch(self):
         conditions = [condition.convert_to_elasticsearch() for condition in self.conditions]
         conditions = self._merge_es_range_queries(conditions)
         return {"bool": {"must": conditions}}
-
+
+    def convert_to_sql(self):
+        conditions = [MetaDocumentORM.document_id.in_(condition.convert_to_sql()) for condition in self.conditions]
+        return select(MetaDocumentORM.document_id).filter(and_(*conditions))
+
 
 class OrOperation(LogicalFilterClause):
     """
@@ -219,6 +247,10 @@ def convert_to_elasticsearch(self):
         conditions = [condition.convert_to_elasticsearch() for condition in self.conditions]
         conditions = self._merge_es_range_queries(conditions)
         return {"bool": {"should": conditions}}
+
+    def convert_to_sql(self):
+        conditions = [MetaDocumentORM.document_id.in_(condition.convert_to_sql()) for condition in self.conditions]
+        return select(MetaDocumentORM.document_id).filter(or_(*conditions))
 
 
 class EqOperation(ComparisonOperation):
@@ -228,6 +260,9 @@ class EqOperation(ComparisonOperation):
 
     def convert_to_elasticsearch(self):
         return {"term": {self.field_name: self.comparison_value}}
+
+    def convert_to_sql(self):
+        return select([MetaDocumentORM.document_id]).where(MetaDocumentORM.name == self.field_name, MetaDocumentORM.value == self.comparison_value)
 
 
 class InOperation(ComparisonOperation):
@@ -237,6 +272,9 @@ class InOperation(ComparisonOperation):
 
     def convert_to_elasticsearch(self):
         return {"terms": {self.field_name: self.comparison_value}}
+
+    def convert_to_sql(self):
+        return select([MetaDocumentORM.document_id]).where(MetaDocumentORM.name == self.field_name, MetaDocumentORM.value.in_(self.comparison_value))
 
 
 class NeOperation(ComparisonOperation):
@@ -247,6 +285,9 @@ class NeOperation(ComparisonOperation):
     def convert_to_elasticsearch(self):
         return {"bool": {"must_not": {"term": {self.field_name: self.comparison_value}}}}
 
+    def convert_to_sql(self):
+        return select([MetaDocumentORM.document_id]).where(MetaDocumentORM.name == self.field_name, MetaDocumentORM.value != self.comparison_value)
+
 
 class NinOperation(ComparisonOperation):
     """
@@ -255,6 +296,9 @@ class NinOperation(ComparisonOperation):
 
     def convert_to_elasticsearch(self):
         return {"bool": {"must_not": {"terms": {self.field_name: self.comparison_value}}}}
+
+    def convert_to_sql(self):
+        return select([MetaDocumentORM.document_id]).where(MetaDocumentORM.name == self.field_name, MetaDocumentORM.value.notin_(self.comparison_value))
 
 
 class GtOperation(ComparisonOperation):
@@ -264,6 +308,9 @@ class GtOperation(ComparisonOperation):
 
     def convert_to_elasticsearch(self):
         return {"range": {self.field_name: {"gt": self.comparison_value}}}
+
+    def convert_to_sql(self):
+        return select([MetaDocumentORM.document_id]).where(MetaDocumentORM.name == self.field_name, MetaDocumentORM.value > self.comparison_value)
 
 
 class GteOperation(ComparisonOperation):
@@ -273,6 +320,9 @@ class GteOperation(ComparisonOperation):
 
     def convert_to_elasticsearch(self):
         return {"range": {self.field_name: {"gte": self.comparison_value}}}
+
+    def convert_to_sql(self):
+        return select([MetaDocumentORM.document_id]).where(MetaDocumentORM.name == self.field_name, MetaDocumentORM.value >= self.comparison_value)
 
 
 class LtOperation(ComparisonOperation):
@@ -282,6 +332,9 @@ class LtOperation(ComparisonOperation):
 
     def convert_to_elasticsearch(self):
         return {"range": {self.field_name: {"lt": self.comparison_value}}}
+
+    def convert_to_sql(self):
+        return select([MetaDocumentORM.document_id]).where(MetaDocumentORM.name == self.field_name, MetaDocumentORM.value < self.comparison_value)
 
 
 class LteOperation(ComparisonOperation):
@@ -291,3 +344,6 @@ class LteOperation(ComparisonOperation):
 
     def convert_to_elasticsearch(self):
         return {"range": {self.field_name: {"lte": self.comparison_value}}}
+
+    def convert_to_sql(self):
+        return select([MetaDocumentORM.document_id]).where(MetaDocumentORM.name == self.field_name, MetaDocumentORM.value <= self.comparison_value)

haystack/document_stores/sql.py

@@ -105,6 +105,7 @@ class MetaLabelORM(ORMBase):
         {},
     )  # type: ignore
 
+from haystack.document_stores.filter_utils import LogicalFilterClause # Necessary to avoid circular import
 
 class SQLDocumentStore(BaseDocumentStore):
     def __init__(
@@ -294,12 +295,10 @@ def _query(
         ).filter_by(index=index)
 
         if filters:
-            for key, values in filters.items():
-                documents_query = documents_query.join(MetaDocumentORM, aliased=True).filter(
-                    MetaDocumentORM.name == key,
-                    MetaDocumentORM.value.in_(values),
-                )
-
+            parsed_filter = LogicalFilterClause.parse(filters)
+            select_ids = parsed_filter.convert_to_sql()
+            documents_query = documents_query.filter(DocumentORM.id.in_(select_ids))
+
         if only_documents_without_embedding:
             documents_query = documents_query.filter(DocumentORM.vector_id.is_(None))
         if vector_ids:

test/test_document_store.py

@@ -216,7 +216,7 @@ def test_get_all_documents_with_incorrect_filter_value(document_store_with_docs)
     assert len(documents) == 0
 
 
-@pytest.mark.parametrize("document_store_with_docs", ["elasticsearch"], indirect=True)
+@pytest.mark.parametrize("document_store_with_docs", ["elasticsearch", "sql"], indirect=True)
 def test_extended_filter(document_store_with_docs):
     # Test comparison operators individually
     documents = document_store_with_docs.get_all_documents(filters={"meta_field": {"$eq": "test1"}})