Coverage for haystack/components/joiners/document_joiner.py: 100%
101 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +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 :param top_k:
111 The maximum number of documents to return. Must be `None` or greater than 0.
112 :param sort_by_score:
113 If `True`, sorts the documents by score in descending order.
114 If a document has no score, it is handled as if its score is -infinity.
116 :raises ValueError:
117 If `top_k` is not `None` and is less than or equal to 0.
118 """
119 if top_k is not None and top_k <= 0:
120 raise ValueError("top_k must be greater than 0.")
121 if isinstance(join_mode, str):
122 join_mode = JoinMode.from_str(join_mode)
123 join_mode_functions = {
124 JoinMode.CONCATENATE: DocumentJoiner._concatenate,
125 JoinMode.MERGE: self._merge,
126 JoinMode.RECIPROCAL_RANK_FUSION: self._rrf,
127 JoinMode.DISTRIBUTION_BASED_RANK_FUSION: DocumentJoiner._distribution_based_rank_fusion,
128 }
129 self.join_mode_function = join_mode_functions[join_mode]
130 self.join_mode = join_mode
131 if weights:
132 weight_sum = sum(weights)
133 if weight_sum == 0:
134 raise ValueError("The provided `weights` must not sum to zero.")
135 self.weights: list[float] | None = [float(i) / weight_sum for i in weights]
136 else:
137 self.weights = None
138 self.top_k = top_k
139 self.sort_by_score = sort_by_score
141 @component.output_types(documents=list[Document])
142 def run(self, documents: Variadic[list[Document]], top_k: int | None = None) -> dict[str, Any]:
143 """
144 Joins multiple lists of Documents into a single list depending on the `join_mode` parameter.
146 :param documents:
147 List of list of documents to be merged.
148 :param top_k:
149 The maximum number of documents to return. Overrides the instance's `top_k` if provided.
150 A value of 0 returns no documents. Must not be negative.
152 :returns:
153 A dictionary with the following keys:
154 - `documents`: Merged list of Documents
156 :raises ValueError:
157 If `top_k` is negative.
158 """
159 documents = list(documents)
160 output_documents = self.join_mode_function(documents)
162 if self.sort_by_score:
163 output_documents = sorted(
164 output_documents, key=lambda doc: doc.score if doc.score is not None else -inf, reverse=True
165 )
166 if any(doc.score is None for doc in output_documents):
167 logger.info(
168 "Some of the Documents DocumentJoiner got have score=None. It was configured to sort Documents by "
169 "score, so those with score=None were sorted as if they had a score of -infinity."
170 )
172 if top_k is not None:
173 if top_k < 0:
174 raise ValueError("top_k must not be negative.")
175 output_documents = output_documents[:top_k]
176 elif self.top_k is not None:
177 output_documents = output_documents[: self.top_k]
179 return {"documents": output_documents}
181 @staticmethod
182 def _concatenate(document_lists: list[list[Document]]) -> list[Document]:
183 """
184 Concatenate multiple lists of Documents and return only the Document with the highest score for duplicates.
185 """
186 output = []
187 docs_per_id = defaultdict(list)
188 for doc in itertools.chain.from_iterable(document_lists):
189 docs_per_id[doc.id].append(doc)
190 for docs in docs_per_id.values():
191 doc_with_best_score = max(docs, key=lambda doc: doc.score if doc.score is not None else -inf)
192 output.append(doc_with_best_score)
193 return output
195 def _merge(self, document_lists: list[list[Document]]) -> list[Document]:
196 """
197 Merge multiple lists of Documents and calculate a weighted sum of the scores of duplicate Documents.
198 """
199 # This check prevents a division by zero when no documents are passed
200 if not document_lists:
201 return []
203 scores_map: dict = defaultdict(int)
204 documents_map = {}
205 weights = self.weights if self.weights else [1 / len(document_lists)] * len(document_lists)
207 for documents, weight in zip(document_lists, weights, strict=True):
208 for doc in documents:
209 scores_map[doc.id] += (doc.score if doc.score is not None else 0) * weight
210 documents_map[doc.id] = doc
212 return [replace(doc, score=scores_map[doc.id]) for doc in documents_map.values()]
214 def _rrf(self, document_lists: list[list[Document]]) -> list[Document]:
215 """
216 Merge multiple lists of Documents and assign scores based on reciprocal rank fusion.
217 """
218 return _reciprocal_rank_fusion(document_lists, weights=self.weights)
220 @staticmethod
221 def _distribution_based_rank_fusion(document_lists: list[list[Document]]) -> list[Document]:
222 """
223 Merge multiple lists of Documents and assign scores based on Distribution-Based Score Fusion.
225 (https://medium.com/plain-simple-software/distribution-based-score-fusion-dbsf-a-new-approach-to-vector-search-ranking-f87c37488b18)
226 If a Document is in more than one retriever, the one with the highest score is used.
227 """
228 rescaled_lists: list[list[Document]] = []
229 for documents in document_lists:
230 if len(documents) == 0:
231 rescaled_lists.append(documents)
232 continue
234 scores_list = [doc.score if doc.score is not None else 0 for doc in documents]
236 mean_score = sum(scores_list) / len(scores_list)
237 std_dev = (sum((x - mean_score) ** 2 for x in scores_list) / len(scores_list)) ** 0.5
238 min_score = mean_score - 3 * std_dev
239 max_score = mean_score + 3 * std_dev
240 delta_score = max_score - min_score
242 # if all docs have the same score delta_score is 0, the docs are uninformative for the query
243 rescaled_lists.append(
244 [
245 replace(
246 doc,
247 score=((doc.score if doc.score is not None else 0) - min_score) / delta_score
248 if delta_score != 0.0
249 else 0.0,
250 )
251 for doc in documents
252 ]
253 )
255 return DocumentJoiner._concatenate(document_lists=rescaled_lists)
257 def to_dict(self) -> dict[str, Any]:
258 """
259 Serializes the component to a dictionary.
261 :returns:
262 Dictionary with serialized data.
263 """
264 return default_to_dict(
265 self,
266 join_mode=str(self.join_mode),
267 weights=self.weights,
268 top_k=self.top_k,
269 sort_by_score=self.sort_by_score,
270 )
272 @classmethod
273 def from_dict(cls, data: dict[str, Any]) -> "DocumentJoiner":
274 """
275 Deserializes the component from a dictionary.
277 :param data:
278 The dictionary to deserialize from.
279 :returns:
280 The deserialized component.
281 """
282 return default_from_dict(cls, data)