Coverage for haystack/components/preprocessors/document_splitter.py: 100%
206 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
5from collections.abc import Callable
6from copy import deepcopy
7from typing import Any, Literal
9from more_itertools import windowed
11from haystack import Document, component, logging
12from haystack.components.preprocessors.sentence_tokenizer import Language, SentenceSplitter, nltk_imports
13from haystack.core.serialization import default_from_dict, default_to_dict
14from haystack.utils import deserialize_callable, serialize_callable
16logger = logging.getLogger(__name__)
18# mapping of split by character, 'function' and 'sentence' don't split by character
19_CHARACTER_SPLIT_BY_MAPPING = {"page": "\f", "passage": "\n\n", "period": ".", "word": " ", "line": "\n"}
22@component
23class DocumentSplitter:
24 """
25 Splits long documents into smaller chunks.
27 This is a common preprocessing step during indexing. It helps Embedders create meaningful semantic representations
28 and prevents exceeding language model context limits.
30 The DocumentSplitter is compatible with the following DocumentStores:
31 - [Astra](https://docs.haystack.deepset.ai/docs/astradocumentstore)
32 - [Chroma](https://docs.haystack.deepset.ai/docs/chromadocumentstore) limited support, overlapping information is
33 not stored
34 - [Elasticsearch](https://docs.haystack.deepset.ai/docs/elasticsearch-document-store)
35 - [OpenSearch](https://docs.haystack.deepset.ai/docs/opensearch-document-store)
36 - [Pgvector](https://docs.haystack.deepset.ai/docs/pgvectordocumentstore)
37 - [Pinecone](https://docs.haystack.deepset.ai/docs/pinecone-document-store) limited support, overlapping
38 information is not stored
39 - [Qdrant](https://docs.haystack.deepset.ai/docs/qdrant-document-store)
40 - [Weaviate](https://docs.haystack.deepset.ai/docs/weaviatedocumentstore)
42 ### Usage example
44 ```python
45 from haystack import Document
46 from haystack.components.preprocessors import DocumentSplitter
48 doc = Document(content="Moonlight shimmered softly, wolves howled nearby, night enveloped everything.")
50 splitter = DocumentSplitter(split_by="word", split_length=3, split_overlap=0)
51 result = splitter.run(documents=[doc])
52 ```
53 """
55 def __init__(
56 self,
57 split_by: Literal["function", "page", "passage", "period", "word", "line", "sentence"] = "word",
58 split_length: int = 200,
59 split_overlap: int = 0,
60 split_threshold: int = 0,
61 splitting_function: Callable[[str], list[str]] | None = None,
62 respect_sentence_boundary: bool = False,
63 language: Language = "en",
64 use_split_rules: bool = True,
65 extend_abbreviations: bool = True,
66 *,
67 skip_empty_documents: bool = True,
68 ) -> None:
69 """
70 Initialize DocumentSplitter.
72 :param split_by: The unit for splitting your documents. Choose from:
73 - `word` for splitting by spaces (" ")
74 - `period` for splitting by periods (".")
75 - `page` for splitting by form feed ("\\f")
76 - `passage` for splitting by double line breaks ("\\n\\n")
77 - `line` for splitting each line ("\\n")
78 - `sentence` for splitting by NLTK sentence tokenizer
80 :param split_length: The maximum number of units in each split.
81 :param split_overlap: The number of overlapping units for each split.
82 :param split_threshold: The minimum number of units per split. If a split has fewer units
83 than the threshold, it's attached to the previous split.
84 :param splitting_function: Necessary when `split_by` is set to "function".
85 This is a function which must accept a single `str` as input and return a `list` of `str` as output,
86 representing the chunks after splitting.
87 :param respect_sentence_boundary: Choose whether to respect sentence boundaries when splitting by "word".
88 If True, uses NLTK to detect sentence boundaries, ensuring splits occur only between sentences.
89 :param language: Choose the language for the NLTK tokenizer. The default is English ("en").
90 :param use_split_rules: Choose whether to use additional split rules when splitting by `sentence`.
91 :param extend_abbreviations: Choose whether to extend NLTK's PunktTokenizer abbreviations with a list
92 of curated abbreviations, if available. This is currently supported for English ("en") and German ("de").
93 :param skip_empty_documents: Choose whether to skip documents with empty content. Default is True.
94 Set to False when downstream components in the Pipeline (like LLMDocumentContentExtractor) can extract text
95 from non-textual documents.
96 """
98 self.split_by = split_by
99 self.split_length = split_length
100 self.split_overlap = split_overlap
101 self.split_threshold = split_threshold
102 self.splitting_function = splitting_function
103 self.respect_sentence_boundary = respect_sentence_boundary
104 self.language = language
105 self.use_split_rules = use_split_rules
106 self.extend_abbreviations = extend_abbreviations
107 self.skip_empty_documents = skip_empty_documents
109 self._init_checks(
110 split_by=split_by,
111 split_length=split_length,
112 split_overlap=split_overlap,
113 splitting_function=splitting_function,
114 respect_sentence_boundary=respect_sentence_boundary,
115 )
116 self._use_sentence_splitter = split_by == "sentence" or (respect_sentence_boundary and split_by == "word")
117 if self._use_sentence_splitter:
118 nltk_imports.check()
119 self.sentence_splitter: SentenceSplitter | None = None
121 def _init_checks(
122 self,
123 *,
124 split_by: str,
125 split_length: int,
126 split_overlap: int,
127 splitting_function: Callable | None,
128 respect_sentence_boundary: bool,
129 ) -> None:
130 """
131 Validates initialization parameters for DocumentSplitter.
133 :param split_by: The unit for splitting documents
134 :param split_length: The maximum number of units in each split
135 :param split_overlap: The number of overlapping units for each split
136 :param splitting_function: Custom function for splitting when split_by="function"
137 :param respect_sentence_boundary: Whether to respect sentence boundaries when splitting
138 :raises ValueError: If any parameter is invalid
139 """
140 valid_split_by = ["function", "page", "passage", "period", "word", "line", "sentence"]
141 if split_by not in valid_split_by:
142 raise ValueError(f"split_by must be one of {', '.join(valid_split_by)}.")
144 if split_by == "function" and splitting_function is None:
145 raise ValueError("When 'split_by' is set to 'function', a valid 'splitting_function' must be provided.")
147 if split_length <= 0:
148 raise ValueError("split_length must be greater than 0.")
150 if split_overlap < 0:
151 raise ValueError("split_overlap must be greater than or equal to 0.")
153 if split_overlap >= split_length:
154 raise ValueError("split_overlap must be less than split_length.")
156 if respect_sentence_boundary and split_by != "word":
157 logger.warning(
158 "The 'respect_sentence_boundary' option is only supported for `split_by='word'`. "
159 "The option `respect_sentence_boundary` will be set to `False`."
160 )
161 self.respect_sentence_boundary = False
163 def warm_up(self) -> None:
164 """
165 Warm up the DocumentSplitter by loading the sentence tokenizer.
166 """
167 if self._use_sentence_splitter and self.sentence_splitter is None:
168 self.sentence_splitter = SentenceSplitter(
169 language=self.language,
170 use_split_rules=self.use_split_rules,
171 extend_abbreviations=self.extend_abbreviations,
172 keep_white_spaces=True,
173 )
175 @component.output_types(documents=list[Document])
176 def run(self, documents: list[Document]) -> dict[str, list[Document]]:
177 """
178 Split documents into smaller parts.
180 Splits documents by the unit expressed in `split_by`, with a length of `split_length`
181 and an overlap of `split_overlap`.
183 :param documents: The documents to split.
184 :returns: A dictionary with the following key:
185 - `documents`: List of documents with the split texts. Each document includes:
186 - A metadata field `source_id` to track the original document.
187 - A metadata field `page_number` to track the original page number.
188 - All other metadata copied from the original document.
190 :raises TypeError: if the input is not a list of Documents.
191 :raises ValueError: if the content of a document is None.
192 """
193 if self._use_sentence_splitter and self.sentence_splitter is None:
194 self.warm_up()
196 if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
197 raise TypeError("DocumentSplitter expects a List of Documents as input.")
199 split_docs: list[Document] = []
200 for doc in documents:
201 if doc.content is None:
202 raise ValueError(
203 f"DocumentSplitter only works with text documents but content for document ID {doc.id} is None."
204 )
205 if doc.content == "" and self.skip_empty_documents:
206 logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
207 continue
209 split_docs += self._split_document(doc)
210 return {"documents": split_docs}
212 def _split_document(self, doc: Document) -> list[Document]:
213 if self.split_by == "sentence" or self.respect_sentence_boundary:
214 return self._split_by_nltk_sentence(doc)
216 if self.split_by == "function" and self.splitting_function is not None:
217 return self._split_by_function(doc)
219 return self._split_by_character(doc)
221 def _split_by_nltk_sentence(self, doc: Document) -> list[Document]:
222 split_docs = []
224 result = self.sentence_splitter.split_sentences(doc.content) # type: ignore # None check is done in run()
225 units = [sentence["sentence"] for sentence in result]
227 if self.respect_sentence_boundary:
228 text_splits, splits_pages, splits_start_idxs = self._concatenate_sentences_based_on_word_amount(
229 sentences=units, split_length=self.split_length, split_overlap=self.split_overlap
230 )
231 else:
232 text_splits, splits_pages, splits_start_idxs = self._concatenate_units(
233 elements=units,
234 split_length=self.split_length,
235 split_overlap=self.split_overlap,
236 split_threshold=self.split_threshold,
237 )
238 metadata = deepcopy(doc.meta)
239 metadata["source_id"] = doc.id
240 split_docs += self._create_docs_from_splits(
241 text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata
242 )
244 return split_docs
246 def _split_by_character(self, doc: Document) -> list[Document]:
247 split_at = _CHARACTER_SPLIT_BY_MAPPING[self.split_by]
248 units = doc.content.split(split_at) # type: ignore[union-attr]
249 # Add the delimiter back to all units except the last one
250 for i in range(len(units) - 1):
251 units[i] += split_at
252 text_splits, splits_pages, splits_start_idxs = self._concatenate_units(
253 units, self.split_length, self.split_overlap, self.split_threshold
254 )
255 metadata = deepcopy(doc.meta)
256 metadata["source_id"] = doc.id
257 return self._create_docs_from_splits(
258 text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata
259 )
261 def _split_by_function(self, doc: Document) -> list[Document]:
262 # the check for None is done already in the run method
263 splits = self.splitting_function(doc.content) # type: ignore
264 docs: list[Document] = []
265 for s in splits:
266 meta = deepcopy(doc.meta)
267 meta["source_id"] = doc.id
268 docs.append(Document(content=s, meta=meta))
269 return docs
271 def _concatenate_units(
272 self, elements: list[str], split_length: int, split_overlap: int, split_threshold: int
273 ) -> tuple[list[str], list[int], list[int]]:
274 """
275 Concatenates the elements into parts of split_length units.
277 Keeps track of the original page number that each element belongs. If the length of the current units is less
278 than the pre-defined `split_threshold`, it does not create a new split. Instead, it concatenates the current
279 units with the last split, preventing the creation of excessively small splits.
280 """
282 text_splits: list[str] = []
283 splits_pages: list[int] = []
284 splits_start_idxs: list[int] = []
285 cur_start_idx = 0
286 cur_page = 1
287 segments = windowed(elements, n=split_length, step=split_length - split_overlap)
289 for seg in segments:
290 current_units = [unit for unit in seg if unit is not None]
291 txt = "".join(current_units)
293 # check if length of current units is below split_threshold
294 if len(current_units) < split_threshold and len(text_splits) > 0:
295 # concatenate the last split with the current one. The first `split_overlap` units
296 # of this segment are already part of the previous split (that is what the overlap
297 # is), so only append the units beyond the overlap; otherwise the overlapping text
298 # would be duplicated and the merged chunk would not be present in the source.
299 text_splits[-1] += "".join(current_units[split_overlap:])
301 # NOTE: If skip_empty_documents is True, this line skips documents that have content=""
302 elif not self.skip_empty_documents or len(txt) > 0:
303 text_splits.append(txt)
304 splits_pages.append(cur_page)
305 splits_start_idxs.append(cur_start_idx)
307 processed_units = current_units[: split_length - split_overlap]
308 cur_start_idx += len("".join(processed_units))
310 if self.split_by == "page":
311 num_page_breaks = len(processed_units)
312 else:
313 num_page_breaks = sum(processed_unit.count("\f") for processed_unit in processed_units)
315 cur_page += num_page_breaks
317 return text_splits, splits_pages, splits_start_idxs
319 def _create_docs_from_splits(
320 self, text_splits: list[str], splits_pages: list[int], splits_start_idxs: list[int], meta: dict[str, Any]
321 ) -> list[Document]:
322 """
323 Creates Document objects from splits enriching them with page number and the metadata of the original document.
324 """
325 documents: list[Document] = []
327 for i, (txt, split_idx) in enumerate(zip(text_splits, splits_start_idxs, strict=True)):
328 copied_meta = deepcopy(meta)
329 copied_meta["page_number"] = splits_pages[i]
330 copied_meta["split_id"] = i
331 copied_meta["split_idx_start"] = split_idx
332 doc = Document(content=txt, meta=copied_meta)
333 documents.append(doc)
335 if self.split_overlap <= 0:
336 continue
338 doc.meta["_split_overlap"] = []
340 if i == 0:
341 continue
343 doc_start_idx = splits_start_idxs[i]
344 previous_doc = documents[i - 1]
345 previous_doc_start_idx = splits_start_idxs[i - 1]
346 self._add_split_overlap_information(doc, doc_start_idx, previous_doc, previous_doc_start_idx)
348 return documents
350 @staticmethod
351 def _add_split_overlap_information(
352 current_doc: Document, current_doc_start_idx: int, previous_doc: Document, previous_doc_start_idx: int
353 ) -> None:
354 """
355 Adds split overlap information to the current and previous Document's meta.
357 :param current_doc: The Document that is being split.
358 :param current_doc_start_idx: The starting index of the current Document.
359 :param previous_doc: The Document that was split before the current Document.
360 :param previous_doc_start_idx: The starting index of the previous Document.
361 """
362 overlapping_range = (current_doc_start_idx - previous_doc_start_idx, len(previous_doc.content)) # type: ignore
364 if overlapping_range[0] < overlapping_range[1]:
365 overlapping_str = previous_doc.content[overlapping_range[0] : overlapping_range[1]] # type: ignore
367 if current_doc.content.startswith(overlapping_str): # type: ignore
368 # add split overlap information to this Document regarding the previous Document
369 current_doc.meta["_split_overlap"].append({"doc_id": previous_doc.id, "range": overlapping_range})
371 # add split overlap information to previous Document regarding this Document
372 overlapping_range = (0, overlapping_range[1] - overlapping_range[0])
373 previous_doc.meta["_split_overlap"].append({"doc_id": current_doc.id, "range": overlapping_range})
375 def to_dict(self) -> dict[str, Any]:
376 """
377 Serializes the component to a dictionary.
378 """
379 serialized = default_to_dict(
380 self,
381 split_by=self.split_by,
382 split_length=self.split_length,
383 split_overlap=self.split_overlap,
384 split_threshold=self.split_threshold,
385 respect_sentence_boundary=self.respect_sentence_boundary,
386 language=self.language,
387 use_split_rules=self.use_split_rules,
388 extend_abbreviations=self.extend_abbreviations,
389 skip_empty_documents=self.skip_empty_documents,
390 )
391 if self.splitting_function:
392 serialized["init_parameters"]["splitting_function"] = serialize_callable(self.splitting_function)
393 return serialized
395 @classmethod
396 def from_dict(cls, data: dict[str, Any]) -> "DocumentSplitter":
397 """
398 Deserializes the component from a dictionary.
399 """
400 init_params = data.get("init_parameters", {})
402 splitting_function = init_params.get("splitting_function", None)
403 if splitting_function:
404 init_params["splitting_function"] = deserialize_callable(splitting_function)
406 return default_from_dict(cls, data)
408 @staticmethod
409 def _concatenate_sentences_based_on_word_amount(
410 sentences: list[str], split_length: int, split_overlap: int
411 ) -> tuple[list[str], list[int], list[int]]:
412 """
413 Groups the sentences into chunks of `split_length` words while respecting sentence boundaries.
415 This function is only used when splitting by `word` and `respect_sentence_boundary` is set to `True`, i.e.:
416 with NLTK sentence tokenizer.
418 :param sentences: The list of sentences to split.
419 :param split_length: The maximum number of words in each split.
420 :param split_overlap: The number of overlapping words in each split.
421 :returns: A tuple containing the concatenated sentences, the start page numbers, and the start indices.
422 """
423 # chunk information
424 chunk_word_count = 0
425 chunk_starting_page_number = 1
426 chunk_start_idx = 0
427 current_chunk: list[str] = []
428 # output lists
429 split_start_page_numbers = []
430 list_of_splits: list[list[str]] = []
431 split_start_indices = []
433 for sentence_idx, sentence in enumerate(sentences):
434 current_chunk.append(sentence)
435 chunk_word_count += len(sentence.split())
436 next_sentence_word_count = (
437 len(sentences[sentence_idx + 1].split()) if sentence_idx < len(sentences) - 1 else 0
438 )
440 # Number of words in the current chunk plus the next sentence is larger than the split_length,
441 # or we reached the last sentence
442 if (chunk_word_count + next_sentence_word_count) > split_length or sentence_idx == len(sentences) - 1:
443 # Save current chunk and start a new one
444 list_of_splits.append(current_chunk)
445 split_start_page_numbers.append(chunk_starting_page_number)
446 split_start_indices.append(chunk_start_idx)
448 # Get the number of sentences that overlap with the next chunk
449 num_sentences_to_keep = DocumentSplitter._number_of_sentences_to_keep(
450 sentences=current_chunk, split_length=split_length, split_overlap=split_overlap
451 )
452 # Set up information for the new chunk
453 if num_sentences_to_keep > 0:
454 # Processed sentences are the ones that are not overlapping with the next chunk
455 processed_sentences = current_chunk[:-num_sentences_to_keep]
456 chunk_starting_page_number += sum(sent.count("\f") for sent in processed_sentences)
457 chunk_start_idx += len("".join(processed_sentences))
458 # Next chunk starts with the sentences that were overlapping with the previous chunk
459 current_chunk = current_chunk[-num_sentences_to_keep:]
460 chunk_word_count = sum(len(s.split()) for s in current_chunk)
461 else:
462 # Here processed_sentences is the same as current_chunk since there is no overlap
463 chunk_starting_page_number += sum(sent.count("\f") for sent in current_chunk)
464 chunk_start_idx += len("".join(current_chunk))
465 current_chunk = []
466 chunk_word_count = 0
468 # Concatenate the sentences together within each split
469 text_splits = []
470 for split in list_of_splits:
471 text = "".join(split)
472 if len(text) > 0:
473 text_splits.append(text)
475 return text_splits, split_start_page_numbers, split_start_indices
477 @staticmethod
478 def _number_of_sentences_to_keep(sentences: list[str], split_length: int, split_overlap: int) -> int:
479 """
480 Returns the number of sentences to keep in the next chunk based on the `split_overlap` and `split_length`.
482 :param sentences: The list of sentences to split.
483 :param split_length: The maximum number of words in each split.
484 :param split_overlap: The number of overlapping words in each split.
485 :returns: The number of sentences to keep in the next chunk.
486 """
487 # If the split_overlap is 0, we don't need to keep any sentences
488 if split_overlap == 0:
489 return 0
491 num_sentences_to_keep = 0
492 num_words = 0
493 # Next overlapping Document should not start exactly the same as the previous one, so we skip the first sentence
494 for sent in reversed(sentences[1:]):
495 num_words += len(sent.split())
496 # If the number of words is larger than the split_length then don't add any more sentences
497 if num_words > split_length:
498 break
499 num_sentences_to_keep += 1
500 if num_words > split_overlap:
501 break
502 return num_sentences_to_keep