Coverage for src/dataknobs_data/pandas/batch_ops.py: 0%

1"""Batch operations for DataKnobs-Pandas integration.""" 

2 

3import logging 

4from dataclasses import dataclass 

5from typing import Any, Callable, Dict, Generator, List, Optional, Union 

6 

7import pandas as pd 

8import numpy as np 

9 

10from dataknobs_data.records import Record 

11from dataknobs_data.query import Query 

12from dataknobs_data.database import AsyncDatabase, SyncDatabase 

13from .converter import DataFrameConverter, ConversionOptions 

14 

15logger = logging.getLogger(__name__) 

16 

17 

18@dataclass 

19class BatchConfig: 

20 """Configuration for batch operations.""" 

21 chunk_size: int = 1000 

22 parallel: bool = False 

23 max_workers: int = 4 

24 progress_callback: Optional[Callable[[int, int], None]] = None 

25 error_handling: str = "raise" # "raise", "skip", "log" 

26 memory_efficient: bool = True 

27 

28 def __post_init__(self): 

29 """Validate configuration parameters.""" 

30 if self.chunk_size <= 0: 

31 raise ValueError("chunk_size must be greater than 0") 

32 

33 if self.error_handling not in ("raise", "skip", "log"): 

34 raise ValueError("error_handling must be one of: 'raise', 'skip', 'log'") 

35 

36 

37class ChunkedProcessor: 

38 """Process DataFrames in chunks for memory efficiency.""" 

39 

40 def __init__(self, chunk_size: int = 1000): 

41 """Initialize chunked processor. 

42  

43 Args: 

44 chunk_size: Size of each chunk 

45 """ 

46 self.chunk_size = chunk_size 

47 

48 def process_dataframe( 

49 self, 

50 df: pd.DataFrame, 

51 processor: Callable[[pd.DataFrame], Any], 

52 combine: Optional[Callable[[List[Any]], Any]] = None 

53 ) -> Any: 

54 """Process DataFrame in chunks. 

55  

56 Args: 

57 df: DataFrame to process 

58 processor: Function to process each chunk 

59 combine: Function to combine results 

60  

61 Returns: 

62 Combined results or list of chunk results 

63 """ 

64 results = [] 

65 

66 for chunk in self.iter_chunks(df): 

67 result = processor(chunk) 

68 results.append(result) 

69 

70 if combine: 

71 return combine(results) 

72 return results 

73 

74 def iter_chunks(self, df: pd.DataFrame) -> Generator[pd.DataFrame, None, None]: 

75 """Iterate over DataFrame in chunks. 

76  

77 Args: 

78 df: DataFrame to chunk 

79  

80 Yields: 

81 DataFrame chunks 

82 """ 

83 for start_idx in range(0, len(df), self.chunk_size): 

84 end_idx = min(start_idx + self.chunk_size, len(df)) 

85 yield df.iloc[start_idx:end_idx] 

86 

87 def read_csv_chunked( 

88 self, 

89 filepath: str, 

90 processor: Callable[[pd.DataFrame], Any], 

91 **read_kwargs 

92 ) -> List[Any]: 

93 """Read CSV file in chunks and process. 

94  

95 Args: 

96 filepath: Path to CSV file 

97 processor: Function to process each chunk 

98 **read_kwargs: Additional arguments for pd.read_csv 

99  

100 Returns: 

101 List of processed results 

102 """ 

103 results = [] 

104 

105 for chunk in pd.read_csv(filepath, chunksize=self.chunk_size, **read_kwargs): 

106 result = processor(chunk) 

107 results.append(result) 

108 

109 return results 

110 

111 

112class BatchOperations: 

113 """Batch operations for DataKnobs databases using DataFrames.""" 

114 

115 def __init__( 

116 self, 

117 database: Union[AsyncDatabase, SyncDatabase], 

118 converter: Optional[DataFrameConverter] = None 

119 ): 

120 """Initialize batch operations. 

121  

122 Args: 

123 database: Target database 

124 converter: DataFrame converter 

125 """ 

