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

1"""Features processor base class. 

2 

3This module contains the FeaturesProcessor base class that all 

4feature extraction processors must inherit from. 

5""" 

6 

7import pyarrow as pa 

8 

9from dqm_ml_core.api.processor import Processor 

10from dqm_ml_core.utils.matching import resolve_include_exclude 

11 

12 

13class FeaturesProcessor(Processor): 

14 """ 

15 Base class for all feature extraction processors. 

16 

17 Feature processors transform raw data into per-sample features 

18 (e.g., image luminosity, embeddings). The primary lifecycle method 

19 is ``compute_features``, which produces a dict of feature arrays 

20 from each batch of raw data. 

21 """ 

22 

23 def _resolve_output_name(self, col: str, feature_key: str) -> str: 

24 """Build the output column name with prefix and suffix applied. 

25 

26 Args: 

27 col: The input column name. 

28 feature_key: The feature key to append after ``col_``. 

29 

30 Returns: 

31 The fully qualified output column name. 

32 """ 

33 from dqm_ml_core.models.columns import ColumnsConfig 

34 

35 cfg: ColumnsConfig | None = getattr(self, "columns_config", None) 

36 base = f"{col}_{feature_key}" 

37 p = cfg.prefix if cfg else "" 

38 s = cfg.suffix if cfg else "" 

39 return f"{p}{base}{s}" 

40 

41 def _check_image_fail_fast(self, exc: Exception, *error_attrs: str) -> None: 

42 if not (self.errors_config and self.errors_config.images): 

43 return 

44 image_errors = self.errors_config.images 

45 for attr in error_attrs: 

46 if getattr(image_errors, attr, None) == "fail_fast": 

47 raise exc 

48 

49 def generated_features(self) -> list[str]: 

50 """ 

51 Return the list of columns generated by this processor during feature extraction. 

52 

53 Returns: 

54 A list of feature names. 

55 """ 

56 outputs = getattr(self, "output_features", {}) 

57 return list(outputs.values()) 

58 

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

60 """ 

61 Transform a raw data batch into features. 

62 

63 Args: 

64 batch: The input pyarrow RecordBatch. 

65 prev_features: Features already computed by preceding processors. 

66 

67 Returns: 

68 A dictionary mapping feature names to pyarrow Arrays. 

69 """ 

70 features = {} 

71 

72 available = batch.schema.names 

73 cols = resolve_include_exclude( 

74 self.input_columns, 

75 self.exclude_columns, 

76 available, 

77 ) 

78 

79 for col in cols: 

80 if col in prev_features: 

81 continue 

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

83 

84 return features