fix: add standard split metadata when DocumentSplitter uses split_by="function"

haystack/components/preprocessors/document_splitter.py

@@ -261,12 +261,47 @@ def _split_by_character(self, doc: Document) -> list[Document]:
     def _split_by_function(self, doc: Document) -> list[Document]:
         # the check for None is done already in the run method
         splits = self.splitting_function(doc.content)  # type: ignore
-        docs: list[Document] = []
-        for s in splits:
-            meta = deepcopy(doc.meta)
-            meta["source_id"] = doc.id
-            docs.append(Document(content=s, meta=meta))
-        return docs
+        content: str = doc.content  # type: ignore[assignment]
+
+        text_splits: list[str] = []
+        splits_pages: list[int] = []
+        splits_start_idxs: list[int] = []
+        cur_start_idx = 0
+        cur_page = 1
+        prev_start_idx = 0
+
+        for split in splits:
+            if self.skip_empty_documents and len(split) == 0:
+                continue
+
+            # A splitting function may transform the text, so a split is not guaranteed to appear
+            # verbatim in the source. Locate it when possible and fall back to the running offset.
+            found_idx = content.find(split, cur_start_idx)
+            if found_idx == -1 and text_splits:
+                # The split may legitimately start before the end of the previous split, e.g. an
+                # overlapping/sliding-window splitting function. Retry allowing it to be found
+                # anywhere after the start of the previous split, so it isn't missed just because
+                # the forward-only search above already moved past it.
+                found_idx = content.find(split, prev_start_idx + 1)
+
+            if found_idx != -1:
+                cur_start_idx = found_idx
+                cur_page = 1 + content.count("\f", 0, found_idx)
+
+            text_splits.append(split)
+            splits_start_idxs.append(cur_start_idx)
+            splits_pages.append(cur_page)
+
+            prev_start_idx = cur_start_idx
+            cur_start_idx += len(split)
+            if found_idx == -1:
+                cur_page += split.count("\f")
+
+        metadata = deepcopy(doc.meta)
+        metadata["source_id"] = doc.id
+        return self._create_docs_from_splits(
+            text_splits=text_splits, splits_pages=splits_pages, splits_start_idxs=splits_start_idxs, meta=metadata
+        )
 
     def _concatenate_units(
         self, elements: list[str], split_length: int, split_overlap: int, split_threshold: int

releasenotes/notes/document-splitter-function-meta-8b0a9894c49b67c2.yaml

@@ -0,0 +1,8 @@
+---
+fixes:
+  - |
+    ``DocumentSplitter`` now adds ``split_id``, ``split_idx_start`` and ``page_number`` to the metadata of the chunks it
+    creates when ``split_by="function"``, like it already did for all other split modes. Without this metadata,
+    components that rely on it, such as ``SentenceWindowRetriever``, could not be used with chunks produced by a
+    custom splitting function. Empty chunks returned by the splitting function are now also skipped, unless
+    ``skip_empty_documents=False``.

test/components/preprocessors/test_document_splitter.py

@@ -8,6 +8,8 @@
 
 from haystack import Document
 from haystack.components.preprocessors import DocumentSplitter
+from haystack.components.retrievers import SentenceWindowRetriever
+from haystack.document_stores.in_memory import InMemoryDocumentStore
 from haystack.utils import deserialize_callable, serialize_callable
 
 
@@ -217,28 +219,107 @@ def test_split_by_function(self):
 
         assert len(docs) == 4
         assert docs[0].content == "This"
-        assert docs[0].meta == {"key": "value", "source_id": "1"}
+        assert docs[0].meta == {"key": "value", "source_id": "1", "split_id": 0, "split_idx_start": 0, "page_number": 1}
         assert docs[1].content == "Is"
-        assert docs[1].meta == {"key": "value", "source_id": "1"}
+        assert docs[1].meta == {"key": "value", "source_id": "1", "split_id": 1, "split_idx_start": 5, "page_number": 1}
         assert docs[2].content == "A"
-        assert docs[2].meta == {"key": "value", "source_id": "1"}
+        assert docs[2].meta == {"key": "value", "source_id": "1", "split_id": 2, "split_idx_start": 8, "page_number": 1}
         assert docs[3].content == "Test"
-        assert docs[3].meta == {"key": "value", "source_id": "1"}
+        assert docs[3].meta == {
+            "key": "value",
+            "source_id": "1",
+            "split_id": 3,
+            "split_idx_start": 10,
+            "page_number": 1,
+        }
 
         splitting_function = lambda s: re.split(r"[\s]{2,}", s)
         splitter = DocumentSplitter(split_by="function", splitting_function=splitting_function)
         text = "This       Is\n  A  Test"
         result = splitter.run(documents=[Document(id="1", content=text, meta={"key": "value"})])
         docs = result["documents"]
         assert len(docs) == 4
-        assert docs[0].content == "This"
-        assert docs[0].meta == {"key": "value", "source_id": "1"}
-        assert docs[1].content == "Is"
-        assert docs[1].meta == {"key": "value", "source_id": "1"}
-        assert docs[2].content == "A"
-        assert docs[2].meta == {"key": "value", "source_id": "1"}
-        assert docs[3].content == "Test"
-        assert docs[3].meta == {"key": "value", "source_id": "1"}
+        for split_id, (content, doc) in enumerate(zip(["This", "Is", "A", "Test"], docs, strict=True)):
+            assert doc.content == content
+            assert doc.meta["key"] == "value"
+            assert doc.meta["source_id"] == "1"
+            assert doc.meta["split_id"] == split_id
+            assert doc.meta["split_idx_start"] == text.index(content)
+            assert doc.meta["page_number"] == 1
+
+    def test_split_by_function_tracks_page_numbers(self):
+        splitter = DocumentSplitter(split_by="function", splitting_function=lambda s: s.split("\f"))
+        text = "First chunk.\fSecond chunk.\fThird chunk."
+        docs = splitter.run(documents=[Document(content=text)])["documents"]
+
+        assert [doc.meta["page_number"] for doc in docs] == [1, 2, 3]
+        assert [doc.meta["split_id"] for doc in docs] == [0, 1, 2]
+        assert [doc.meta["split_idx_start"] for doc in docs] == [text.index(doc.content) for doc in docs]
+
+    def test_split_by_function_with_transformed_splits(self):
+        # The splits don't appear verbatim in the source, so they cannot be located in it
+        splitter = DocumentSplitter(split_by="function", splitting_function=lambda s: [t.upper() for t in s.split(".")])
+        docs = splitter.run(documents=[Document(content="one.two")])["documents"]
+
+        assert [doc.content for doc in docs] == ["ONE", "TWO"]
+        assert [doc.meta["split_id"] for doc in docs] == [0, 1]
+        assert [doc.meta["split_idx_start"] for doc in docs] == [0, 3]
+        assert [doc.meta["page_number"] for doc in docs] == [1, 1]
+
+    def test_split_by_function_skips_empty_splits(self):
+        splitting_function = lambda s: s.split("\f")
+        splitter = DocumentSplitter(split_by="function", splitting_function=splitting_function)
+        docs = splitter.run(documents=[Document(content="a\f\fb")])["documents"]
+
+        assert [doc.content for doc in docs] == ["a", "b"]
+        assert [doc.meta["page_number"] for doc in docs] == [1, 3]
+
+        splitter = DocumentSplitter(
+            split_by="function", splitting_function=splitting_function, skip_empty_documents=False
+        )
+        docs = splitter.run(documents=[Document(content="a\f\fb")])["documents"]
+
+        assert [doc.content for doc in docs] == ["a", "", "b"]
+
+    def test_split_by_function_output_usable_by_sentence_window_retriever(self):
+        splitter = DocumentSplitter(split_by="function", splitting_function=lambda s: s.split("|"))
+        text = "first part|second part|third part"
+        docs = splitter.run(documents=[Document(content=text)])["documents"]
+
+        document_store = InMemoryDocumentStore()
+        document_store.write_documents(docs)
+        retriever = SentenceWindowRetriever(document_store=document_store, window_size=1)
+        result = retriever.run(retrieved_documents=[docs[1]])
+
+        assert result["context_windows"] == ["first partsecond partthird part"]
+        assert [doc.content for doc in result["context_documents"]] == ["first part", "second part", "third part"]
+
+    def test_split_by_function_with_overlapping_sliding_window(self):
+        # Hand-rolled sliding-window splitting function: each split shares words with the next one,
+        # so the splits genuinely overlap in the source text.
+        def sliding_window(text):
+            words = text.split(" ")
+            return [" ".join(words[i : i + 3]) for i in range(0, len(words), 2)]
+
+        text = "the quick brown fox jumps"
+        assert sliding_window(text) == ["the quick brown", "brown fox jumps", "jumps"]
+        # "brown fox jumps" really starts at index 10 in the source text
+        assert text.find("brown fox jumps") == 10
+
+        splitter = DocumentSplitter(split_by="function", splitting_function=sliding_window, split_overlap=1)
+        docs = splitter.run(documents=[Document(content=text)])["documents"]
+
+        assert [doc.content for doc in docs] == ["the quick brown", "brown fox jumps", "jumps"]
+
+        # content.find(split, cur_start_idx) only searches forward from the end of the previous
+        # split, so a split that genuinely starts *before* that point (as any overlapping split
+        # does) can never be located, and silently falls back to the wrong cumulative offset.
+        assert docs[1].meta["split_idx_start"] == 10  # real index of "brown fox jumps" in the source
+        assert docs[2].meta["split_idx_start"] == 20  # real index of "jumps" in the source
+
+        # The real overlap ("brown") between docs[0] and docs[1] should be detected since
+        # split_overlap=1 was requested.
+        assert docs[1].meta["_split_overlap"] != []
 
     def test_split_by_word_with_overlap(self):
         splitter = DocumentSplitter(split_by="word", split_length=10, split_overlap=2)