126 self.database = database 

127 self.converter = converter or DataFrameConverter() 

128 self.is_async = hasattr(database, 'create') and asyncio.iscoroutinefunction(database.create) 

129 

130 def bulk_insert_dataframe( 

131 self, 

132 df: pd.DataFrame, 

133 config: Optional[BatchConfig] = None, 

134 conversion_options: Optional[ConversionOptions] = None 

135 ) -> Dict[str, Any]: 

136 """Bulk insert DataFrame rows into database. 

137  

138 Args: 

139 df: DataFrame to insert 

140 config: Batch configuration 

141 conversion_options: Options for DataFrame conversion 

142  

143 Returns: 

144 Insert statistics 

145 """ 

146 config = config or BatchConfig() 

147 conversion_options = conversion_options or ConversionOptions() 

148 

149 stats = { 

150 "total_rows": len(df), 

151 "inserted": 0, 

152 "failed": 0, 

153 "errors": [] 

154 } 

155 

156 # Process in chunks if memory efficient mode 

157 if config.memory_efficient and len(df) > config.chunk_size: 

158 processor = ChunkedProcessor(config.chunk_size) 

159 

160 def process_chunk(chunk_df: pd.DataFrame) -> Dict[str, int]: 

161 return self._insert_chunk(chunk_df, config, conversion_options) 

162 

163 chunk_results = processor.process_dataframe(df, process_chunk) 

164 

165 # Aggregate results 

166 for result in chunk_results: 

167 stats["inserted"] += result["inserted"] 

168 stats["failed"] += result["failed"] 

169 if "errors" in result: 

170 stats["errors"].extend(result["errors"]) 

171 else: 

172 # Process entire DataFrame at once 

173 stats = self._insert_chunk(df, config, conversion_options) 

174 

175 return stats 

176 

177 def query_as_dataframe( 

178 self, 

179 query: Query, 

180 conversion_options: Optional[ConversionOptions] = None 

181 ) -> pd.DataFrame: 

182 """Execute query and return results as DataFrame. 

183  

184 Args: 

185 query: Query to execute 

186 conversion_options: Options for conversion 

187  

188 Returns: 

189 Query results as DataFrame 

190 """ 

191 conversion_options = conversion_options or ConversionOptions() 

192 

193 # Execute query 

194 if self.is_async: 

195 import asyncio 

196 records = asyncio.run(self.database.search(query)) 

197 else: 

198 records = self.database.search(query) 

199 

200 # Convert to DataFrame 

201 return self.converter.records_to_dataframe(records, conversion_options) 

202 

203 def update_from_dataframe( 

204 self, 

205 df: pd.DataFrame, 

206 id_column: str, 

207 config: Optional[BatchConfig] = None, 

208 conversion_options: Optional[ConversionOptions] = None 

209 ) -> Dict[str, Any]: 

210 """Update records from DataFrame using ID column. 

211  

212 Args: 

213 df: DataFrame with updates 

214 id_column: Column containing record IDs 

215 config: Batch configuration 

216 conversion_options: Conversion options 

217  

218 Returns: 

219 Update statistics 

220 """ 

221 config = config or BatchConfig() 

222 conversion_options = conversion_options or ConversionOptions() 

223 

224 stats = { 

225 "total_rows": len(df), 

226 "updated": 0, 

227 "failed": 0, 

228 "not_found": 0, 

229 "errors": [] 

230 } 

231 

232 # Ensure ID column exists 

233 if id_column not in df.columns: 

234 raise ValueError(f"ID column '{id_column}' not found in DataFrame") 

235 

236 # Convert DataFrame to records  

237 records = self.converter.dataframe_to_records(df, conversion_options) 

238 

239 # Prepare updates as (id, record) tuples 

240 updates = [] 

241 for i, record in enumerate(records): 

242 record_id = str(df.iloc[i][id_column]) 

243 updates.append((record_id, record)) 

244 

245 # Process updates in chunks 

