Repository: deepset-ai/haystack

Issue:
feat: Add `page_number` to meta of Documents in `DocumentSplitter`
**Is your feature request related to a problem? Please describe.**
In Haystack v1 we had an option in the Preprocessor to add the original `page_number` to a Document's meta data when it was split into a chunk. This feature made down stream applications of visualizing the retrieved text from original files (e.g. PDFs) very easy and straightforward, so I'd like to see it in Haystack v2 as well. 

**Describe the solution you'd like**
I would like to add the option to store the `page_number` in the meta info of the Document in the DocumentSplitter component. I believe we can use a similar/same implementation of calculating this like we did for the Preprocessor. 



Discussion:
This issue https://github.com/deepset-ai/haystack/issues/6706 is related since we currently do not keep page break information when converting a PDF file to a Haystack Document. 
Hi @sjrl :)

This is my first issue. I'm trying to understand the requisites better. It seems to me that to keep the page number and the associated text, I suppose we have to keep the chunks in the metadata, e.g.:

```
units = self._split_into_units(doc.content, self.split_by)
text_splits = self._concatenate_units(units, self.split_length, self.split_overlap)
metadata = deepcopy(doc.meta)
metadata["source_id"] = doc.id
metadata["page_number"] = units
split_docs += [Document(content=txt, meta=metadata) for txt in text_splits]
```

This has a few drawbacks:
- duplicated text, `doc.content` and `doc.metadata['page_number']` now have the same information, a possible solution would be to have `self._concatenate_units()` being triggered only when `doc.content` is called/needed
- the `metadata["page_number"]` has the page number 0 - but this can be easily fixed

Hi @davidsbatista! 

Thanks for taking on this issue :) 

I don't think we need to keep the associated text for the use case I am imagining. Basically what we are interested in Haystack terms would like this
1. Load a PDF File
2. Convert a PDF file to a single Document object --> PyPDFToDocument
3. Split the single Document into Chunked Documents (so Document to List of Documents) --> Document Splitter
  - In this final step I would like to insert a `page_number` into the each Doc's metadata in the List of Documents that would tell me which page the chunked doc came from  based on the original single Document. This tracking of `page_number` was done in Haystack v1 by counting and keeping track of page breaks (`\f`)

Does this make more sense? 
Yes, that helps! So essentially, the `DocumentSplitter` should return a `List[Document]`.

But since it has as input:

`def run(self, documents: List[Document]):`

It should return a `List[List[Documents]]`

do you agree?






Hmm I'm not entirely sure. Initially I would say that it makes sense to return `List[List[Documents]]`, but often we want a flattened list to be returned since we will often directly write these documents to a document store which I believe expects `List[Document]` as input. 

So I think to keep that workflow working we should return `List[Document]` or have some way of flattening the list. What do you think?
Hi Sebastian, I can pick up this again after finishing some high-priority issues I need to handle - maybe by the end of the week. Just to let you know, I haven't forgot it
No problem! Thanks for the update
Would be interested by a follow up about this 👀 If something I could do ?
@lambda-science there hasn't been any follow up, feel free to start working on it if you feel like it
This is my first issue, I will try add this feature, but don't know if I will manage to.  :)
@CarlosFerLo let me know if you need help with this

Resolving patch:
diff --git a/haystack/components/preprocessors/document_splitter.py b/haystack/components/preprocessors/document_splitter.py
--- a/haystack/components/preprocessors/document_splitter.py
+++ b/haystack/components/preprocessors/document_splitter.py
@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import List, Literal
+from typing import Dict, List, Literal, Tuple
 
 from more_itertools import windowed
 
@@ -53,7 +53,7 @@ def run(self, documents: List[Document]):
 
         :returns: A dictionary with the following key:
             - `documents`: List of documents with the split texts. A metadata field "source_id" is added to each
-            document to keep track of the original document that was split. Other metadata are copied from the original
+            document to keep track of the original document that was split. Another metadata field "page_number" is added to each number to keep track of the page it belonged to in the original document. Other metadata are copied from the original
             document.
 
         :raises TypeError: if the input is not a list of Documents.
