Coverage for packages/dqm-ml-core/src/dqm_ml_core/api/metrics_processor.py: 94%

1"""Metrics processor base class. 

2 

3This module contains the MetricsProcessor base class that all 

4metric computation processors must inherit from. 

5""" 

6 

7import logging 

8from typing import Any 

9 

10import pyarrow as pa 

11 

12from dqm_ml_core.api.processor import Processor 

13from dqm_ml_core.utils.matching import resolve_include_exclude 

14 

15logger = logging.getLogger(__name__) 

16 

17 

18class MetricsProcessor(Processor): 

19 """ 

20 Base class for all metric computation processors. 

21 

22 Metric processors compute dataset-level scores (e.g., completeness, 

23 diversity, representativeness). The primary lifecycle methods are 

24 ``select_columns`` (per-batch column selection), ``compute_batch_metric`` 

25 (batch aggregation), and ``compute`` (final dataset-level computation). 

26 """ 

27 

28 def generated_metrics(self) -> list[str]: 

29 """ 

30 Return the names of the final metrics produced by this processor. 

31 

32 Returns: 

33 A list of metric names. 

34 """ 

35 outputs = getattr(self, "output_metrics", {}) 

36 return list(outputs.values()) 

37 

38 def select_columns(self, batch: pa.RecordBatch, prev_features: dict[str, pa.Array]) -> dict[str, pa.Array]: 

39 """ 

40 Select relevant columns from a raw batch for metric computation. 

41 

42 Args: 

43 batch: The input pyarrow RecordBatch. 

44 prev_features: Features already computed by preceding processors. 

45 

46 Returns: 

47 A dictionary mapping column names to pyarrow Arrays. 

48 """ 

49 features = {} 

50 

51 available = batch.schema.names 

52 cols = resolve_include_exclude( 

53 self.input_columns, 

54 self.exclude_columns, 

55 available, 

56 ) 

57 

58 for col in cols: 

59 if col in prev_features: 

60 continue 

61 

62 if col not in available: 

63 if ( 

64 self.errors_config 

65 and self.errors_config.tabular 

66 and self.errors_config.tabular.on_missing_column == "fail_fast" 

67 ): 

68 raise KeyError(f"Column '{col}' not found in batch") 

69 logger.warning(f"[{self.name}] column '{col}' not found in batch") 

70 continue 

71 features[col] = batch.column(col) 

72 

73 return features 

74 

75 def compute_batch_metric(self, features: dict[str, pa.Array]) -> dict[str, pa.Array]: 

76 """ 

77 Aggregate features into intermediate statistics for the current batch. 

78 

79 This method is critical for scalability. It should return a compact 

80 representation of the data (e.g., partial sums) that can be 

81 efficiently combined later. 

82 

83 Args: 

84 features: Dictionary of feature arrays computed on the batch. 

85 

86 Returns: 

87 A dictionary of aggregated statistics. 

88 """ 

89 return {} 

90 

91 def compute(self, batch_metrics: dict[str, pa.Array]) -> dict[str, Any]: # NOSONAR 

92 """ 

93 Perform the final dataset-level metric calculation. 

94 

95 Args: 

96 batch_metrics: The aggregated intermediate statistics from all batches. 

97 

98 Returns: 

99 A dictionary containing the final metrics. 

100 """ 

101 # SonarQube raises a warning because batch_metrics is not used. 

102 # It is irrelevant because compute is implemented in child classes which use batch_metrics. 

103 return {}