Coverage for haystack/components/preprocessors/csv_document_splitter.py: 96%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from copy import deepcopy 

6from io import StringIO 

7from typing import Any, Literal, get_args 

8 

9from haystack import Document, component, logging 

10from haystack.lazy_imports import LazyImport 

11 

12with LazyImport("Run 'pip install pandas'") as pandas_import: 

13 import pandas as pd 

14 

15logger = logging.getLogger(__name__) 

16 

17SplitMode = Literal["threshold", "row-wise"] 

18 

19 

20@component 

21class CSVDocumentSplitter: 

22 """ 

23 A component for splitting CSV documents into sub-tables based on split arguments. 

24 

25 The splitter supports two modes of operation: 

26 - identify consecutive empty rows or columns that exceed a given threshold 

27 and uses them as delimiters to segment the document into smaller tables. 

28 - split each row into a separate sub-table, represented as a Document. 

29 

30 """ 

31 

32 def __init__( 

33 self, 

34 row_split_threshold: int | None = 2, 

35 column_split_threshold: int | None = 2, 

36 read_csv_kwargs: dict[str, Any] | None = None, 

37 split_mode: SplitMode = "threshold", 

38 ) -> None: 

39 """ 

40 Initializes the CSVDocumentSplitter component. 

41 

42 :param row_split_threshold: The minimum number of consecutive empty rows required to trigger a split. 

43 :param column_split_threshold: The minimum number of consecutive empty columns required to trigger a split. 

44 :param read_csv_kwargs: Additional keyword arguments to pass to `pandas.read_csv`. 

45 By default, the component with options: 

46 - `header=None` 

47 - `skip_blank_lines=False` to preserve blank lines 

48 - `dtype=object` to prevent type inference (e.g., converting numbers to floats). 

49 See https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html for more information. 

50 :param split_mode: 

51 If `threshold`, the component will split the document based on the number of 

52 consecutive empty rows or columns that exceed the `row_split_threshold` or `column_split_threshold`. 

53 If `row-wise`, the component will split each row into a separate sub-table. 

54 """ 

55 pandas_import.check() 

56 if split_mode not in get_args(SplitMode): 

57 raise ValueError( 

58 f"Split mode '{split_mode}' not recognized. Choose one among: {', '.join(get_args(SplitMode))}." 

59 ) 

60 if row_split_threshold is not None and row_split_threshold < 1: 

61 raise ValueError("row_split_threshold must be greater than 0") 

62 

63 if column_split_threshold is not None and column_split_threshold < 1: 

64 raise ValueError("column_split_threshold must be greater than 0") 

65 

66 if row_split_threshold is None and column_split_threshold is None: 

67 raise ValueError("At least one of row_split_threshold or column_split_threshold must be specified.") 

68 

69 self.row_split_threshold = row_split_threshold 

70 self.column_split_threshold = column_split_threshold 

71 self.read_csv_kwargs = read_csv_kwargs or {} 

72 self.split_mode = split_mode 

73 

74 @component.output_types(documents=list[Document]) 

75 def run(self, documents: list[Document]) -> dict[str, list[Document]]: 

76 """ 

77 Processes and splits a list of CSV documents into multiple sub-tables. 

78 

79 **Splitting Process:** 

80 1. Applies a row-based split if `row_split_threshold` is provided. 

81 2. Applies a column-based split if `column_split_threshold` is provided. 

82 3. If both thresholds are specified, performs a recursive split by rows first, then columns, ensuring 

83 further fragmentation of any sub-tables that still contain empty sections. 

84 4. Sorts the resulting sub-tables based on their original positions within the document. 

85 

86 :param documents: A list of Documents containing CSV-formatted content. 

87 Each document is assumed to contain one or more tables separated by empty rows or columns. 

88 

89 :return: 

90 A dictionary with a key `"documents"`, mapping to a list of new `Document` objects, 

91 each representing an extracted sub-table from the original CSV. 

92 The metadata of each document includes: 

93 - A field `source_id` to track the original document. 

94 - A field `row_idx_start` to indicate the starting row index of the sub-table in the original table. 

95 - A field `col_idx_start` to indicate the starting column index of the sub-table in the original table. 

96 - A field `split_id` to indicate the order of the split in the original document. 

97 - All other metadata copied from the original document. 

98 

99 - If a document cannot be processed, it is returned unchanged. 

100 - The `meta` field from the original document is preserved in the split documents. 

101 """ 

