feat: data quality for multiclass + fix class metrics

spark/jobs/current_job.py

@@ -6,10 +6,10 @@
 import orjson
 from pyspark.sql.types import StructType, StructField, StringType
 
-from jobs.metrics.statistics import calculate_statistics_current
-from jobs.models.current_dataset import CurrentDataset
-from jobs.models.reference_dataset import ReferenceDataset
-from utils.current import CurrentMetricsService
+from metrics.statistics import calculate_statistics_current
+from models.current_dataset import CurrentDataset
+from models.reference_dataset import ReferenceDataset
+from utils.current_binary import CurrentMetricsService
 from utils.models import JobStatus, ModelOut
 from utils.db import update_job_status, write_to_db
 

spark/jobs/metrics/data_quality_calculator.py

@@ -0,0 +1,200 @@
+from typing import List
+
+import pyspark.sql.functions as F
+from pandas import DataFrame
+
+from models.data_quality import (
+    NumericalFeatureMetrics,
+    Histogram,
+    CategoricalFeatureMetrics,
+    ClassMetrics,
+)
+from utils.misc import split_dict
+from utils.models import ModelOut
+from utils.spark import check_not_null
+
+
+class DataQualityCalculator:
+    @staticmethod
+    def numerical_metrics(
+        model: ModelOut, dataframe: DataFrame, dataframe_count: int
+    ) -> List[NumericalFeatureMetrics]:
+        numerical_features = [
+            numerical.name for numerical in model.get_numerical_features()
+        ]
+
+        mean_agg = [
+            (F.mean(check_not_null(x))).alias(f"{x}-mean") for x in numerical_features
+        ]
+
+        max_agg = [
+            (F.max(check_not_null(x))).alias(f"{x}-max") for x in numerical_features
+        ]
+
+        min_agg = [
+            (F.min(check_not_null(x))).alias(f"{x}-min") for x in numerical_features
+        ]
+
+        median_agg = [
+            (F.median(check_not_null(x))).alias(f"{x}-median")
+            for x in numerical_features
+        ]
+
+        perc_25_agg = [
+            (F.percentile(check_not_null(x), 0.25)).alias(f"{x}-perc_25")
+            for x in numerical_features
+        ]
+
+        perc_75_agg = [
+            (F.percentile(check_not_null(x), 0.75)).alias(f"{x}-perc_75")
+            for x in numerical_features
+        ]
+
+        std_agg = [
+            (F.std(check_not_null(x))).alias(f"{x}-std") for x in numerical_features
+        ]
+
+        missing_values_agg = [
+            (F.count(F.when(F.col(x).isNull() | F.isnan(x), x))).alias(
+                f"{x}-missing_values"
+            )
+            for x in numerical_features
+        ]
+
+        missing_values_perc_agg = [
+            (
+                (F.count(F.when(F.col(x).isNull() | F.isnan(x), x)) / dataframe_count)
+                * 100
+            ).alias(f"{x}-missing_values_perc")
+            for x in numerical_features
+        ]
+
+        # Global
+        global_stat = dataframe.select(numerical_features).agg(
+            *(
+                mean_agg
+                + max_agg
+                + min_agg
+                + median_agg
+                + perc_25_agg
+                + perc_75_agg
+                + std_agg
+                + missing_values_agg
+                + missing_values_perc_agg
+            )
+        )
+
+        global_dict = global_stat.toPandas().iloc[0].to_dict()
+        global_data_quality = split_dict(global_dict)
+
+        # TODO probably not so efficient but I haven't found another way
+        histograms = {
+            column: dataframe.select(column).rdd.flatMap(lambda x: x).histogram(10)
+            for column in numerical_features
+        }
+
+        dict_of_hist = {
+            k: Histogram(buckets=v[0], reference_values=v[1])
+            for k, v in histograms.items()
+        }
+
+        numerical_features_metrics = [
+            NumericalFeatureMetrics.from_dict(
+                feature_name,
+                metrics,
+                histogram=dict_of_hist.get(feature_name),
+            )
+            for feature_name, metrics in global_data_quality.items()
+        ]
+
+        return numerical_features_metrics
+
+    @staticmethod
+    def categorical_metrics(
+        model: ModelOut, dataframe: DataFrame, dataframe_count: int
+    ) -> List[CategoricalFeatureMetrics]:
+        categorical_features = [
+            categorical.name for categorical in model.get_categorical_features()
+        ]
+
+        missing_values_agg = [
+            (F.count(F.when(F.col(x).isNull(), x))).alias(f"{x}-missing_values")
+            for x in categorical_features
+        ]
+
+        missing_values_perc_agg = [
+            ((F.count(F.when(F.col(x).isNull(), x)) / dataframe_count) * 100).alias(
+                f"{x}-missing_values_perc"
+            )
+            for x in categorical_features
+        ]
+
+        distinct_values = [
+            (F.countDistinct(check_not_null(x))).alias(f"{x}-distinct_values")
+            for x in categorical_features
+        ]
+
+        global_stat = dataframe.select(categorical_features).agg(
+            *(missing_values_agg + missing_values_perc_agg + distinct_values)
+        )
+
+        global_dict = global_stat.toPandas().iloc[0].to_dict()
+        global_data_quality = split_dict(global_dict)
+
+        # FIXME by design this is not efficient
+        # FIXME understand if we want to divide by whole or by number of not null
+
+        count_distinct_categories = {
+            column: dict(
+                dataframe.select(column)
+                .filter(F.isnotnull(column))
+                .groupBy(column)
+                .agg(*[F.count(check_not_null(column)).alias("count")])
+                .withColumn(
+                    "freq",
+                    F.col("count") / dataframe_count,
+                )
+                .toPandas()
+                .set_index(column)
+                .to_dict()
+            )
+            for column in categorical_features
+        }
+
+        categorical_features_metrics = [
+            CategoricalFeatureMetrics.from_dict(
+                feature_name=feature_name,
+                global_metrics=metrics,
+                categories_metrics=count_distinct_categories.get(feature_name),
+            )
+            for feature_name, metrics in global_data_quality.items()
+        ]
+
+        return categorical_features_metrics
+
+    @staticmethod
+    def class_metrics(
+        class_column: str, dataframe: DataFrame, dataframe_count: int
+    ) -> List[ClassMetrics]:
+        class_metrics_dict = (
+            dataframe.select(class_column)
+            .filter(F.isnotnull(class_column))
+            .groupBy(class_column)
+            .agg(*[F.count(check_not_null(class_column)).alias("count")])
+            .withColumn(
+                "percentage",
+                (F.col("count") / dataframe_count) * 100,
+            )
+            .toPandas()
+            .set_index(class_column)
+            .to_dict(orient="index")
+        )
+
+        return [
+            ClassMetrics(
+                name=str(label),
+                count=metrics["count"],
+                percentage=metrics["percentage"],
+            )
+            for label, metrics in class_metrics_dict.items()
+        ]