@@ -70,10 +70,12 @@ def run(self, documents: List[Document]):
                     f"DocumentSplitter only works with text documents but document.content for document ID {doc.id} is None."
                 )
             units = self._split_into_units(doc.content, self.split_by)
-            text_splits = self._concatenate_units(units, self.split_length, self.split_overlap)
+            text_splits, splits_pages = self._concatenate_units(units, self.split_length, self.split_overlap)
             metadata = deepcopy(doc.meta)
             metadata["source_id"] = doc.id
-            split_docs += [Document(content=txt, meta=metadata) for txt in text_splits]
+            split_docs += self._create_docs_from_splits(
+                text_splits=text_splits, splits_pages=splits_pages, meta=metadata
+            )
         return {"documents": split_docs}
 
     def _split_into_units(self, text: str, split_by: Literal["word", "sentence", "passage", "page"]) -> List[str]:
@@ -95,15 +97,40 @@ def _split_into_units(self, text: str, split_by: Literal["word", "sentence", "pa
             units[i] += split_at
         return units
 
-    def _concatenate_units(self, elements: List[str], split_length: int, split_overlap: int) -> List[str]:
+    def _concatenate_units(
+        self, elements: List[str], split_length: int, split_overlap: int
+    ) -> Tuple[List[str], List[int]]:
         """
-        Concatenates the elements into parts of split_length units.
+        Concatenates the elements into parts of split_length units keeping track of the original page number that each element belongs.
         """
         text_splits = []
+        splits_pages = []
+        cur_page = 1
         segments = windowed(elements, n=split_length, step=split_length - split_overlap)
         for seg in segments:
             current_units = [unit for unit in seg if unit is not None]
             txt = "".join(current_units)
             if len(txt) > 0:
                 text_splits.append(txt)
-        return text_splits
+                splits_pages.append(cur_page)
+                processed_units = current_units[: split_length - split_overlap]
+                if self.split_by == "page":
+                    num_page_breaks = len(processed_units)
+                else:
+                    num_page_breaks = sum(processed_unit.count("\f") for processed_unit in processed_units)
+                cur_page += num_page_breaks
+        return text_splits, splits_pages
+
+    @staticmethod
+    def _create_docs_from_splits(text_splits: List[str], splits_pages: List[int], meta: Dict) -> List[Document]:
+        """
+        Creates Document objects from text splits enriching them with page number and the metadata of the original document.
+        """
+        documents: List[Document] = []
+
+        for i, txt in enumerate(text_splits):
+            meta = deepcopy(meta)
+            doc = Document(content=txt, meta=meta)
+            doc.meta["page_number"] = splits_pages[i]
+            documents.append(doc)
+        return documents


Test patch:
diff --git a/test/components/preprocessors/test_document_splitter.py b/test/components/preprocessors/test_document_splitter.py
--- a/test/components/preprocessors/test_document_splitter.py
+++ b/test/components/preprocessors/test_document_splitter.py
@@ -141,3 +141,98 @@ def test_copy_metadata(self):
         for doc, split_doc in zip(documents, result["documents"]):
             assert doc.meta.items() <= split_doc.meta.items()
             assert split_doc.content == "Text."
+
+    def test_add_page_number_to_metadata_with_no_overlap_word_split(self):
+        splitter = DocumentSplitter(split_by="word", split_length=2)
+        doc1 = Document(content="This is some text.\f This text is on another page.")
+        doc2 = Document(content="This content has two.\f\f page brakes.")
+        result = splitter.run(documents=[doc1, doc2])
+
+        expected_pages = [1, 1, 2, 2, 2, 1, 1, 3]
+        for doc, p in zip(result["documents"], expected_pages):
+            assert doc.meta["page_number"] == p
+
+    def test_add_page_number_to_metadata_with_no_overlap_sentence_split(self):
+        splitter = DocumentSplitter(split_by="sentence", split_length=1)
+        doc1 = Document(content="This is some text.\f This text is on another page.")
+        doc2 = Document(content="This content has two.\f\f page brakes.")
+        result = splitter.run(documents=[doc1, doc2])
+
+        expected_pages = [1, 1, 1, 1]
+        for doc, p in zip(result["documents"], expected_pages):
+            assert doc.meta["page_number"] == p
+
+    def test_add_page_number_to_metadata_with_no_overlap_passage_split(self):
+        splitter = DocumentSplitter(split_by="passage", split_length=1)
+        doc1 = Document(
+            content="This is a text with some words.\f There is a second sentence.\n\nAnd there is a third sentence.\n\nAnd more passages.\n\n\f And another passage."
+        )
+        result = splitter.run(documents=[doc1])
+
+        expected_pages = [1, 2, 2, 2]
+        for doc, p in zip(result["documents"], expected_pages):
+            assert doc.meta["page_number"] == p
+
+    def test_add_page_number_to_metadata_with_no_overlap_page_split(self):
+        splitter = DocumentSplitter(split_by="page", split_length=1)
+        doc1 = Document(
+            content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
+        )
+        result = splitter.run(documents=[doc1])
+        expected_pages = [1, 2, 3]
+        for doc, p in zip(result["documents"], expected_pages):
+            assert doc.meta["page_number"] == p
+
+        splitter = DocumentSplitter(split_by="page", split_length=2)
+        doc1 = Document(
+            content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
+        )
+        result = splitter.run(documents=[doc1])
+        expected_pages = [1, 3]
+
+        for doc, p in zip(result["documents"], expected_pages):
+            assert doc.meta["page_number"] == p
+
+    def test_add_page_number_to_metadata_with_overlap_word_split(self):
+        splitter = DocumentSplitter(split_by="word", split_length=3, split_overlap=1)
+        doc1 = Document(content="This is some text. And\f this text is on another page.")
+        doc2 = Document(content="This content has two.\f\f page brakes.")
+        result = splitter.run(documents=[doc1, doc2])
+
+        expected_pages = [1, 1, 1, 2, 2, 1, 1, 3]
+        for doc, p in zip(result["documents"], expected_pages):
+            print(doc.content, doc.meta, p)
+            assert doc.meta["page_number"] == p
+
+    def test_add_page_number_to_metadata_with_overlap_sentence_split(self):
+        splitter = DocumentSplitter(split_by="sentence", split_length=2, split_overlap=1)
+        doc1 = Document(content="This is some text. And this is more text.\f This text is on another page. End.")
+        doc2 = Document(content="This content has two.\f\f page brakes. More text.")
+        result = splitter.run(documents=[doc1, doc2])
+
+        expected_pages = [1, 1, 1, 2, 1, 1]
+        for doc, p in zip(result["documents"], expected_pages):
+            print(doc.content, doc.meta, p)
+            assert doc.meta["page_number"] == p
+
+    def test_add_page_number_to_metadata_with_overlap_passage_split(self):
+        splitter = DocumentSplitter(split_by="passage", split_length=2, split_overlap=1)
+        doc1 = Document(
+            content="This is a text with some words.\f There is a second sentence.\n\nAnd there is a third sentence.\n\nAnd more passages.\n\n\f And another passage."
+        )
+        result = splitter.run(documents=[doc1])
+
+        expected_pages = [1, 2, 2]
+        for doc, p in zip(result["documents"], expected_pages):
+            assert doc.meta["page_number"] == p
+
+    def test_add_page_number_to_metadata_with_overlap_page_split(self):
+        splitter = DocumentSplitter(split_by="page", split_length=2, split_overlap=1)
+        doc1 = Document(
+            content="This is a text with some words. There is a second sentence.\f And there is a third sentence.\f And another passage."
+        )
+        result = splitter.run(documents=[doc1])
+        expected_pages = [1, 2, 3]
+
+        for doc, p in zip(result["documents"], expected_pages):
+            assert doc.meta["page_number"] == p