From 16427071fa53ff36b4791e3db1a49562b789b10d Mon Sep 17 00:00:00 2001 From: SynapseML CI Date: Fri, 14 Aug 2026 13:54:52 -0700 Subject: [PATCH] fix(core): include reference-only distribution categories ## Summary Fix DistributionBalanceMeasure so custom comparisons cover the union of observed and positive-probability reference categories, with zero observed values for reference-only categories. Add validation, persistence coverage, focused regressions, and a semantics clarification. ## Prompting Intent Own GitHub issue #2010 end-to-end with a focused, scalable DataFrame fix. Preserve legacy outputs; cover reference-only and observed-only categories, ordering and supported types, zeros, invalid inputs, persistence, and generated APIs; validate only through the JDK 11 wrapper and the shared SBT mutex. ## Linked Sources - GitHub issue: https://github.com/microsoft/SynapseML/issues/2010 - Data balance documentation: docs/Explore Algorithms/Responsible AI/Data Balance Analysis.md ## Rationale Materialize only the already-driver-resident reference parameter as a typed DataFrame, then union and aggregate it with distributed observed statistics. This aligns categories without collecting observed support or using RDD APIs. Positive-probability keys extend support, explicit zero entries behave like omitted keys, and typed-key validation prevents ambiguous numeric categories while preserving the public API and output schema. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../DistributionBalanceMeasure.scala | 141 ++++++++++++--- .../DistributionBalanceMeasureSuite.scala | 166 ++++++++++++++++++ .../Responsible AI/Data Balance Analysis.md | 2 +- 3 files changed, 282 insertions(+), 27 deletions(-) diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/exploratory/DistributionBalanceMeasure.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/exploratory/DistributionBalanceMeasure.scala index 190f7e8ca2d..faa3e66ad1f 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/exploratory/DistributionBalanceMeasure.scala +++ b/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) + } + } } } } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/exploratory/DistributionBalanceMeasureSuite.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/exploratory/DistributionBalanceMeasureSuite.scala index bd8ff485029..0c614091e7a 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/exploratory/DistributionBalanceMeasureSuite.scala +++ b/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) diff --git a/docs/Explore Algorithms/Responsible AI/Data Balance Analysis.md b/docs/Explore Algorithms/Responsible AI/Data Balance Analysis.md index c8437cc65b7..0e933812470 100644 --- a/docs/Explore Algorithms/Responsible AI/Data Balance Analysis.md +++ b/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: