| 361 | ) |
| 362 | |
| 363 | async def read(self, batch_size: Optional[int] = None) -> List[Task]: |
| 364 | await self.prepare() |
| 365 | if self.stopped: |
| 366 | raise StopAsyncIteration() |
| 367 | if self.offset > self.total_samples: |
| 368 | raise StopAsyncIteration() |
| 369 | batch_size = self.batch_size if batch_size is None else batch_size |
| 370 | |
| 371 | table_cls = self.table_model_cls |
| 372 | |
| 373 | async def operation(session: AsyncSession): |
| 374 | stmt = ( |
| 375 | select(table_cls) |
| 376 | .where(table_cls.id > self.offset) |
| 377 | .order_by(asc(table_cls.id)) |
| 378 | .limit(batch_size) |
| 379 | ) |
| 380 | result = await session.execute(stmt) |
| 381 | results = result.scalars().all() |
| 382 | if len(results) == 0: |
| 383 | raise StopAsyncIteration() |
| 384 | if not self.is_eval and len(results) < batch_size: |
| 385 | raise StopAsyncIteration() |
| 386 | return results[-1].id, [self.formatter.format(item.raw_task) for item in results] |
| 387 | |
| 388 | self.offset, tasks = await async_run_with_retry_session( |
| 389 | self.session, operation, self.max_retry_times, self.max_retry_interval |
| 390 | ) |
| 391 | return tasks |
| 392 | |
| 393 | @classmethod |
| 394 | async def load_from_dataset(cls, dataset: Dataset, config: StorageConfig) -> "SQLTaskStorage": |