102 if len(documents) == 0: 

103 return {"documents": documents} 

104 

105 resolved_read_csv_kwargs = {"header": None, "skip_blank_lines": False, "dtype": object, **self.read_csv_kwargs} 

106 

107 split_documents = [] 

108 split_dfs = [] 

109 for document in documents: 

110 try: 

111 df = pd.read_csv(StringIO(document.content), **resolved_read_csv_kwargs) 

112 except Exception as e: 

113 logger.exception( 

114 "Error processing document {document_id}. Keeping it, but skipping splitting. Error: {error}", 

115 document_id=document.id, 

116 error=e, 

117 ) 

118 split_documents.append(document) 

119 continue 

120 

121 if self.split_mode == "row-wise": 

122 # each row is a separate sub-table 

123 split_dfs = self._split_by_row(df=df) 

124 

125 elif self.split_mode == "threshold": 

126 if self.row_split_threshold is not None and self.column_split_threshold is None: 

127 # split by rows 

128 split_dfs = self._split_dataframe(df=df, split_threshold=self.row_split_threshold, axis="row") 

129 elif self.column_split_threshold is not None and self.row_split_threshold is None: 

130 # split by columns 

131 split_dfs = self._split_dataframe(df=df, split_threshold=self.column_split_threshold, axis="column") 

132 else: 

133 # recursive split 

134 split_dfs = self._recursive_split( 

135 df=df, 

136 row_split_threshold=self.row_split_threshold, # type: ignore 

137 column_split_threshold=self.column_split_threshold, # type: ignore 

138 ) 

139 

140 # check if no sub-tables were found 

141 if len(split_dfs) == 0: 

142 logger.warning( 

143 "No sub-tables found while splitting CSV Document with id {doc_id}. Skipping document.", 

144 doc_id=document.id, 

145 ) 

146 continue 

147 

148 # Sort split_dfs first by row index, then by column index 

149 split_dfs.sort(key=lambda dataframe: (dataframe.index[0], dataframe.columns[0])) 

150 

151 for split_id, split_df in enumerate(split_dfs): 

152 split_documents.append( 

153 Document( 

154 content=split_df.to_csv(index=False, header=False, lineterminator="\n"), 

155 meta={ 

156 **deepcopy(document.meta), 

157 "source_id": document.id, 

158 "row_idx_start": int(split_df.index[0]), 

159 "col_idx_start": int(split_df.columns[0]), 

160 "split_id": split_id, 

161 }, 

162 ) 

163 ) 

164 

165 return {"documents": split_documents} 

166 

167 @staticmethod 

168 def _find_split_indices( 

169 df: "pd.DataFrame", split_threshold: int, axis: Literal["row", "column"] 

170 ) -> list[tuple[int, int]]: 

171 """ 

172 Finds the indices of consecutive empty rows or columns in a DataFrame. 

173 

174 :param df: DataFrame to split. 

175 :param split_threshold: Minimum number of consecutive empty rows or columns to trigger a split. 

176 :param axis: Axis along which to find empty elements. Either "row" or "column". 

177 :return: List of zero-based positional indices for consecutive empty rows or columns. 

178 """ 

179 if axis == "row": 

180 empty_elements = [i for i, is_empty in enumerate(df.isnull().all(axis=1).tolist()) if is_empty] 

181 else: 

182 empty_elements = [i for i, is_empty in enumerate(df.isnull().all(axis=0).tolist()) if is_empty] 

183 

184 # If no empty elements found, return empty list 

185 if len(empty_elements) == 0: 

186 return [] 

187 

188 # Identify groups of consecutive empty elements 

189 split_indices = [] 

190 consecutive_count = 1 

191 start_index = empty_elements[0] 

192 

193 for i in range(1, len(empty_elements)): 

194 if empty_elements[i] == empty_elements[i - 1] + 1: 

195 consecutive_count += 1 

