Coverage for haystack/components/rankers/meta_field_grouping_ranker.py: 100%
39 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-09 14:34 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5from collections import defaultdict
7from haystack import Document, component, logging
8from haystack.utils.misc import _deduplicate_documents
10logger = logging.getLogger(__name__)
13@component
14class MetaFieldGroupingRanker:
15 """
16 Reorders the documents by grouping them based on metadata keys.
18 The MetaFieldGroupingRanker can group documents by a primary metadata key `group_by`, and subgroup them with an optional
19 secondary key, `subgroup_by`.
20 Within each group or subgroup, it can also sort documents by a metadata key `sort_docs_by`.
22 The output is a flat list of documents ordered by `group_by` and `subgroup_by` values.
23 Any documents without a group are placed at the end of the list.
25 The proper organization of documents helps improve the efficiency and performance of subsequent processing by an LLM.
27 ### Usage example
29 ```python
30 from haystack.components.rankers import MetaFieldGroupingRanker
31 from haystack.dataclasses import Document
34 docs = [
35 Document(content="Javascript is a popular programming language", meta={"group": "42", "split_id": 7, "subgroup": "subB"}),
36 Document(content="Python is a popular programming language",meta={"group": "42", "split_id": 4, "subgroup": "subB"}),
37 Document(content="A chromosome is a package of DNA", meta={"group": "314", "split_id": 2, "subgroup": "subC"}),
38 Document(content="An octopus has three hearts", meta={"group": "11", "split_id": 2, "subgroup": "subD"}),
39 Document(content="Java is a popular programming language", meta={"group": "42", "split_id": 3, "subgroup": "subB"})
40 ]
42 ranker = MetaFieldGroupingRanker(group_by="group",subgroup_by="subgroup", sort_docs_by="split_id")
43 result = ranker.run(documents=docs)
44 print(result["documents"])
46 # >> [
47 # >> Document(id=d665bbc83e52c08c3d8275bccf4f22bf2bfee21c6e77d78794627637355b8ebc,
48 # >> content: 'Java is a popular programming language', meta: {'group': '42', 'split_id': 3, 'subgroup': 'subB'}),
49 # >> Document(id=a20b326f07382b3cbf2ce156092f7c93e8788df5d48f2986957dce2adb5fe3c2,
50 # >> content: 'Python is a popular programming language', meta: {'group': '42', 'split_id': 4, 'subgroup': 'subB'}),
51 # >> Document(id=ce12919795d22f6ca214d0f161cf870993889dcb146f3bb1b3e1ffdc95be960f,
52 # >> content: 'Javascript is a popular programming language', meta: {'group': '42', 'split_id': 7, 'subgroup': 'subB'}),
53 # >> Document(id=d9fc857046c904e5cf790b3969b971b1bbdb1b3037d50a20728fdbf82991aa94,
54 # >> content: 'A chromosome is a package of DNA', meta: {'group': '314', 'split_id': 2, 'subgroup': 'subC'}),
55 # >> Document(id=6d3b7bdc13d09aa01216471eb5fb0bfdc53c5f2f3e98ad125ff6b85d3106c9a3,
56 # >> content: 'An octopus has three hearts', meta: {'group': '11', 'split_id': 2, 'subgroup': 'subD'})
57 ```
58 """ # noqa: E501
60 def __init__(self, group_by: str, subgroup_by: str | None = None, sort_docs_by: str | None = None) -> None:
61 """
62 Creates an instance of MetaFieldGroupingRanker.
64 :param group_by: The metadata key to aggregate the documents by.
65 :param subgroup_by: The metadata key to aggregate the documents within a group that was created by the
66 `group_by` key.
67 :param sort_docs_by: Determines which metadata key is used to sort the documents. If not provided, the
68 documents within the groups or subgroups are not sorted and are kept in the same order as
69 they were inserted in the subgroups.
71 """
72 self.group_by = group_by
73 self.sort_docs_by = sort_docs_by
74 self.subgroup_by = subgroup_by
76 @component.output_types(documents=list[Document])
77 def run(self, documents: list[Document]) -> dict[str, list[Document]]:
78 """
79 Groups the provided list of documents based on the `group_by` parameter and optionally the `subgroup_by`.
81 Before grouping, documents are deduplicated by their id, retaining only the document with the highest score
82 if a score is present.
84 The output is a list of documents reordered based on how they were grouped.
86 :param documents: The list of documents to group.
87 :returns:
88 A dictionary with the following keys:
89 - documents: The list of documents ordered by the `group_by` and `subgroup_by` metadata values.
90 """
92 if not documents:
93 return {"documents": []}
95 document_groups: dict[str, dict[str, list[Document]]] = defaultdict(lambda: defaultdict(list))
96 no_group_docs = []
98 deduplicated_documents = _deduplicate_documents(documents)
99 for doc in deduplicated_documents:
100 # A value of None counts as missing, as it does for `sort_docs_by`
101 raw_group_value = doc.meta.get(self.group_by)
102 group_value = "" if raw_group_value is None else str(raw_group_value)
104 # If no group value, add to no_group_docs and continue
105 if not group_value:
106 no_group_docs.append(doc)
107 continue
109 # Get subgroup value or use a default if not specified
110 subgroup_value = "no_subgroup"
111 if self.subgroup_by and doc.meta.get(self.subgroup_by) is not None:
112 subgroup_value = str(doc.meta[self.subgroup_by])
114 document_groups[group_value][subgroup_value].append(doc)
116 # use a non-optional key for type checking; "" disables sorting.
117 sort_field = self.sort_docs_by or ""
119 ordered_docs = []
120 for subgroups in document_groups.values():
121 for docs in subgroups.values():
122 if sort_field:
123 # Sort by the field value, placing documents with a missing value last.
124 # The (is_missing, value) tuple keeps documents with a missing value out of the
125 # value comparison, but two present values of mutually non-comparable types
126 # (e.g. an int and a str) would still raise a TypeError. In that case we keep the
127 # group's insertion order instead of crashing, mirroring MetaFieldRanker.
128 try:
129 docs = sorted(docs, key=lambda d: (d.meta.get(sort_field) is None, d.meta.get(sort_field)))
130 except TypeError as error:
131 logger.warning(
132 "Tried to sort Documents with IDs {document_ids}, but got TypeError with the "
133 "message: {error}\nKeeping the original order of the Documents in this group "
134 "since sorting by '{sort_field}' is not possible.",
135 document_ids=",".join([doc.id for doc in docs]),
136 error=error,
137 sort_field=sort_field,
138 )
139 ordered_docs.extend(docs)
141 ordered_docs.extend(no_group_docs)
143 return {"documents": ordered_docs}