Coverage for packages/dqm-ml-core/src/dqm_ml_core/metrics/completeness.py: 100%

1"""Completeness metric processor for evaluating data completeness. 

2 

3This module contains the CompletenessProcessor class that evaluates 

4the completeness of tabular data by computing non-null value ratios. 

5""" 

6 

7import json 

8import logging 

9from typing import Any 

10 

11import numpy as np 

12import pyarrow as pa 

13 

14# COMPATIBILITY : from typing import Any, override # When support of 3.10 and 3.11 will be removed 

15from typing_extensions import override 

16 

17from dqm_ml_core.api.metrics_processor import MetricsProcessor 

18from dqm_ml_core.models.processors import CompletenessProcessorConfig 

19 

20logger = logging.getLogger(__name__) 

21 

22 

23class CompletenessProcessor(MetricsProcessor): 

24 """ 

25 Data completeness processor that evaluates the completeness of tabular data. 

26 

27 This processor calculates completeness scores (ratio of non-null to 

28 total values) for specified columns and provides overall dataset 

29 completeness metrics. 

30 

31 The processor operates at multiple levels: 

32 - Batch level: Aggregated counts for streaming processing 

33 - Dataset level: Final completeness scores and statistics 

34 """ 

35 

36 def __init__(self, name: str = "completeness", config: dict[str, Any] | None = None) -> None: 

37 """ 

38 Initialize the completeness processor. 

39 

40 Args: 

41 name: Name of the processor. 

42 config: Configuration dictionary containing: 

43 - input_columns: List of columns to analyze. 

44 - output_metrics: Mapping of metric names to column names. 

45 - include_per_column: Include per-column completeness scores. 

46 - include_overall: Include overall completeness score. 

47 """ 

48 super().__init__(name, config) 

49 

50 cfg = CompletenessProcessorConfig.model_validate({**self.config, "name": self.name}) 

51 

52 # Which completeness levels to compute: per-column, overall, and metadata 

53 self.include_per_column: bool = cfg.include_per_column 

54 self.include_overall: bool = cfg.include_overall 

55 self.include_metadata: bool = cfg.include_metadata 

56 

57 # Output column mappings 

58 self.output_metrics: dict[str, str] = {} 

59 

60 @override 

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

62 """ 

63 Return the list of metric columns that will be generated. 

64 

65 Returns: 

66 List of output metric column names 

67 """ 

68 # TODO : manage output metrics names with configuration 

69 metrics = [] 

70 

71 if self.include_overall: 

72 overall_key = self.output_metrics.get("overall_completeness", "completeness_overall") 

73 metrics.append(overall_key) 

74 

75 if self.include_per_column: 

76 for col in self.input_columns: 

77 col_key = self.output_metrics.get(f"completeness_{col}", f"completeness_{col}") 

78 metrics.append(col_key) 

79 

80 return metrics 

81 

82 @override 

83 def select_columns( 

84 self, 

85 batch: pa.RecordBatch, 

86 prev_features: dict[str, pa.Array] | None = None, 

87 ) -> dict[str, pa.Array]: 

88 """ 

89 Extract the needed columns from the batch for completeness analysis. 

90 

91 This method simply passes through the columns we need to analyze, 

92 as completeness calculation is done at batch and dataset levels. 

93 

94 Args: 

95 batch: Input batch of data 

96 prev_features: Previous features (not used in this processor) 

97 

98 Returns: 

99 Dictionary containing the columns to analyze 

100 """ 

101 features = {} 

102 

103 columns_to_analyze = self.input_columns if self.input_columns else batch.column_names 

104 

105 for col in columns_to_analyze: 

106 if col not in batch.schema.names: 

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

108 continue 

109 

110 # Simply pass through the column data for batch-level processing 

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

112 

113 return features 

114 

115 @override 

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

117 """ 

118 Compute batch-level completeness counts for streaming aggregation. 

119 

120 This counts total and non-null values per column in this batch, 

121 which will be aggregated across all batches for final dataset completeness. 

122 

123 Args: 

124 features: Dictionary of column arrays from this batch 

125 

126 Returns: 

127 Dictionary of batch-level completeness counts 

128 """ 

129 batch_metrics = {} 

130 

131 for col, col_array in features.items(): 

132 total_count = len(col_array) 

133 

134 # Count complete (non-null, non-NaN) samples in this batch. 

135 # pa.compute.is_valid() returns True for non-null values, but 

136 # float NaN from numpy is preserved as a valid float in Arrow 

137 # (not as a null), so we subtract NaN positions for float cols. 

138 valid_count = pa.compute.sum(pa.compute.is_valid(col_array)).as_py() 

139 if pa.types.is_floating(col_array.type): 

140 nan_count = pa.compute.sum(pa.compute.is_nan(col_array)).as_py() 

141 complete_count = valid_count - nan_count 

142 else: 