spark/jobs/metrics/statistics.py

@@ -14,6 +14,7 @@
 
 
 # FIXME use pydantic struct like data quality
+# FIXME generalize to one method
 def calculate_statistics_reference(
     reference_dataset: ReferenceDataset,
 ) -> dict[str, float]:

spark/jobs/utils/data_quality.py → spark/jobs/models/data_quality.py

@@ -57,8 +57,6 @@ def from_dict(
         cls,
         feature_name: str,
         global_dict: Dict,
-        true_feature_dict: Dict,
-        false_feature_dict: Dict,
         histogram: Histogram,
     ) -> "NumericalFeatureMetrics":
         return NumericalFeatureMetrics(
@@ -76,26 +74,7 @@ def from_dict(
                 perc_25=global_dict.get("perc_25"),
                 perc_75=global_dict.get("perc_75"),
             ),
-            class_median_metrics=[
-                ClassMedianMetrics(
-                    name="true",
-                    mean=true_feature_dict.get("mean"),
-                    median_metrics=MedianMetrics(
-                        median=true_feature_dict.get("median"),
-                        perc_25=true_feature_dict.get("perc_25"),
-                        perc_75=true_feature_dict.get("perc_75"),
-                    ),
-                ),
-                ClassMedianMetrics(
-                    name="false",
-                    mean=false_feature_dict.get("mean"),
-                    median_metrics=MedianMetrics(
-                        median=false_feature_dict.get("median"),
-                        perc_25=false_feature_dict.get("perc_25"),
-                        perc_75=false_feature_dict.get("perc_75"),
-                    ),
-                ),
-            ],
+            class_median_metrics=[],
             histogram=histogram,
         )
 
