Coverage for packages/dqm-ml-job/src/dqm_ml_job/dataloaders/proto.py: 100%

1"""Protocol definitions for data loaders and selections. 

2 

3This module contains the DataLoader and DataSelection protocol classes 

4that define the interface for data loading implementations. 

5""" 

6 

7from typing import Any, Protocol, runtime_checkable 

8 

9 

10@runtime_checkable 

11class DataSelection(Protocol): 

12 """ 

13 Protocol for a specific subset of data discovered by a DataLoader. 

14 

15 A DataSelection represents a concrete set of samples (e.g., a 

16 specific folder, a filtered view of a database, or a single file) 

17 and provides an iterator over data batches. 

18 """ 

19 

20 name: str 

21 

22 def bootstrap(self, columns_list: list[str]) -> None: 

23 """Perform initial setup for the selection before iteration. 

24 

25 Args: 

26 columns_list: List of column names to load. 

27 """ 

28 

29 def get_nb_batches(self) -> int: 

30 """ 

31 Return the estimated number of batches in this selection. 

32 

33 Used primarily for progress bar estimation. 

34 """ 

35 

36 def __iter__(self) -> Any: 

37 """ 

38 Iterate over the selection, yielding pyarrow.RecordBatch objects. 

39 """ 

40 

41 

42@runtime_checkable 

43class DataLoader(Protocol): 

44 """ 

45 Protocol for Data Loader factories. 

46 

47 A DataLoader is responsible for scanning a source (disk, DB, S3) and 

48 discovering available DataSelections based on its configuration. 

49 """ 

50 

51 def get_selections(self) -> list[DataSelection]: 

52 """ 

53 Discover and return the list of available selections for this loader. 

54 

55 Returns: 

56 A list of initialized DataSelection instances. 

57 """