143 complete_count = valid_count 

144 

145 # store counts for aggregation across batches 

146 batch_metrics[f"{col}_total_count"] = pa.array([total_count], type=pa.int64()) 

147 batch_metrics[f"{col}_complete_count"] = pa.array([complete_count], type=pa.int64()) 

148 

149 return batch_metrics 

150 

151 @staticmethod 

152 def _select_columns_from_metrics( 

153 batch_metrics: dict[str, pa.Array], 

154 ) -> list[str]: 

155 """Extract column names from batch metrics keys ending in _total_count. 

156 

157 Args: 

158 batch_metrics: Dictionary of batch-level metrics. 

159 

160 Returns: 

161 List of column names. 

162 """ 

163 columns = [] 

164 for key in batch_metrics: 

165 if key.endswith("_total_count"): 

166 columns.append(key.replace("_total_count", "")) 

167 return columns 

168 

169 @staticmethod 

170 def _compute_column_completeness(col: str, batch_metrics: dict[str, pa.Array]) -> float | None: 

171 """Compute completeness score for a single column. 

172 

173 Args: 

174 col: Column name. 

175 batch_metrics: Dictionary of batch-level metrics. 

176 

177 Returns: 

178 Completeness score (0.0-1.0) or None if metrics are missing. 

179 """ 

180 total_key = f"{col}_total_count" 

181 complete_key = f"{col}_complete_count" 

182 if total_key not in batch_metrics or complete_key not in batch_metrics: 

183 return None 

184 

185 col_total = int(np.sum(batch_metrics[total_key].to_numpy())) 

186 col_complete = int(np.sum(batch_metrics[complete_key].to_numpy())) 

187 return col_complete / col_total if col_total > 0 else 0.0 

188 

189 def _write_output_metrics( 

190 self, 

191 results: dict[str, Any], 

192 per_column_completeness: dict[str, float], 

193 ) -> None: 

194 """Write per-column and overall completeness metrics to results. 

195 

196 Args: 

197 results: Output dict to populate. 

198 per_column_completeness: Dict of column -> completeness score. 

199 """ 

200 if self.include_per_column: 

201 for col, score in per_column_completeness.items(): 

202 output_key = self.output_metrics.get(f"completeness_{col}", f"completeness_{col}") 

203 results[output_key] = score 

204 

205 if self.include_overall: 

206 overall = ( 

207 sum(per_column_completeness.values()) / len(per_column_completeness) if per_column_completeness else 0.0 

208 ) 

209 output_key = self.output_metrics.get("overall_completeness", "completeness_overall") 

210 results[output_key] = overall 

211 

212 @override 

213 def compute(self, batch_metrics: dict[str, pa.Array] | None = None) -> dict[str, Any]: 

214 """ 

215 Compute final dataset-level completeness metrics. 

216 

217 This aggregates the batch-level counts to compute final completeness scores 

218 for each column and overall dataset completeness. 

219 

220 Args: 

221 batch_metrics: Dictionary of batch-level metrics to aggregate 

222 

223 Returns: 

224 Dictionary of final completeness metrics 

225 """ 

226 if not batch_metrics: 

227 return {"_metadata": {"error": "No batch metrics provided"}} 

228 

229 columns_analyzed = self._select_columns_from_metrics(batch_metrics) 

230 if not columns_analyzed: 

231 logger.warning(f"[{self.name}] No columns found in batch metrics") 

232 return {"_metadata": {"error": "No columns found in batch metrics"}} 

233 

234 per_column_completeness: dict[str, float] = {} 

235 total_samples = 0 

236 

237 for col in columns_analyzed: 

238 score = self._compute_column_completeness(col, batch_metrics) 

239 if score is None: 

240 logger.warning(f"[{self.name}] Missing batch metrics for column '{col}'") 

241 continue 

242 per_column_completeness[col] = score 

243 total_key = f"{col}_total_count" 

244 total_samples += int(np.sum(batch_metrics[total_key].to_numpy())) 

245 

246 results: dict[str, Any] = {} 

247 self._write_output_metrics(results, per_column_completeness) 

248 

249 if self.include_metadata: 

250 metadata = { 

251 "columns_analyzed": columns_analyzed, 

252 "total_samples_per_column": total_samples // len(columns_analyzed) if columns_analyzed else 0, 

253 "per_column_scores": per_column_completeness, 

254 "overall_score": sum(per_column_completeness.values()) / len(per_column_completeness) 

255 if per_column_completeness 

256 else 0.0, 

257 } 

258 results["_metadata"] = json.dumps(metadata) 

259 

260 return results 

261 

262 @override 

263 def reset(self) -> None: 

264 """Reset processor state for new processing run. 

265 

266 The completeness processor has no persistent state across runs, 

267 so this is a no-op. Provided for interface compliance. 

268 """ 

269 # No persistent state to reset for completeness processor