diff --git a/src/chonkie/chunker/table.py b/src/chonkie/chunker/table.py index 7f6335787..606df4467 100644 --- a/src/chonkie/chunker/table.py +++ b/src/chonkie/chunker/table.py @@ -144,24 +144,34 @@ def chunk(self, text: str) -> list[Chunk]: ), ] else: - # Track character position for data rows (after header) - header_len = len(header) - current_char_index = header_len + # Track character position over the original, unstripped input. + # The first span absorbs any prefix before the parsed header, and + # the final span absorbs the footer and trailing whitespace so the + # spans remain contiguous over the complete table text. + header_start = text.find(header) + header_span_len = len(header) + (header_start if header_start != -1 else 0) + current_char_index = 0 for i in range(0, len(data_rows), self.chunk_size): chunk_rows = data_rows[i : i + self.chunk_size] - chunk_text = header + "".join(chunk_rows) + footer - data_rows_len = len("".join(chunk_rows)) + chunk_rows_text = "".join(chunk_rows) + chunk_text = header + chunk_rows_text + footer + is_last_chunk = i + self.chunk_size >= len(data_rows) + span_len = ( + len(text) - current_char_index + if is_last_chunk + else len(chunk_rows_text) + (header_span_len if i == 0 else 0) + ) chunks.append( Chunk( text=chunk_text, token_count=len(chunk_rows), start_index=current_char_index, - end_index=current_char_index + data_rows_len, + end_index=current_char_index + span_len, ), ) - current_char_index += data_rows_len + current_char_index += span_len return chunks diff --git a/tests/chunkers/test_table_chunker.py b/tests/chunkers/test_table_chunker.py index 7a138904d..7c285f231 100644 --- a/tests/chunkers/test_table_chunker.py +++ b/tests/chunkers/test_table_chunker.py @@ -861,3 +861,35 @@ def test_table_chunker_markdown_document_no_tables_preserves_chunks() -> None: assert len(result.chunks) == 1 assert result.chunks[0].text == "Some prose text here." + + +def test_table_chunker_row_tokenizer_indices(): + table = "| A | B |\n|---|---|\n| 1 | 2 |\n| 3 | 4 |\n| 5 | 6 |\n| 7 | 8 |" + chunker = TableChunker(tokenizer="row", chunk_size=2) + chunks = chunker.chunk(table) + assert len(chunks) > 1 + # First chunk must start at 0 (matches the character path and every other chunker). + assert chunks[0].start_index == 0 + # Spans are contiguous and reconstruct the full table (no lost/duplicated regions). + for i in range(len(chunks) - 1): + assert chunks[i].end_index == chunks[i + 1].start_index + assert chunks[-1].end_index == len(table) + assert "".join(table[c.start_index : c.end_index] for c in chunks) == table + + +def test_table_chunker_row_tokenizer_html_indices_cover_full_input() -> None: + table = ( + " " + "" + "
A
1
2
3
\n" + ) + chunker = TableChunker(tokenizer="row", chunk_size=1) + + chunks = chunker.chunk(table) + + assert len(chunks) == 3 + assert chunks[0].start_index == 0 + for current, following in zip(chunks, chunks[1:]): + assert current.end_index == following.start_index + assert chunks[-1].end_index == len(table) + assert "".join(table[chunk.start_index : chunk.end_index] for chunk in chunks) == table