196 else: 

197 if consecutive_count >= split_threshold: 

198 split_indices.append((start_index, empty_elements[i - 1])) 

199 consecutive_count = 1 

200 start_index = empty_elements[i] 

201 

202 # Handle the last group of consecutive elements 

203 if consecutive_count >= split_threshold: 

204 split_indices.append((start_index, empty_elements[-1])) 

205 

206 return split_indices 

207 

208 def _split_dataframe( 

209 self, df: "pd.DataFrame", split_threshold: int, axis: Literal["row", "column"] 

210 ) -> list["pd.DataFrame"]: 

211 """ 

212 Splits a DataFrame into sub-tables based on consecutive empty rows or columns exceeding `split_threshold`. 

213 

214 :param df: DataFrame to split. 

215 :param split_threshold: Minimum number of consecutive empty rows or columns to trigger a split. 

216 :param axis: Axis along which to split. Either "row" or "column". 

217 :return: List of split DataFrames. 

218 """ 

219 # Find indices of consecutive empty rows or columns 

220 split_indices = self._find_split_indices(df=df, split_threshold=split_threshold, axis=axis) 

221 

222 # If no split_indices are found, return the original DataFrame 

223 if len(split_indices) == 0: 

224 return [df] 

225 

226 # Split the DataFrame at identified indices 

227 sub_tables = [] 

228 table_start_idx = 0 

229 df_length = df.shape[0] if axis == "row" else df.shape[1] 

230 for empty_start_idx, empty_end_idx in split_indices + [(df_length, df_length)]: 

231 # Avoid empty splits 

232 if empty_start_idx - table_start_idx >= 1: 

233 if axis == "row": 

234 sub_table = df.iloc[table_start_idx:empty_start_idx] 

235 else: 

236 sub_table = df.iloc[:, table_start_idx:empty_start_idx] 

237 if not sub_table.empty: 

238 sub_tables.append(sub_table) 

239 table_start_idx = empty_end_idx + 1 

240 

241 return sub_tables 

242 

243 def _recursive_split( 

244 self, df: "pd.DataFrame", row_split_threshold: int, column_split_threshold: int 

245 ) -> list["pd.DataFrame"]: 

246 """ 

247 Recursively splits a DataFrame. 

248 

249 Recursively splits a DataFrame first by empty rows, then by empty columns, and repeats the process 

250 until no more splits are possible. Returns a list of DataFrames, each representing a fully separated sub-table. 

251 

252 :param df: A Pandas DataFrame representing a table (or multiple tables) extracted from a CSV. 

253 :param row_split_threshold: The minimum number of consecutive empty rows required to trigger a split. 

254 :param column_split_threshold: The minimum number of consecutive empty columns to trigger a split. 

255 """ 

256 

257 # Step 1: Split by rows 

258 new_sub_tables = self._split_dataframe(df=df, split_threshold=row_split_threshold, axis="row") 

259 

260 # Step 2: Split by columns 

261 final_tables = [] 

262 for table in new_sub_tables: 

263 final_tables.extend(self._split_dataframe(df=table, split_threshold=column_split_threshold, axis="column")) 

264 

265 # Step 3: Recursively reapply splitting checked by whether any new empty rows appear after column split 

266 result = [] 

267 for table in final_tables: 

268 # Check if there are consecutive rows >= row_split_threshold now present 

269 if len(self._find_split_indices(df=table, split_threshold=row_split_threshold, axis="row")) > 0: 

270 result.extend( 

271 self._recursive_split( 

272 df=table, row_split_threshold=row_split_threshold, column_split_threshold=column_split_threshold 

273 ) 

274 ) 

275 else: 

276 result.append(table) 

277 

278 return result 

279 

280 def _split_by_row(self, df: "pd.DataFrame") -> list["pd.DataFrame"]: 

281 """Split each CSV row into a separate subtable""" 

282 split_dfs = [] 

283 for idx, row in enumerate(df.itertuples(index=False)): 

284 split_df = pd.DataFrame(row).T 

285 split_df.index = [idx] # Set the index of the new DataFrame to idx 

286 split_dfs.append(split_df) 

287 return split_dfs