diff --git a/haystack/components/preprocessors/csv_document_splitter.py b/haystack/components/preprocessors/csv_document_splitter.py index ea1dfa78c5..3850e4ba82 100644 --- a/haystack/components/preprocessors/csv_document_splitter.py +++ b/haystack/components/preprocessors/csv_document_splitter.py @@ -6,9 +6,10 @@ from io import StringIO from typing import Any, Literal, get_args -from haystack import Document, component, logging from haystack.lazy_imports import LazyImport +from haystack import Document, component, logging + with LazyImport("Run 'pip install pandas'") as pandas_import: import pandas as pd @@ -145,8 +146,15 @@ def run(self, documents: list[Document]) -> dict[str, list[Document]]: ) continue - # Sort split_dfs first by row index, then by column index - split_dfs.sort(key=lambda dataframe: (dataframe.index[0], dataframe.columns[0])) + # Columns are only positional when ``header=None``. A caller passing + # ``read_csv_kwargs={"header": 0}`` gets the first row as string labels, + # so ``int(label)`` failed for every sub-table and sorting on the label + # ordered the sub-tables alphabetically. Map labels back to their + # position in the original frame instead. + column_positions = {label: position for position, label in enumerate(df.columns)} + + # Sort split_dfs first by row index, then by column position + split_dfs.sort(key=lambda dataframe: (dataframe.index[0], column_positions[dataframe.columns[0]])) for split_id, split_df in enumerate(split_dfs): split_documents.append( @@ -156,7 +164,7 @@ def run(self, documents: list[Document]) -> dict[str, list[Document]]: **deepcopy(document.meta), "source_id": document.id, "row_idx_start": int(split_df.index[0]), - "col_idx_start": int(split_df.columns[0]), + "col_idx_start": column_positions[split_df.columns[0]], "split_id": split_id, }, ) diff --git a/releasenotes/notes/csv-splitter-header-kwargs-fc8ce7f31997395f.yaml b/releasenotes/notes/csv-splitter-header-kwargs-fc8ce7f31997395f.yaml new file mode 100644 index 0000000000..5045e559c4 --- /dev/null +++ b/releasenotes/notes/csv-splitter-header-kwargs-fc8ce7f31997395f.yaml @@ -0,0 +1,9 @@ +--- +fixes: + - | + ``CSVDocumentSplitter`` now reports the column position of each sub-table and orders + sub-tables by column position instead of by column label. Passing + ``read_csv_kwargs={"header": 0}`` used to raise ``ValueError: invalid literal for int() + with base 10`` for every sub-table, because ``header=None`` was the only case in which + the labels were positional; with a caller-supplied header the labels are strings, and + sorting on them reordered the sub-tables alphabetically. diff --git a/test/components/preprocessors/test_csv_document_splitter_header_kwargs.py b/test/components/preprocessors/test_csv_document_splitter_header_kwargs.py new file mode 100644 index 0000000000..0e27ff35d1 --- /dev/null +++ b/test/components/preprocessors/test_csv_document_splitter_header_kwargs.py @@ -0,0 +1,28 @@ +"""read_csv_kwargs that sets a header must not break CSV splitting.""" + +from haystack import Document +from haystack.components.preprocessors.csv_document_splitter import CSVDocumentSplitter + +ROWS = "name,score\nAda,9\n\nBob,8" +COLUMNS = "z,_,a\n1,,2\n3,,4\n" + + +def test_default_header_none_keeps_positional_columns(): + splitter = CSVDocumentSplitter(row_split_threshold=1, column_split_threshold=None) + result = splitter.run([Document(content=ROWS)]) + assert [(d.meta["row_idx_start"], d.meta["col_idx_start"]) for d in result["documents"]] == [(0, 0), (3, 0)] + + +def test_read_csv_kwargs_header_reports_column_positions(): + """With a caller-supplied header the columns are labels, not positions.""" + splitter = CSVDocumentSplitter(row_split_threshold=1, column_split_threshold=None, read_csv_kwargs={"header": 0}) + result = splitter.run([Document(content=ROWS)]) + assert [(d.meta["row_idx_start"], d.meta["col_idx_start"]) for d in result["documents"]] == [(0, 0), (2, 0)] + + +def test_header_labels_do_not_reorder_sub_tables(): + """Sub-tables keep the original column order, not the alphabetical one.""" + splitter = CSVDocumentSplitter(row_split_threshold=None, column_split_threshold=1, read_csv_kwargs={"header": 0}) + result = splitter.run([Document(content=COLUMNS)]) + assert [(d.meta["col_idx_start"], d.meta["split_id"]) for d in result["documents"]] == [(0, 0), (2, 1)] + assert [d.content.strip().replace("\n", "/") for d in result["documents"]] == ["1/3", "2/4"]