process solr grouped results

Author: ciur

poetry.lock

@@ -10,6 +10,7 @@ python = ">=3.11, <4.0"
 pydantic = "^2.7"
 requests = "^2.31.0"
 xapianpy = {version = "1.4.22.post2406040406", optional = true}
+glom = "^23.5.0"
 
 [tool.poetry.extras]
 xapian = ["xapianpy"]

pyproject.toml

@@ -16,17 +16,30 @@ def __init__(self, url: URL):
 
 class ClientRW(Base):
 
-    def search(self, sq: SearchQuery):
+    def search(self, sq: SearchQuery, user_id: str | None = None):
         payload = {
-            'q': sq.query.original_query
+            'q': sq.query.original_query,
+            'group': 'true',
+            'group.field': 'document_id',
+            'rows': sq.rows,
+            'start': sq.start,
+            'group.limit': sq.group_limit,
+            'group.offset': sq.group_offset,
+            'group.sort': 'page_number asc'
         }
 
+        if user_id:
+            payload['q'] = f"{payload['q']} AND user_id:{user_id}"
+
         response = requests.get(
             self.http_select_url,
             params=payload
         )
+        logger.debug(payload)
+
+        result = response.json()
 
-        return response.json()
+        return result
 
     def add(self, some_dict):
         # change data specific for add

salinic/backends/solr/client.py

@@ -1,11 +1,13 @@
 import json
 import logging
 
+from glom import glom
 from pydantic import BaseModel
 
 from salinic.field import Field
 from salinic.query import SearchQuery
-from salinic.utils import filter_keys, first, trim_suffixes
+from salinic.schema import Document, Folder, Page
+from salinic.utils import first
 
 logger = logging.getLogger(__name__)
 
@@ -15,34 +17,131 @@ def __init__(self, client, schema):
         self.client = client
         self.schema = schema
 
-    def search(self, sq: SearchQuery):
+    def search(
+        self,
+        sq: SearchQuery,
+        user_id: str | None = None
+    ) -> list[Document | Folder]:
+        """Query index
+
+        Solr results are grouped by `document_id` field: this way
+        all folder entries will be part of group with `document_id=null`,
+        while all page entities will be grouped per document i.e.
+        pages which belong together are all in the same group.
+
+        {
+          "responseHeader":{
+            ...
+          "grouped":{
+            "document_id":{
+              "matches":26,
+              "groups":[
+                      "groupValue":null,
+                      "doclist":{"numFound":4,"start":0,"numFoundExact":true,"docs":[
+                          {
+                            "id":"0b663599-32b1-4396-8dbe-ae7cd327cec6",
+                            "lang":"en",
+                            "user_id":"4cee7c39-7c34-4cc5-8543-42a8c88c9fe6",
+                            "entity_type":"folder",
+                            "title_txt_en":"A2  updated",
+                            "_version_":1801539995817738240},
+                          {
+                            "id":"768c6841-d37a-4d02-857f-ab7eaf69b27e",
+                            "lang":"en",
+                            "user_id":"4cee7c39-7c34-4cc5-8543-42a8c88c9fe6",
+                            "entity_type":"folder",
+                            "title_txt_en":".inbox",
+                            "_version_":1801539995692957696}]
+                      }},
+                {
+                  "groupValue":"9bc57688-302e-4e1f-840a-c747dcccb362",
+                  "doclist":{"numFound":5,"start":0,"numFoundExact":true,"docs":[
+                      {
+                        "id":"a6e4916f-dea6-414b-aa38-f5b9ea375725",
+                        "document_id":"9bc57688-302e-4e1f-840a-c747dcccb362",
+                        "lang":"en",
+                        "user_id":"4cee7c39-7c34-4cc5-8543-42a8c88c9fe6",
+                        "page_number":1,
+                        "entity_type":"page",
+                        "title_txt_en":"brother_004603.pdf",
+                        "_version_":1801539996374532096},
+                      {
+                        "id":"72f6ca9e-af4b-4235-a56c-a62508e24efe",
+                        "document_id":"9bc57688-302e-4e1f-840a-c747dcccb362",
+                        "lang":"en",
+                        "user_id":"4cee7c39-7c34-4cc5-8543-42a8c88c9fe6",
+                        "page_number":2,
+                        "entity_type":"page",
+                        "title_txt_en":"brother_004603.pdf",
+                        "_version_":1801539996403892224},]
+                  }},
+                {
+                  "groupValue":"200b0201-cfcd-43df-b41f-f1732568a0d2",
+                  "doclist":{"numFound":2,"start":0,"numFoundExact":true,"docs":[
+                      {
+                        "id":"9fa936e6-fe94-46bf-ad01-d8591cc290d4",
+                        "document_id":"200b0201-cfcd-43df-b41f-f1732568a0d2",
+                        "lang":"en",
+                        "user_id":"4cee7c39-7c34-4cc5-8543-42a8c88c9fe6",
+                        "page_number":1,
+                        "entity_type":"page",
+                        "title_txt_en":"brother_004598.pdf",
+                        "_version_":1801539995874361344},
+                      {
+                        "id":"c364994c-eab5-4c6a-842a-6f40537f7a2e",
+                        "document_id":"200b0201-cfcd-43df-b41f-f1732568a0d2",
+                        "lang":"en",
+                        "user_id":"4cee7c39-7c34-4cc5-8543-42a8c88c9fe6",
+                        "page_number":2,
+                        "entity_type":"page",
+                        "title_txt_en":"brother_004598.pdf",
+                        "_version_":1801539995910012928}]
+                  }},
+                  }}]}}}
         """
