fix(core): include reference-only distribution categories

core/src/main/scala/com/microsoft/azure/synapse/ml/exploratory/DistributionBalanceMeasure.scala

@@ -20,7 +20,7 @@ import scala.collection.JavaConverters._
 import scala.language.postfixOps
 
 /** This transformer computes data balance measures based on a reference distribution.
-  * For now, we only support a uniform reference distribution.
+  * A uniform reference distribution is used by default, and custom reference distributions are supported.
   *
   * The output is a dataframe that contains two columns:
   *   - The sensitive feature name.
@@ -62,7 +62,8 @@ class DistributionBalanceMeasure(override val uid: String)
   val referenceDistribution = new ArrayMapParam(
     this,
     "referenceDistribution",
-    "An ordered list of reference distributions that correspond to each of the sensitive columns."
+    "An ordered list of reference distributions that correspond to each of the sensitive columns. " +
+      "Each non-empty distribution must sum to 1. Omitted observed categories have reference probability 0."
   )
 
   val emptyReferenceDistribution: Array[Map[String, Double]] = Array.empty
@@ -85,22 +86,89 @@ class DistributionBalanceMeasure(override val uid: String)
     outputCol -> "DistributionBalanceMeasure"
   )
 
-  private val uniformDistribution: Int => String => Double = {
-    n: Int => {
-      _: String =>
-        1d / n
+  private val referenceDistributionTolerance = 1e-8
+
+  private def parseReferenceKey(category: String, dataType: DataType, sensitiveCol: String): Any = {
+    if (category == null) {
+      throw new IllegalArgumentException(
+        s"Reference distribution keys for sensitive column '$sensitiveCol' cannot be null.")
+    }
+
+    try {
+      dataType match {
+        case ByteType => java.lang.Byte.valueOf(category)
+        case ShortType => java.lang.Short.valueOf(category)
+        case IntegerType => java.lang.Integer.valueOf(category)
+        case LongType => java.lang.Long.valueOf(category)
+        case StringType => category
+        case _ => throw new IllegalArgumentException(
+          s"Unsupported sensitive column type ${dataType.simpleString} for '$sensitiveCol'.")
+      }
+    } catch {
+      case _: NumberFormatException => throw new IllegalArgumentException(
+        s"Reference distribution key '$category' cannot be converted to ${dataType.simpleString} " +
+          s"for sensitive column '$sensitiveCol'.")
     }
   }
 
-  private val customDistribution: Map[String, Double] => String => Double = {
-    dist: Map[String, Double] => {
-      // NOTE: If the custom distribution doesn't have the col value, return a default probability of 0
-      // This assumes that the reference distribution does not contain the col value at all
-      s: String =>
-        dist.getOrElse(s, 0d)
+  private def validateCustomReferenceDistribution(sensitiveCol: String,
+                                                  dataType: DataType,
+                                                  distribution: Map[String, Double]): Unit = {
+    distribution.foreach { case (category, probability) =>
+      if (!java.lang.Double.isFinite(probability) || probability < 0d || probability > 1d) {
+        throw new IllegalArgumentException(
+          s"Reference probability for category '$category' in sensitive column '$sensitiveCol' " +
+            s"must be finite and between 0 and 1, but found $probability.")
+      }
+    }
+
+    val probabilitySum = distribution.values.sum
+    if (math.abs(probabilitySum - 1d) > referenceDistributionTolerance) {
+      throw new IllegalArgumentException(
+        s"Reference distribution for sensitive column '$sensitiveCol' must sum to 1, but found $probabilitySum.")
+    }
+
+    val typedCategories = distribution.keys.toSeq.map(parseReferenceKey(_, dataType, sensitiveCol))
+    if (typedCategories.distinct.length != typedCategories.length) {
+      throw new IllegalArgumentException(
+        s"Reference distribution keys for sensitive column '$sensitiveCol' must identify distinct " +
+          s"${dataType.simpleString} categories.")
     }
   }
 
+  private def createReferenceSupport(observed: DataFrame,
+                                     sensitiveCol: String,
+                                     obsFeatureProbCol: String,
+                                     obsFeatureCountCol: String,
+                                     refFeatureProbCol: String,
+                                     refFeatureCountCol: String,
+                                     numRows: Double,
+                                     distribution: Map[String, Double]): DataFrame = {
+    val sensitiveType = observed.schema(sensitiveCol).dataType
+    val referenceRows = distribution.toSeq.collect {
+      case (category, probability) if probability > 0d =>
+        Row(parseReferenceKey(category, sensitiveType, sensitiveCol), probability)
+    }
+    val referenceSchema = StructType(Seq(
+      StructField(sensitiveCol, sensitiveType, nullable = false),
+      StructField(refFeatureProbCol, DoubleType, nullable = false)
+    ))
+    val reference = observed.sparkSession.createDataFrame(referenceRows.asJava, referenceSchema)
+      .withColumn(obsFeatureProbCol, lit(0d))
+      .withColumn(obsFeatureCountCol, lit(0d))
+
+    observed
+      .withColumn(refFeatureProbCol, lit(0d))
+      .unionByName(reference)
+      .groupBy(col(sensitiveCol))
+      .agg(
+        sum(obsFeatureProbCol).alias(obsFeatureProbCol),
+        sum(obsFeatureCountCol).alias(obsFeatureCountCol),
+        sum(refFeatureProbCol).alias(refFeatureProbCol)
+      )
+      .withColumn(refFeatureCountCol, col(refFeatureProbCol) * lit(numRows))
+  }
+
   override def transform(dataset: Dataset[_]): DataFrame = {
     logTransform[DataFrame]({
       validateSchema(dataset.schema)
@@ -137,19 +205,31 @@ class DistributionBalanceMeasure(override val uid: String)
           .groupBy(sensitiveCol)
           .agg(sum(obsFeatureProbCol).alias(obsFeatureProbCol), sum(obsFeatureCountCol).alias(obsFeatureCountCol))
 
-        val numFeatures = observed.count.toInt
         val refFeatureProbCol = DatasetExtensions.findUnusedColumnName("refFeatureProb", featureStats.schema)
         val refFeatureCountCol = DatasetExtensions.findUnusedColumnName("refFeatureCount", featureStats.schema)
-
-        val refDist: String => Double =
-          if (!isDefined(referenceDistribution) || getReferenceDistribution(i).isEmpty) uniformDistribution(numFeatures)
-          else customDistribution(getReferenceDistribution(i))
-        val refDistFunc = udf(refDist)
-
-        val observedWithRef = observed
-          .withColumn(refFeatureProbCol, refDistFunc(col(sensitiveCol)))
-          .withColumn(refFeatureCountCol, refDistFunc(col(sensitiveCol)) * lit(numRows))
-          .cache
+        val reference =
+          if (isDefined(referenceDistribution)) getReferenceDistribution(i) else Map.empty[String, Double]
+
+        val (observedWithRef, numFeatures) = if (reference.isEmpty) {
+          val observedFeatureCount = observed.count.toInt
+          val uniformProbability = 1d / observedFeatureCount
+          (observed
+            .withColumn(refFeatureProbCol, lit(uniformProbability))
+            .withColumn(refFeatureCountCol, lit(uniformProbability * numRows))
+            .cache, observedFeatureCount)
+        } else {
+          val support = createReferenceSupport(
+            observed,
+            sensitiveCol,
+            obsFeatureProbCol,
+            obsFeatureCountCol,
+            refFeatureProbCol,
+            refFeatureCountCol,
+            numRows,
+            reference
+          ).cache
+          (support, support.count.toInt)
+        }
 
         val metrics =
           DistributionMetrics(numFeatures, obsFeatureProbCol, obsFeatureCountCol, refFeatureProbCol, refFeatureCountCol)
@@ -183,9 +263,18 @@ class DistributionBalanceMeasure(override val uid: String)
   override def validateSchema(schema: StructType): Unit = {
     super.validateSchema(schema)
 
-    if (isDefined(referenceDistribution) && getReferenceDistribution.length != getSensitiveCols.length) {
-      throw new Exception("The reference distribution must have the same length and order as the sensitive columns: "
-        + getSensitiveCols.mkString(", "))
+    if (isDefined(referenceDistribution)) {
+      val distributions = getReferenceDistribution
+      if (distributions.length != getSensitiveCols.length) {
+        throw new Exception("The reference distribution must have the same length and order as the sensitive columns: "
+          + getSensitiveCols.mkString(", "))
+      }
+
+      getSensitiveCols.zip(distributions).foreach { case (sensitiveCol, distribution) =>
+        if (distribution.nonEmpty) {
+          validateCustomReferenceDistribution(sensitiveCol, schema(sensitiveCol).dataType, distribution)
+        }
+      }
     }
   }
 }

core/src/test/scala/com/microsoft/azure/synapse/ml/exploratory/DistributionBalanceMeasureSuite.scala

@@ -8,6 +8,8 @@ import org.apache.spark.ml.util.MLReadable
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions.{array, col}
 
+import scala.collection.immutable.ListMap
+
 class DistributionBalanceMeasureSuite extends DataBalanceTestBase with TransformerFuzzing[DistributionBalanceMeasure] {
 
   override def testObjects(): Seq[TestObject[DistributionBalanceMeasure]] = Seq(
@@ -24,6 +26,38 @@ class DistributionBalanceMeasureSuite extends DataBalanceTestBase with Transform
       .setSensitiveCols(features)
       .setVerbose(true)
 
+  private def metricsFor(input: DataFrame,
+                         sensitiveCol: String,
+                         reference: Map[String, Double]): Map[String, Double] =
+    (METRICS zip new DistributionBalanceMeasure()
+      .setSensitiveCols(Array(sensitiveCol))
+      .setReferenceDistribution(Array(reference))
+      .transform(input)
+      .filter(col("FeatureName") === sensitiveCol)
+      .select(array(col("DistributionBalanceMeasure.*")))
+      .as[Array[Double]]
+      .head).toMap
+
+  private def assertMetric(actual: Double, expected: Double): Unit = {
+    if (java.lang.Double.isNaN(expected)) {
+      assert(java.lang.Double.isNaN(actual))
+    } else if (java.lang.Double.isInfinite(expected)) {
+      assert(actual === expected)
+    } else {
+      assert(math.abs(actual - expected) < errorTolerance)
+    }
+  }
+
+  private def assertMetrics(actual: Map[String, Double], expected: DistributionMetricsCalculator): Unit = {
+    assertMetric(actual(KLDIVERGENCE), expected.klDivergence)
+    assertMetric(actual(JSDISTANCE), expected.jsDistance)
+    assertMetric(actual(INFNORMDISTANCE), expected.infNormDistance)
+    assertMetric(actual(TOTALVARIATIONDISTANCE), expected.totalVariationDistance)
+    assertMetric(actual(WASSERSTEINDISTANCE), expected.wassersteinDistance)
+    assertMetric(actual(CHISQUAREDTESTSTATISTIC), expected.chiSquaredTestStatistic)
+    assertMetric(actual(CHISQUAREDPVALUE), expected.chiSquaredPValue)
+  }
+
   test("DistributionBalanceMeasure can calculate Distribution Balance Measures end-to-end") {
     val df = distributionBalanceMeasure.transform(sensitiveFeaturesDf)
     df.show(truncate = false)
@@ -157,6 +191,138 @@ class DistributionBalanceMeasureSuite extends DataBalanceTestBase with Transform
     }
   }
 
+  test("DistributionBalanceMeasure includes reference-only categories in every metric") {
+    val source = Seq("red", "red", "red", "green", "blue").toDF("color")
+    val reference = Map("red" -> 0.4, "green" -> 0.2, "blue" -> 0.2, "yellow" -> 0.2)
+    val actual = metricsFor(source, "color", reference)
+    val expected = DistributionMetricsCalculator(
+      refFeatureProbabilities = Array(0.4, 0.2, 0.2, 0.2),
+      refFeatureCounts = Array(2d, 1d, 1d, 1d),
+      obsFeatureProbabilities = Array(0.6, 0.2, 0.2, 0d),
+      obsFeatureCounts = Array(3d, 1d, 1d, 0d),
+      numFeatures = 4d
+    )
+
+    assertMetrics(actual, expected)
+    assertMetric(actual(JSDISTANCE), 0.28174895710781067)
+  }
+
+  test("DistributionBalanceMeasure preserves results when every reference category is observed") {
+    val source = Seq("red", "red", "red", "green", "blue").toDF("color")
+    val reference = Map("red" -> 0.4, "green" -> 0.2, "blue" -> 0.4)
+
+    assertMetric(metricsFor(source, "color", reference)(JSDISTANCE), 0.1644921288538882)
+  }
+
+  test("DistributionBalanceMeasure includes observed-only categories with zero reference probability") {
+    val source = Seq("red", "red", "red", "green", "blue").toDF("color")
+    val actual = metricsFor(source, "color", Map("red" -> 0.75, "green" -> 0.25))
+    val expected = DistributionMetricsCalculator(
+      refFeatureProbabilities = Array(0.75, 0.25, 0d),
+      refFeatureCounts = Array(3.75, 1.25, 0d),
+      obsFeatureProbabilities = Array(0.6, 0.2, 0.2),
+      obsFeatureCounts = Array(3d, 1d, 1d),
+      numFeatures = 3d
+    )
+
+    assertMetrics(actual, expected)
+  }
+
+  test("DistributionBalanceMeasure aligns unordered reference keys with string and integral categories") {
+    val source = Seq(
+      (1L, "red"),
+      (1L, "red"),
+      (2L, "blue")
+    ).toDF("code", "color")
+    val references: Array[Map[String, Double]] = Array(
+      ListMap("3" -> 0.5, "1" -> 0.5),
+      ListMap("green" -> 0.5, "red" -> 0.5)
+    )
+    val reversedReferences: Array[Map[String, Double]] = Array(
+      ListMap("1" -> 0.5, "3" -> 0.5),
+      ListMap("red" -> 0.5, "green" -> 0.5)
+    )
+    val expected = DistributionMetricsCalculator(
+      refFeatureProbabilities = Array(0.5, 0d, 0.5),
+      refFeatureCounts = Array(1.5, 0d, 1.5),
+      obsFeatureProbabilities = Array(2d / 3d, 1d / 3d, 0d),
+      obsFeatureCounts = Array(2d, 1d, 0d),
+      numFeatures = 3d
+    )
+
+    Seq(references, reversedReferences).foreach { reference =>
+      val result = new DistributionBalanceMeasure()
+        .setSensitiveCols(Array("code", "color"))
+        .setReferenceDistribution(reference)
+        .transform(source)
+
+      assertMetrics(
+        (METRICS zip result.filter(col("FeatureName") === "code")
+          .select(array(col("DistributionBalanceMeasure.*")))
+          .as[Array[Double]].head).toMap,
+        expected
+      )
+      assertMetrics(
+        (METRICS zip result.filter(col("FeatureName") === "color")
+          .select(array(col("DistributionBalanceMeasure.*")))
+          .as[Array[Double]].head).toMap,
+        expected
+      )
+    }
+  }
+
+  test("DistributionBalanceMeasure accepts explicit zero probabilities without extending support") {
+    val source = Seq("red", "red", "blue").toDF("color")
+    val withoutZeroCategory = metricsFor(source, "color", Map("red" -> 1d))
+    val withZeroCategory = metricsFor(source, "color", Map("red" -> 1d, "unused" -> 0d))
+
+    METRICS.foreach(metric => assertMetric(withZeroCategory(metric), withoutZeroCategory(metric)))
+    assert(withZeroCategory(KLDIVERGENCE) === Double.PositiveInfinity)
+    assert(withZeroCategory(CHISQUAREDTESTSTATISTIC) === Double.PositiveInfinity)
+  }
+
+  test("DistributionBalanceMeasure validates custom reference probabilities") {
+    val source = Seq("red", "blue").toDF("color")
+    val sumError = intercept[IllegalArgumentException] {
+      metricsFor(source, "color", Map("red" -> 0.4, "blue" -> 0.4))
+    }
+    assert(sumError.getMessage.contains("must sum to 1"))
+
+    val probabilityError = intercept[IllegalArgumentException] {
+      metricsFor(source, "color", Map("red" -> 1.1, "blue" -> -0.1))
+    }
+    assert(probabilityError.getMessage.contains("must be finite and between 0 and 1"))
+  }
+
+  test("DistributionBalanceMeasure validates reference keys against integral column types") {
+    val source = Seq(1, 2).toDF("code")
+    val conversionError = intercept[IllegalArgumentException] {
+      metricsFor(source, "code", Map("not-an-integer" -> 1d))
+    }
+    assert(conversionError.getMessage.contains("cannot be converted to int"))
+
+    val duplicateError = intercept[IllegalArgumentException] {
+      metricsFor(source, "code", Map("1" -> 0.5, "01" -> 0.5))
+    }
+    assert(duplicateError.getMessage.contains("must identify distinct int categories"))
+  }
+
+  test("DistributionBalanceMeasure persists custom fractional distributions") {
+    val source = Seq("red", "red", "blue").toDF("color")
+    val original = new DistributionBalanceMeasure()
+      .setSensitiveCols(Array("color"))
+      .setReferenceDistribution(Array(Map("red" -> 0.6, "blue" -> 0.4)))
+    val path = tmpDir.resolve("distribution-balance-measure").toString
+
+    original.write.overwrite().save(path)
+    val loaded = DistributionBalanceMeasure.load(path)
+
+    assert(loaded.getReferenceDistribution.sameElements(original.getReferenceDistribution))
+    val expected = metricsFor(source, "color", original.getReferenceDistribution.head)
+    val actual = metricsFor(source, "color", loaded.getReferenceDistribution.head)
+    METRICS.foreach(metric => assertMetric(actual(metric), expected(metric)))
+  }
+
   private def actualCustomDist: DataFrame =
     new DistributionBalanceMeasure()
       .setSensitiveCols(features)

docs/Explore Algorithms/Responsible AI/Data Balance Analysis.md

@@ -134,7 +134,7 @@ Note: Many of these metrics were influenced by this paper [Measuring Model Biase
 
 ### Distribution Balance Measures
 
-Distribution Balance Measures allow us to compare our data with a reference distribution (currently only uniform distribution is supported as a reference distribution). They are calculated per sensitive column and don't depend on the label column.
+Distribution Balance Measures allow us to compare our data with a reference distribution. They use a uniform reference distribution by default, or a custom distribution can be supplied with `setReferenceDistribution`. Custom probabilities must sum to 1. The measures use the union of the observed and positive-probability reference categories: an observed-only category has reference probability 0, and a reference-only category has observed probability and count 0. They are calculated per sensitive column and don't depend on the label column.
 
 For example, let's assume we have a dataset with nine rows and a Gender column, and we observe that: