Coverage for haystack/components/joiners/document_joiner.py: 100%
104 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import itertools
6from collections import defaultdict
7from dataclasses import replace
8from enum import Enum
9from math import inf
10from typing import Any
12from haystack import Document, component, default_from_dict, default_to_dict, logging
13from haystack.core.component.types import Variadic
14from haystack.utils.misc import _reciprocal_rank_fusion
16logger = logging.getLogger(__name__)
19class JoinMode(Enum):
20 """
21 Enum for join mode.
22 """
24 CONCATENATE = "concatenate"
25 MERGE = "merge"
26 RECIPROCAL_RANK_FUSION = "reciprocal_rank_fusion"
27 DISTRIBUTION_BASED_RANK_FUSION = "distribution_based_rank_fusion"
29 def __str__(self) -> str:
30 return self.value
32 @staticmethod
33 def from_str(string: str) -> "JoinMode":
34 """
35 Convert a string to a JoinMode enum.
36 """
37 enum_map = {e.value: e for e in JoinMode}
38 mode = enum_map.get(string)
39 if mode is None:
40 msg = f"Unknown join mode '{string}'. Supported modes in DocumentJoiner are: {list(enum_map.keys())}"
41 raise ValueError(msg)
42 return mode
45@component
46class DocumentJoiner:
47 """
48 Joins multiple lists of documents into a single list.
50 It supports different join modes:
51 - concatenate: Keeps the highest-scored document in case of duplicates.
52 - merge: Calculates a weighted sum of scores for duplicates and merges them.
53 - reciprocal_rank_fusion: Merges and assigns scores based on reciprocal rank fusion.
54 - distribution_based_rank_fusion: Merges and assigns scores based on scores distribution in each Retriever.
56 ### Usage example:
58 ```python
59 from haystack import Pipeline, Document
60 from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder
61 from haystack.components.joiners import DocumentJoiner
62 from haystack.components.retrievers import InMemoryBM25Retriever
63 from haystack.components.retrievers import InMemoryEmbeddingRetriever
64 from haystack.document_stores.in_memory import InMemoryDocumentStore
66 document_store = InMemoryDocumentStore()
67 docs = [Document(content="Paris"), Document(content="Berlin"), Document(content="London")]
68 embedder = OpenAIDocumentEmbedder()
69 docs_embeddings = embedder.run(docs)
70 document_store.write_documents(docs_embeddings['documents'])
72 p = Pipeline()
73 p.add_component(instance=InMemoryBM25Retriever(document_store=document_store), name="bm25_retriever")
74 p.add_component(
75 instance=OpenAITextEmbedder(),
76 name="text_embedder",
77 )
78 p.add_component(instance=InMemoryEmbeddingRetriever(document_store=document_store), name="embedding_retriever")
79 p.add_component(instance=DocumentJoiner(), name="joiner")
80 p.connect("bm25_retriever", "joiner")
81 p.connect("embedding_retriever", "joiner")
82 p.connect("text_embedder.embedding", "embedding_retriever.query_embedding")
83 query = "What is the capital of France?"
84 p.run(data={"query": query, "text": query, "top_k": 1})
85 ```
86 """
88 def __init__(
89 self,
90 join_mode: str | JoinMode = JoinMode.CONCATENATE,
91 weights: list[float] | None = None,
92 top_k: int | None = None,
93 sort_by_score: bool = True,
94 ) -> None:
95 """
96 Creates a DocumentJoiner component.
98 :param join_mode:
99 Specifies the join mode to use. Available modes:
100 - `concatenate`: Keeps the highest-scored document in case of duplicates.
101 - `merge`: Calculates a weighted sum of scores for duplicates and merges them.
102 - `reciprocal_rank_fusion`: Merges and assigns scores based on reciprocal rank fusion.
103 - `distribution_based_rank_fusion`: Merges and assigns scores based on scores
104 distribution in each Retriever.
105 :param weights:
106 Assign importance to each list of documents to influence how they're joined.
107 This parameter is ignored for
108 `concatenate` or `distribution_based_rank_fusion` join modes.
109 Weight for each list of documents must match the number of inputs.
110 Each weight must be a non-negative number.
111 :param top_k:
112 The maximum number of documents to return. Must be `None` or greater than 0.
113 :param sort_by_score:
114 If `True`, sorts the documents by score in descending order.
115 If a document has no score, it is handled as if its score is -infinity.
117 :raises ValueError:
118 If `top_k` is not `None` and is less than or equal to 0,
119 or if any value in `weights` is negative.
120 """
121 if top_k is not None and top_k <= 0:
122 raise ValueError("top_k must be greater than 0.")
123 if isinstance(join_mode, str):
124 join_mode = JoinMode.from_str(join_mode)
125 join_mode_functions = {
126 JoinMode.CONCATENATE: DocumentJoiner._concatenate,
127 JoinMode.MERGE: self._merge,
128 JoinMode.RECIPROCAL_RANK_FUSION: self._rrf,
129 JoinMode.DISTRIBUTION_BASED_RANK_FUSION: DocumentJoiner._distribution_based_rank_fusion,
130 }
131 self.join_mode_function = join_mode_functions[join_mode]
132 self.join_mode = join_mode
133 if weights:
134 if any(weight < 0 for weight in weights):
135 raise ValueError("The provided `weights` must not be negative.")
136 weight_sum = sum(weights)
137 if weight_sum == 0:
138 raise ValueError("The provided `weights` must not sum to zero.")
139 self.weights: list[float] | None = [float(i) / weight_sum for i in weights]
140 else:
141 self.weights = None
142 self.top_k = top_k
143 self.sort_by_score = sort_by_score
145 @component.output_types(documents=list[Document])
146 def run(self, documents: Variadic[list[Document]], top_k: int | None = None) -> dict[str, Any]:
147 """
148 Joins multiple lists of Documents into a single list depending on the `join_mode` parameter.
150 :param documents:
151 List of list of documents to be merged.
152 :param top_k:
153 The maximum number of documents to return. Overrides the instance's `top_k` if provided.
154 A value of 0 returns no documents. Must not be negative.
156 :returns:
157 A dictionary with the following keys:
158 - `documents`: Merged list of Documents
160 :raises ValueError:
161 If `top_k` is negative.
162 """
163 documents = list(documents)
164 output_documents = self.join_mode_function(documents)
166 if self.sort_by_score:
167 output_documents = sorted(
168 output_documents, key=lambda doc: doc.score if doc.score is not None else -inf, reverse=True
169 )
170 if any(doc.score is None for doc in output_documents):
171 logger.info(
172 "Some of the Documents DocumentJoiner got have score=None. It was configured to sort Documents by "
173 "score, so those with score=None were sorted as if they had a score of -infinity."
174 )
176 if top_k is not None:
177 if top_k < 0:
178 raise ValueError("top_k must not be negative.")
179 output_documents = output_documents[:top_k]
180 elif self.top_k is not None:
181 output_documents = output_documents[: self.top_k]
183 return {"documents": output_documents}
185 @staticmethod
186 def _concatenate(document_lists: list[list[Document]]) -> list[Document]:
187 """
188 Concatenate multiple lists of Documents and return only the Document with the highest score for duplicates.
189 """
190 output = []
191 docs_per_id = defaultdict(list)
192 for doc in itertools.chain.from_iterable(document_lists):
193 docs_per_id[doc.id].append(doc)
194 for docs in docs_per_id.values():
195 doc_with_best_score = max(docs, key=lambda doc: doc.score if doc.score is not None else -inf)
196 output.append(doc_with_best_score)
197 return output
199 def _merge(self, document_lists: list[list[Document]]) -> list[Document]:
200 """
201 Merge multiple lists of Documents and calculate a weighted sum of the scores of duplicate Documents.
202 """
203 # This check prevents a division by zero when no documents are passed
204 if not document_lists:
205 return []
207 scores_map: dict = defaultdict(int)
208 documents_map = {}
209 weights = self.weights if self.weights else [1 / len(document_lists)] * len(document_lists)
211 for documents, weight in zip(document_lists, weights, strict=True):
212 for doc in documents:
213 scores_map[doc.id] += (doc.score if doc.score is not None else 0) * weight
214 documents_map[doc.id] = doc
216 return [replace(doc, score=scores_map[doc.id]) for doc in documents_map.values()]
218 def _rrf(self, document_lists: list[list[Document]]) -> list[Document]:
219 """
220 Merge multiple lists of Documents and assign scores based on reciprocal rank fusion.
221 """
222 return _reciprocal_rank_fusion(document_lists, weights=self.weights)
224 @staticmethod
225 def _distribution_based_rank_fusion(document_lists: list[list[Document]]) -> list[Document]:
226 """
227 Merge multiple lists of Documents and assign scores based on Distribution-Based Score Fusion.
229 (https://medium.com/plain-simple-software/distribution-based-score-fusion-dbsf-a-new-approach-to-vector-search-ranking-f87c37488b18)
230 If a Document is in more than one retriever, the one with the highest score is used.
231 """
232 rescaled_lists: list[list[Document]] = []
233 for documents in document_lists:
234 if len(documents) == 0:
235 rescaled_lists.append(documents)
236 continue
238 scores_list = [doc.score if doc.score is not None else 0 for doc in documents]
240 mean_score = sum(scores_list) / len(scores_list)
241 std_dev = (sum((x - mean_score) ** 2 for x in scores_list) / len(scores_list)) ** 0.5
242 min_score = mean_score - 3 * std_dev
243 max_score = mean_score + 3 * std_dev
244 delta_score = max_score - min_score
245 scores_vary = min(scores_list) != max(scores_list)
247 # If all docs have the same score, rescaling cannot add information; preserve their original scores.
248 rescaled_lists.append(
249 [
250 replace(
251 doc,
252 score=((doc.score if doc.score is not None else 0) - min_score) / delta_score
253 if scores_vary and delta_score != 0.0
254 else (doc.score if doc.score is not None else 0.0),
255 )
256 for doc in documents
257 ]
258 )
260 return DocumentJoiner._concatenate(document_lists=rescaled_lists)
262 def to_dict(self) -> dict[str, Any]:
263 """
264 Serializes the component to a dictionary.
266 :returns:
267 Dictionary with serialized data.
268 """
269 return default_to_dict(
270 self,
271 join_mode=str(self.join_mode),
272 weights=self.weights,
273 top_k=self.top_k,
274 sort_by_score=self.sort_by_score,
275 )
277 @classmethod
278 def from_dict(cls, data: dict[str, Any]) -> "DocumentJoiner":
279 """
280 Deserializes the component from a dictionary.
282 :param data:
283 The dictionary to deserialize from.
284 :returns:
285 The deserialized component.
286 """
287 return default_from_dict(cls, data)