Coverage for llm_dataset_engine/stages/data_loader_stage.py: 34%

1"""Data loading stage for reading tabular data."""

2

3 from decimal import Decimal

4 from typing import Any

5

6 import pandas as pd

7

8 from llm_dataset_engine.adapters.data_io import create_data_reader

9 from llm_dataset_engine.core.models import CostEstimate, ValidationResult

10 from llm_dataset_engine.core.specifications import DatasetSpec

11 from llm_dataset_engine.stages.pipeline_stage import PipelineStage

12

13

14 class DataLoaderStage(PipelineStage[DatasetSpec, pd.DataFrame]):

15 """

16 Load data from source and validate schema.

17

18 Responsibilities:

19 - Read data from configured source

20 - Validate required columns exist

21 - Apply any filters

22 - Update context with row count

23 """

24

25 def __init__(self, dataframe: pd.DataFrame | None = None):

26 """

27 Initialize data loader stage.

28

29 Args:

30 dataframe: Optional pre-loaded dataframe (for DataFrame source)

31 """

32 super().__init__("DataLoader")

33 self.dataframe = dataframe

34

35 def process(self, spec: DatasetSpec, context: Any) -> pd.DataFrame:

36 """Load data from source."""

37 # Create appropriate reader

38 reader = create_data_reader(

39 source_type=spec.source_type,

40 source_path=spec.source_path,

41 dataframe=self.dataframe,

42 delimiter=spec.delimiter,

43 encoding=spec.encoding,

44 sheet_name=spec.sheet_name,

45 )

46

47 # Read data

48 df = reader.read()

49

50 # Validate columns exist

51 missing_cols = set(spec.input_columns) - set(df.columns)

52 if missing_cols:

53 raise ValueError(f"Missing columns: {missing_cols}")

54

55 # Apply filters if specified

56 if spec.filters:

57 for column, value in spec.filters.items():

58 if column in df.columns:

59 df = df[df[column] == value]

60

61 # Update context with total rows

62 context.total_rows = len(df)

63

64 self.logger.info(f"Loaded {len(df)} rows from {spec.source_type}")

65

66 return df

67

68 def validate_input(self, spec: DatasetSpec) -> ValidationResult:

69 """Validate dataset specification."""

70 result = ValidationResult(is_valid=True)

71

72 # Check file exists for file sources

73 if spec.source_path and not spec.source_path.exists():

74 result.add_error(f"Source file not found: {spec.source_path}")

75

76 # Check input columns specified

77 if not spec.input_columns:

78 result.add_error("No input columns specified")

79

80 # Check output columns specified

81 if not spec.output_columns:

82 result.add_error("No output columns specified")

83

84 return result

85

86 def estimate_cost(self, spec: DatasetSpec) -> CostEstimate:

87 """Data loading has no LLM cost."""

88 return CostEstimate(

89 total_cost=Decimal("0.0"),

90 total_tokens=0,

91 input_tokens=0,

92 output_tokens=0,

93 rows=0,

94 )

95