feat: add data quality regression for reference data

spark/jobs/metrics/data_quality_calculator.py

@@ -13,6 +13,7 @@
     Histogram,
     CategoricalFeatureMetrics,
     ClassMetrics,
+    NumericalTargetMetrics,
 )
 from utils.misc import split_dict
 from utils.models import ModelOut
@@ -383,3 +384,47 @@ def create_histogram(feature: str):
         ]
 
         return numerical_features_metrics
+
+    def regression_target_metrics(
+        target_column: str, dataframe: DataFrame, dataframe_count: int
+    ) -> NumericalTargetMetrics:
+        target_metrics = (
+            dataframe.select(target_column)
+            .filter(F.isnotnull(target_column))
+            .agg(
+                F.mean(target_column).alias("mean"),
+                F.stddev(target_column).alias("std"),
+                F.max(target_column).alias("max"),
+                F.min(target_column).alias("min"),
+                F.median(target_column).alias("median"),
+                F.percentile_approx(target_column, 0.25).alias("perc_25"),
+                F.percentile_approx(target_column, 0.75).alias("perc_75"),
+                F.count(F.when(F.col(target_column).isNull(), target_column)).alias(
+                    "missing_values"
+                ),
+                (
+                    (
+                        F.count(
+                            F.when(
+                                F.col(target_column).isNull() | F.isnan(target_column),
+                                target_column,
+                            )
+                        )
+                        / dataframe_count
+                    )
+                    * 100
+                ).alias("missing_values_perc"),
+            )
+            .toPandas()
+            .iloc[0]
+            .to_dict()
+        )
+
+        _histogram = (
+            dataframe.select(target_column).rdd.flatMap(lambda x: x).histogram(10)
+        )
+        histogram = Histogram(buckets=_histogram[0], reference_values=_histogram[1])
+
+        return NumericalTargetMetrics.from_dict(
+            target_column, target_metrics, histogram
+        )

spark/jobs/models/data_quality.py

@@ -79,6 +79,43 @@ def from_dict(
         )
 
 
+class NumericalTargetMetrics(FeatureMetrics):
+    type: str = "numerical"
+    mean: float
+    std: float
+    min: float
+    max: float
+    median_metrics: MedianMetrics
+    histogram: Histogram
+
+    model_config = ConfigDict(ser_json_inf_nan="null")
+
+    @classmethod
+    def from_dict(
+        cls,
+        feature_name: str,
+        global_dict: Dict,
+        histogram: Histogram,
+    ) -> "NumericalTargetMetrics":
+        return NumericalTargetMetrics(
+            feature_name=feature_name,
+            missing_value=MissingValue(
+                count=global_dict.get("missing_values"),
+                percentage=global_dict.get("missing_values_perc"),
+            ),
+            mean=global_dict.get("mean"),
+            std=global_dict.get("std"),
+            min=global_dict.get("min"),
+            max=global_dict.get("max"),
+            median_metrics=MedianMetrics(
+                median=global_dict.get("median"),
+                perc_25=global_dict.get("perc_25"),
+                perc_75=global_dict.get("perc_75"),
+            ),
+            histogram=histogram,
+        )
+
+
 class CategoryFrequency(BaseModel):
     name: str
     count: int
@@ -135,3 +172,9 @@ class MultiClassDataQuality(BaseModel):
     n_observations: int
     class_metrics: List[ClassMetrics]
     feature_metrics: List[FeatureMetrics]
+
+
+class RegressionDataQuality(BaseModel):
+    n_observations: int
+    target_metrics: NumericalTargetMetrics
+    feature_metrics: List[FeatureMetrics]

spark/jobs/reference_job.py

@@ -89,6 +89,7 @@ def main(
                 reference=reference_dataset
             )
             statistics = calculate_statistics_reference(reference_dataset)
+            data_quality = metrics_service.calculate_data_quality()
             model_quality = metrics_service.calculate_model_quality()
 
             complete_record["STATISTICS"] = statistics.model_dump_json(
@@ -97,6 +98,9 @@ def main(
             complete_record["MODEL_QUALITY"] = model_quality.model_dump_json(
                 serialize_as_any=True
             )
+            complete_record["DATA_QUALITY"] = data_quality.model_dump_json(
+                serialize_as_any=True
+            )
 
     schema = StructType(
         [

spark/jobs/utils/reference_regression.py

@@ -1,6 +1,14 @@
+from typing import List
 from models.regression_model_quality import ModelQualityRegression
 from models.reference_dataset import ReferenceDataset
 from metrics.model_quality_regression_calculator import ModelQualityRegressionCalculator
+from models.data_quality import (
+    CategoricalFeatureMetrics,
+    NumericalFeatureMetrics,
+    NumericalTargetMetrics,
+    RegressionDataQuality,
+)
+from metrics.data_quality_calculator import DataQualityCalculator
 
 
 class ReferenceMetricsRegressionService:
@@ -13,3 +21,37 @@ def calculate_model_quality(self) -> ModelQualityRegression:
             dataframe=self.reference.reference,
             dataframe_count=self.reference.reference_count,
         )
+
+    def calculate_data_quality_numerical(self) -> List[NumericalFeatureMetrics]:
+        return DataQualityCalculator.numerical_metrics(
+            model=self.reference.model,
+            dataframe=self.reference.reference,
+            dataframe_count=self.reference.reference_count,
+        )
+
+    def calculate_data_quality_categorical(self) -> List[CategoricalFeatureMetrics]:
+        return DataQualityCalculator.categorical_metrics(
+            model=self.reference.model,
+            dataframe=self.reference.reference,
+            dataframe_count=self.reference.reference_count,
+        )
+
+    def calculate_target_metrics(self) -> NumericalTargetMetrics:
+        return DataQualityCalculator.regression_target_metrics(
+            target_column=self.reference.model.target.name,
+            dataframe=self.reference.reference,
+            dataframe_count=self.reference.reference_count,
+        )
+
+    def calculate_data_quality(self) -> RegressionDataQuality:
+        feature_metrics = []
+        if self.reference.model.get_numerical_features():
+            feature_metrics.extend(self.calculate_data_quality_numerical())
+        if self.reference.model.get_categorical_features():
+            feature_metrics.extend(self.calculate_data_quality_categorical())
+        target_metrics = self.calculate_target_metrics()
+        return RegressionDataQuality(
+            n_observations=self.reference.reference_count,
+            target_metrics=target_metrics,
+            feature_metrics=feature_metrics,
+        )