Coverage for packages/dqm-ml-core/src/dqm_ml_core/api/gap_processor.py: 100%

1"""Gap processor base class. 

2 

3This module contains the GapProcessor base class that all 

4domain gap 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 GapProcessor(Processor): 

19 """ 

20 Base class for all domain gap processors. 

21 

22 Gap processors compute distribution shift between two datasets 

23 (e.g., MMD, FID, KL divergence). The primary lifecycle methods are 

24 ``select_features`` (per-batch column selection aware of previous features), 

25 ``compute_batch_metric`` (batch aggregation), ``compute`` (final dataset-level 

26 statistics), and ``compute_delta`` (pairwise comparison). 

27 """ 

28 

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

30 """ 

31 Extract relevant columns from a batch, resolving patterns against 

32 both batch columns and previously computed upstream features. 

33 

34 Args: 

35 batch: The input pyarrow RecordBatch. 

36 prev_features: Features already computed by preceding processors. 

37 

38 Returns: 

39 A dictionary mapping column names to pyarrow Arrays. 

40 """ 

41 available = list(prev_features.keys()) + batch.schema.names 

42 cols = resolve_include_exclude( 

43 self.input_columns, 

44 self.exclude_columns, 

45 available, 

46 ) 

47 

48 features: dict[str, pa.Array] = {} 

49 for col in cols: 

50 if col in prev_features: 

51 continue 

52 if col not in batch.schema.names: 

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

54 continue 

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

56 

57 return features 

58 

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

60 """ 

61 Aggregate features into intermediate statistics for the current batch. 

62 

63 Args: 

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

65 

66 Returns: 

67 A dictionary of aggregated statistics. 

68 """ 

69 return {} 

70 

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

72 """ 

73 Perform the final dataset-level aggregation of batch statistics. 

74 

75 Args: 

76 batch_metrics: The aggregated intermediate statistics from all batches. 

77 

78 Returns: 

79 A dictionary containing the final dataset-level statistics. 

80 """ 

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

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

83 return {} 

84 

85 def compute_delta(self, source: dict[str, Any], target: dict[str, Any]) -> dict[str, Any]: 

86 """ 

87 Compare metrics between two different dataselections. 

88 

89 Args: 

90 source: Final metrics from the source dataselection. 

91 target: Final metrics from the target dataselection. 

92 

93 Returns: 

94 A dictionary containing distance or difference scores. 

95 """ 

96 return {}