@@ -150,3 +129,9 @@ class BinaryClassDataQuality(BaseModel):
     n_observations: int
     class_metrics: List[ClassMetrics]
     feature_metrics: List[FeatureMetrics]
+
+
+class MultiClassDataQuality(BaseModel):
+    n_observations: int
+    class_metrics: List[ClassMetrics]
+    feature_metrics: List[FeatureMetrics]

spark/jobs/models/reference_dataset.py

@@ -1,5 +1,6 @@
 from typing import List
 
+from pyspark.ml.feature import StringIndexer
 from pyspark.sql import DataFrame
 from pyspark.sql.types import DoubleType, StructField, StructType
 
@@ -90,3 +91,31 @@ def get_all_variables(self) -> List[ColumnDefinition]:
             + [self.model.timestamp]
             + self.model.outputs.output
         )
+
+    def get_string_indexed_dataframe(self):
+        """
+        Source: https://stackoverflow.com/questions/65911146/how-to-transform-multiple-categorical-columns-to-integers-maintaining-shared-val
+        """
+        predictions_df = self.reference.select(
+            self.model.outputs.prediction.name
+        ).withColumnRenamed(self.model.outputs.prediction.name, "classes")
+        target_df = self.reference.select(self.model.target.name).withColumnRenamed(
+            self.model.target.name, "classes"
+        )
+        prediction_target_df = predictions_df.union(target_df)
+        indexer = StringIndexer(inputCol="classes", outputCol="classes_index")
+        indexer_model = indexer.fit(prediction_target_df)
+        indexer_prediction = indexer_model.setInputCol(
+            self.model.outputs.prediction.name
+        ).setOutputCol(f"{self.model.outputs.prediction.name}-idx")
+        indexed_prediction_df = indexer_prediction.transform(self.reference)
+        indexer_target = indexer_model.setInputCol(self.model.target.name).setOutputCol(
+            f"{self.model.target.name}-idx"
+        )
+        indexed_target_df = indexer_target.transform(indexed_prediction_df)
+
+        index_label_map = {
+            str(float(index)): label
+            for index, label in enumerate(indexer_model.labelsArray[0])
+        }
+        return index_label_map, indexed_target_df

spark/jobs/reference_job.py

@@ -7,14 +7,16 @@
 
 from metrics.statistics import calculate_statistics_reference
 from models.reference_dataset import ReferenceDataset
-from utils.reference import ReferenceMetricsService
+from utils.reference_binary import ReferenceMetricsService
 from utils.models import JobStatus, ModelOut, ModelType
 from utils.db import update_job_status, write_to_db
 
 from pyspark.sql import SparkSession
 
 import logging
 
+from utils.reference_multiclass import ReferenceMetricsMulticlassService
+
 
 def main(
     spark_session: SparkSession,
@@ -46,12 +48,13 @@ def main(
     raw_dataframe = spark_session.read.csv(reference_dataset_path, header=True)
     reference_dataset = ReferenceDataset(model=model, raw_dataframe=raw_dataframe)
 
-    metrics_service = ReferenceMetricsService(reference_dataset.reference, model=model)
-
     complete_record = {"UUID": str(uuid.uuid4()), "REFERENCE_UUID": reference_uuid}
 
     match model.model_type:
         case ModelType.BINARY:
+            metrics_service = ReferenceMetricsService(
+                reference_dataset.reference, model=model
+            )
             model_quality = metrics_service.calculate_model_quality()
             statistics = calculate_statistics_reference(reference_dataset)
             data_quality = metrics_service.calculate_data_quality()
@@ -64,8 +67,15 @@ def main(
             )
         case ModelType.MULTI_CLASS:
             # TODO add data quality and model quality
+            metrics_service = ReferenceMetricsMulticlassService(
+                reference=reference_dataset
+            )
             statistics = calculate_statistics_reference(reference_dataset)
+            data_quality = metrics_service.calculate_data_quality()
             complete_record["STATISTICS"] = orjson.dumps(statistics).decode("utf-8")
+            complete_record["DATA_QUALITY"] = data_quality.model_dump_json(
+                serialize_as_any=True
+            )
 
     schema = StructType(
         [