246 for i in range(0, len(updates), config.chunk_size): 

247 chunk = updates[i:i + config.chunk_size] 

248 

249 try: 

250 # Use batch update for better performance 

251 if self.is_async: 

252 import asyncio 

253 results = asyncio.run(self.database.update_batch(chunk)) 

254 else: 

255 results = self.database.update_batch(chunk) 

256 

257 # Count successes and failures 

258 for success in results: 

259 if success: 

260 stats["updated"] += 1 

261 else: 

262 stats["not_found"] += 1 

263 

264 except Exception as e: 

265 # If batch fails, try individual updates 

266 if config.error_handling == "raise": 

267 raise 

268 

269 for record_id, record in chunk: 

270 try: 

271 if self.is_async: 

272 import asyncio 

273 success = asyncio.run(self.database.update(record_id, record)) 

274 else: 

275 success = self.database.update(record_id, record) 

276 

277 if success: 

278 stats["updated"] += 1 

279 else: 

280 stats["not_found"] += 1 

281 

282 except Exception as e: 

283 stats["failed"] += 1 

284 if config.error_handling == "log": 

285 logger.error(f"Failed to update record {record_id}: {e}") 

286 stats["errors"].append(str(e)) 

287 # else "skip" 

288 

289 # Progress callback 

290 if config.progress_callback: 

291 processed = stats["updated"] + stats["failed"] + stats["not_found"] 

292 config.progress_callback(processed, len(updates)) 

293 

294 return stats 

295 

296 def aggregate( 

297 self, 

298 query: Query, 

299 aggregations: Dict[str, Union[str, Callable]], 

300 group_by: Optional[List[str]] = None 

301 ) -> pd.DataFrame: 

302 """Perform aggregations on query results. 

303  

304 Args: 

305 query: Query to execute 

306 aggregations: Dictionary of column: aggregation function 

307 group_by: Columns to group by 

308  

309 Returns: 

310 Aggregated DataFrame 

311 """ 

312 # Get data as DataFrame 

313 df = self.query_as_dataframe(query) 

314 

315 if df.empty: 

316 return pd.DataFrame() 

317 

318 # Perform aggregations 

319 if group_by: 

320 grouped = df.groupby(group_by) 

321 return grouped.agg(aggregations) 

322 else: 

323 # Single row with aggregations 

324 result = {} 

325 for col, agg_func in aggregations.items(): 

326 if col in df.columns: 

327 if isinstance(agg_func, str): 

328 result[f"{col}_{agg_func}"] = df[col].agg(agg_func) 

329 else: 

330 result[f"{col}_agg"] = agg_func(df[col]) 

331 return pd.DataFrame([result]) 

332 

333 def transform_and_save( 

334 self, 

335 query: Query, 

336 transformer: Callable[[pd.DataFrame], pd.DataFrame], 

337 config: Optional[BatchConfig] = None 

338 ) -> Dict[str, Any]: 

339 """Query, transform with pandas, and save back. 

340  

341 Args: 

342 query: Query to get records 

343 transformer: Function to transform DataFrame 

344 config: Batch configuration 

345  

346 Returns: 

347 Operation statistics 

348 """ 

349 config = config or BatchConfig() 

350 

351 # Get data 

352 df = self.query_as_dataframe(query) 

353 

354 if df.empty: 

355 return {"total_rows": 0, "transformed": 0} 

356 

357 # Apply transformation 

358 transformed_df = transformer(df) 

359 

360 # Save back if index preserved (has record IDs) 

361 if df.index.name == "record_id" and transformed_df.index.name == "record_id": 

362 return self.update_from_dataframe( 

363 transformed_df, 

364 id_column=None, # Use index 

365 config=config 

366 ) 

367 else: 

368 # Insert as new records 

369 return self.bulk_insert_dataframe(transformed_df, config) 

370 

371 def _insert_chunk( 

372 self, 

373 df: pd.DataFrame, 

374 config: BatchConfig, 

375 conversion_options: ConversionOptions 

376 ) -> Dict[str, Any]: 

