Fix chi-square test conditions (#482)

* add population stability index function * add population stability index function * correct style issues * refactor psi formula * update imports * add population stability index function Bug fix of chi-square test condition * correct style issues --------- Co-authored-by: Bever <[email protected]>
awslabs · Apr 16, 2024 · 7800e91 · 7800e91
1 parent 06df314
commit 7800e91
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 1 deletion.
diff --git a/src/main/scala/com/amazon/deequ/analyzers/Distance.scala b/src/main/scala/com/amazon/deequ/analyzers/Distance.scala
@@ -155,7 +155,7 @@ object Distance {
       sample.map(e => (e._1, e._2.toDouble)), expectedNorm, absThresholdYates, percThresholdYates, absThresholdCochran)
 
     // If less than 2 categories remain we cannot conduct the test
-    if (regroupedSample.keySet.size < chisquareMinDimension) {
+    if (regroupedExpected.keySet.size < chisquareMinDimension) {
       Double.NaN
     } else {
       // run chi-square test and return statistics or p-value

diff --git a/src/test/scala/com/amazon/deequ/KLL/KLLDistanceTest.scala b/src/test/scala/com/amazon/deequ/KLL/KLLDistanceTest.scala
@@ -54,6 +54,7 @@ class KLLDistanceTest extends WordSpec with SparkContextSpec
     assert(distance == 0.06015037593984962)
   }
 
+
   "Categorial distance should compute correct linf_robust" in {
     val sample1 = scala.collection.mutable.Map(
       "a" -> 10L, "b" -> 20L, "c" -> 25L, "d" -> 10L, "e" -> 5L)
@@ -183,6 +184,33 @@ class KLLDistanceTest extends WordSpec with SparkContextSpec
     assert(distance.isNaN)
   }
 
+  "Categorial chi-square distance when number of expected categories below minimum" in {
+    val sample = scala.collection.mutable.Map(
+      "a" -> 10L, "b" -> 20L)
+    val expected = scala.collection.mutable.Map(
+      "b" -> 20L)
+    val distance = Distance.categoricalDistance(sample, expected, method = ChisquareMethod())
+    assert(distance.equals(Double.NaN))
+  }
+
+  "Categorial chi-square distance when categories do not match" in {
+    val sample = scala.collection.mutable.Map(
+      "a" -> 15L, "b" -> 20L)
+    val expected = scala.collection.mutable.Map(
+      "c" -> 20L, "d" -> 20L)
+    val distance = Distance.categoricalDistance(sample, expected, method = ChisquareMethod())
+    assert(distance.equals(Double.NaN))
+  }
+
+  "Categorial chi-square distance when number of sample categories is below minimum" in {
+    val sample = scala.collection.mutable.Map(
+      "a" -> 30L)
+    val expected = scala.collection.mutable.Map(
+      "a" -> 20L, "b" -> 20L)
+    val distance = Distance.categoricalDistance(sample, expected, method = ChisquareMethod())
+    assert(distance == 4.3204630539861455E-8)
+  }
+
   "Population Stability Index (PSI) test with deciles " in {
 
     val expected: List[BucketValue] = List(BucketValue(1.0, 1.05, 428), BucketValue(1.05, 1.1, 425),