group results by document_id (#19)

Author: ciur

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "salinic"
-version = "0.4.0"
+version = "0.5.0"
 description = "Search abstraction layer"
 authors = ["Eugen Ciur <[email protected]>"]
 readme = "README.md"
@@ -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,31 +17,131 @@ def __init__(self, client, schema):
         self.client = client
         self.schema = schema
 
-    def search(self, sq: SearchQuery):
-        result = self.client.search(sq)
-        if result['response']['numFound'] == 0:
+    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, 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

@@ -1,7 +1,7 @@
 import logging
 from typing import List, Tuple
 
-from salinic.field import Field, NumericField
+from salinic.field import Field, NumericField, StringField
 from salinic.utils import first
 
 from .types import CopyFieldDump, FieldDump, FieldType
@@ -109,6 +109,10 @@ def _normal_fields(self) -> List[Tuple[str, FieldDump]]:
 
             if isinstance(field_instance, NumericField):
                 _type = FieldType.pint
+            elif isinstance(field_instance, StringField):
+                _type = FieldType.string
+            elif field_instance.group:
+                _type = FieldType.text_gen_sort
             else:
                 _type = FieldType.text_general
 

salinic/backends/solr/schema_manager.py

@@ -11,6 +11,8 @@ def to_dash(string: str) -> str:
 
 class FieldType(str, Enum):
     text_general = 'text_general'
+    text_gen_sort = 'text_gen_sort'
+    string = 'string'
     pint = 'pint'
 
 

salinic/backends/solr/types.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

@@ -9,6 +9,8 @@ class Field(NamedTuple):
     default: any = None
     multi_value: bool = False
     multi_lang: bool = False
+    # enables grouping by this field
+    group: bool = False
 
 
 class KeywordField(Field):
@@ -23,6 +25,10 @@ class NumericField(Field):
     pass
 
 
+class StringField(Field):
+    pass
+
+
 class IdField(Field):
     pass
 

salinic/field.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):