handle and drop nan values

radicalbit · Jan 10, 2025 · 2568df7 · 2568df7
1 parent f58d9b4
commit 2568df7
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 17 deletions.
diff --git a/spark/jobs/metrics/jensen_shannon_distance.py b/spark/jobs/metrics/jensen_shannon_distance.py
@@ -15,17 +15,13 @@ def __init__(self, spark_session, reference_data, current_data) -> None:
         - spark_session (SparkSession): The SparkSession object
         - reference_data (pyspark.sql.DataFrame): The DataFrame containing the reference data
         - current_data (pyspark.sql.DataFrame): The DataFrame containing the current data
-        - reference_data_length (int): The reference length
-        - current_data_length (int): The current length
         - rbit_prefix (str): A prefix to assign to temporary fields
         - percentiles (list): The list with the percentiles to bucketize continuous variables
         - relative_error (float): The error to assign to approxQuantile
         """
         self.spark_session = spark_session
         self.reference_data = reference_data
         self.current_data = current_data
-        self.reference_data_length = self.reference_data.count()
-        self.current_data_length = self.current_data.count()
         self.rbit_prefix = "rbit_spark"
         self.percentiles = [i / 10 for i in range(1, 10)]
         self.relative_error = 0.05
@@ -44,10 +40,13 @@ def __calculate_category_percentages(
         DataFrame with two columns: category and percentage
         """
 
+        df = df.filter(f.col(column_name).isNotNull())
+        total_count = df.count()
+
         category_counts = df.groupBy(column_name).agg(
             f.count("*").alias(f"{self.rbit_prefix}_count")
         )
-        total_count = df.count()
+
         result_df = category_counts.withColumn(
             f"{self.rbit_prefix}_percentage",
             (f.col(f"{self.rbit_prefix}_count") / f.lit(total_count)),
@@ -71,8 +70,13 @@ def __bucketize_continuous_values(
         Returns:
         - A Tuple of DataFrames containing reference and current percentages
         """
-        reference = df_reference.select(column_name)
-        current = df_current.select(column_name)
+        reference = df_reference.select(column_name).filter(
+            f.col(column_name).isNotNull()
+        )
+        current = df_current.select(column_name).filter(f.col(column_name).isNotNull())
+
+        reference_count = reference.count()
+        current_count = current.count()
 
         reference_quantiles = reference.approxQuantile(
             column_name, self.percentiles, 0.01
@@ -96,10 +100,7 @@ def __bucketize_continuous_values(
         reference_bucket_percentages = (
             reference_bucket_counts.withColumn(
                 f"{self.rbit_prefix}_percentage",
-                (
-                    f.col(f"{self.rbit_prefix}_count")
-                    / f.lit(self.reference_data_length)
-                ),
+                (f.col(f"{self.rbit_prefix}_count") / f.lit(reference_count)),
             )
             .select(
                 f.col(f"{self.rbit_prefix}_bucket"),
@@ -108,8 +109,6 @@ def __bucketize_continuous_values(
             .orderBy(f"{self.rbit_prefix}_bucket")
         )
 
-        reference_bucket_percentages.show()
-
         current_with_buckets = bucketizer.setHandleInvalid("keep").transform(current)
 
         current_bucket_counts = current_with_buckets.groupBy(
@@ -119,7 +118,7 @@ def __bucketize_continuous_values(
         current_bucket_percentages = (
             current_bucket_counts.withColumn(
                 f"{self.rbit_prefix}_percentage",
-                (f.col(f"{self.rbit_prefix}_count") / f.lit(self.current_data_length)),
+                (f.col(f"{self.rbit_prefix}_count") / f.lit(current_count)),
             )
             .select(
                 f.col(f"{self.rbit_prefix}_bucket"),

diff --git a/spark/tests/jensen_shannon_distance_test.py b/spark/tests/jensen_shannon_distance_test.py
@@ -1,9 +1,7 @@
 import pytest
 from pyspark.sql import DataFrame
 from pyspark.sql import Row
-from jobs.metrics.jensen_shannon_distance import (
-    JensenShannonDistance
-)
+from jobs.metrics.jensen_shannon_distance import JensenShannonDistance
 
 
 @pytest.fixture(scope="module")
@@ -50,6 +48,7 @@ def test_return_distance_discrete(jensen_shannon_distance):
         0 <= result["JensenShannonDistance"] <= 1
     ), "Distance should be between 0 and 1."
 
+
 def test_calculate_category_percentages(jensen_shannon_distance):
     """Test the __calculate_category_percentages method."""
     percentages = (