-
-        """
-        result = self.client.search(sq)
-        if result['response']['numFound'] == 0:
+        result = self.client.search(sq, user_id)
+        grouped = glom(result, 'grouped.document_id')
+        if glom(grouped, 'matches') == 0:
             return []
 
-        docs_list = [
-            trim_suffixes(doc) for doc in result['response']['docs']
-        ]
-        docs = [
-            filter_keys(some_doc, ['_version_'])
-            for some_doc in docs_list
-        ]
-
         result = []
-        for doc in docs:
-            attrs = {}
-            for field_name, value in doc.items():
-                if '_orig_' in field_name:
-                    continue
-                if self.schema.needs_transform(self.schema, field_name):
-                    attrs[field_name] = json.loads(doc[f'{field_name}_orig_'])
-                else:
-                    attrs[field_name] = value
-
-            result.append(sq.entity(**attrs))
+        for group in glom(grouped, 'groups'):
+            if glom(group, 'groupValue'):
+                # groupValue != null => document
+                document_id = glom(group, 'groupValue')
+                title = ''
+                lang = 'en'
+                tags = []
+                pages = []
+                for page in glom(group, 'doclist.docs'):
+                    lang = page.get('lang', 'en')
+                    title = page.get(f'title_txt_{lang}', None)
+                    text = page.get(f'text_txt_{lang}', None)
+                    tags = page.get('tags', [])
+                    p = Page(
+                        id=page['id'],
+                        page_number=page['page_number'],
+                        text=text
+                    )
+                    pages.append(p)
+                item = Document(
+                    id=document_id,
+                    title=title,
+                    lang=lang,
+                    pages=pages,
+                    tags=tags,
+                )
+                result.append(item)
+            else:
+                for folder in glom(group, 'doclist.docs'):
+                    lang = folder.get('lang', 'en')
+                    title = folder.get(f'title_txt_{lang}', None)
+                    item = Folder(
+                        id=folder['id'],
+                        title=title,
+                        tags=folder.get('tags', []),
+                    )
+                    result.append(item)
 
         return result
 

salinic/backends/solr/index.py

@@ -135,7 +135,7 @@ def remove(self, **kwargs):
 
             self.client.delete_document(id_term)
 
-    def search(self, sq: SearchQuery):
+    def search(self, sq: SearchQuery, user_id: str | None = None):
         results = []
 
         if str(sq.query.free_text):

salinic/backends/xapian/index.py

@@ -8,8 +8,8 @@ def __init__(self, engine, schema):
         self.engine = engine
         self.schema = schema
 
-    def search(self, sq: SearchQuery):
-        return self.backend.search(sq)
+    def search(self, sq: SearchQuery, user_id: str | None = None):
+        return self.backend.search(sq, user_id)
 
 
 class IndexRO(IndexBase):

salinic/index.py

@@ -169,10 +169,26 @@ def __repr__(self):
 
 class SearchQuery:
     query: Query
-
-    def __init__(self, entity, query: str):
+    rows: int
+    start: int
+    group_limit: int
+    group_offset: int
+
+    def __init__(
+        self,
+        entity,
+        q: str,
+        rows: int = 100,
+        start: int = 0,
+        group_limit: int = 100,
+        group_offset: int = 0
+    ):
         self.entity = entity
-        self.query = Query(query)
+        self.query = Query(q)
+        self.rows = rows
+        self.start = start
+        self.group_limit = group_limit
+        self.group_offset = group_offset
 
     def __str__(self):
         return f"SearchQuery(query={self.query}, entity={self.entity})"

salinic/query.py

@@ -1,3 +1,5 @@
+from uuid import UUID
+
 from pydantic import BaseModel, ConfigDict, model_serializer
 
 from .field import Field, IdField
@@ -81,3 +83,31 @@ class Index(Schema):
 
     def needs_transform(self, field_name):
         return hasattr(self, f'get_idx_value__{field_name}')
+
+
+class Page(BaseModel):
+    id: UUID
+    page_number: int
+    text: str | None = None
+
+
+class Document(BaseModel):
+    id: UUID
+    title: str
+    lang: str
+    tags: list[str] = []
+    pages: list[Page]
+    entity_type: str = 'document'
+
+    def __hash__(self):
+        return hash(self.model_dump_json())
+
+
+class Folder(BaseModel):
+    id: UUID
+    title: str
+    tags: list[str] = []
+    entity_type: str = 'folder'
+
+    def __hash__(self):
+        return hash(self.model_dump_json())