377 """Insert a chunk of DataFrame rows. 

378  

379 Args: 

380 df: DataFrame chunk 

381 config: Batch configuration 

382 conversion_options: Conversion options 

383  

384 Returns: 

385 Insert statistics for chunk 

386 """ 

387 stats = { 

388 "total_rows": len(df), 

389 "inserted": 0, 

390 "failed": 0, 

391 "errors": [] 

392 } 

393 

394 # Convert to records 

395 records = self.converter.dataframe_to_records(df, conversion_options) 

396 

397 # Use batch creation for better performance with graceful fallback 

398 if hasattr(self.database, 'create_batch'): 

399 try: 

400 if self.is_async: 

401 import asyncio 

402 ids = asyncio.run(self.database.create_batch(records)) 

403 else: 

404 ids = self.database.create_batch(records) 

405 stats["inserted"] = len(ids) 

406 

407 # Progress callback for successful batch 

408 if config.progress_callback: 

409 config.progress_callback(len(records), len(records)) 

410 

411 except Exception as batch_error: 

412 # Batch failed, try individual records to identify failures 

413 for i, record in enumerate(records): 

414 try: 

415 if self.is_async: 

416 import asyncio 

417 asyncio.run(self.database.create(record)) 

418 else: 

419 self.database.create(record) 

420 stats["inserted"] += 1 

421 

422 except Exception as record_error: 

423 stats["failed"] += 1 

424 

425 # Handle error based on config 

426 if config.error_handling == "raise": 

427 raise 

428 elif config.error_handling == "log": 

429 logger.error(f"Failed to insert row {i}: {record_error}") 

430 stats["errors"].append(str(record_error)) 

431 # else "skip" - just continue 

432 

433 # Progress callback for each record 

434 if config.progress_callback: 

435 config.progress_callback(i + 1, len(records)) 

436 else: 

437 # Fallback to individual inserts if create_batch not available 

438 for i, record in enumerate(records): 

439 try: 

440 if self.is_async: 

441 import asyncio 

442 asyncio.run(self.database.create(record)) 

443 else: 

444 self.database.create(record) 

445 stats["inserted"] += 1 

446 

447 except Exception as e: 

448 stats["failed"] += 1 

449 if config.error_handling == "raise": 

450 raise 

451 elif config.error_handling == "log": 

452 logger.error(f"Failed to insert row {i}: {e}") 

453 stats["errors"].append(str(e)) 

454 # else "skip" 

455 

456 # Progress callback 

457 if config.progress_callback: 

458 config.progress_callback(i + 1, len(records)) 

459 

460 return stats 

461 

462 def export_to_csv( 

463 self, 

464 query: Query, 

465 filepath: str, 

466 conversion_options: Optional[ConversionOptions] = None, 

467 **to_csv_kwargs 

468 ) -> None: 

469 """Export query results to CSV file. 

470  

471 Args: 

472 query: Query to execute 

473 filepath: Output file path 

474 conversion_options: Conversion options 

475 **to_csv_kwargs: Additional arguments for DataFrame.to_csv 

476 """ 

477 df = self.query_as_dataframe(query, conversion_options) 

478 df.to_csv(filepath, **to_csv_kwargs) 

479 

480 def export_to_parquet( 

481 self, 

482 query: Query, 

483 filepath: str, 

484 conversion_options: Optional[ConversionOptions] = None, 

485 **to_parquet_kwargs 

486 ) -> None: 

487 """Export query results to Parquet file. 

488  

489 Args: 

490 query: Query to execute 

491 filepath: Output file path 

492 conversion_options: Conversion options 

493 **to_parquet_kwargs: Additional arguments for DataFrame.to_parquet 

494 """ 

495 df = self.query_as_dataframe(query, conversion_options) 

496 df.to_parquet(filepath, **to_parquet_kwargs) 

497 

498 

499# Import asyncio only if needed 

500try: 

501 import asyncio 

502except ImportError: 

503 asyncio = None