diff --git a/docs/analyzers.md b/docs/analyzers.md index 6d8244e..964f980 100644 --- a/docs/analyzers.md +++ b/docs/analyzers.md @@ -35,6 +35,39 @@ Here are the current supported functionalities of Analyzers. | UniqueValueRatio | UniqueValueRatio(columns) | Done| | AnalyzerContext | successMetricsAsDataFrame(spark_session, analyzerContext) | Done | | | successMetricsAsJson(spark_session, analyzerContext) | Done | +| Distance | Distance(spark_session).categoricalDistance(distribution1, distribution2, method) | Done | +## Distance (categorical feature drift) + +`Distance` wraps Deequ's `com.amazon.deequ.analyzers.Distance` object to compute +the distance (feature drift) between two categorical distributions, using either +the L-infinity or chi-squared method. + +Unlike the other analyzers, Deequ's `Distance` is a static object rather than an +`Analyzer`, so it is *not* added through `addAnalyzer(...)`. Instead, call it +directly with two distributions, each a `dict` of `{category: count}` (for +example, the output of the `Histogram` analyzer). + +```python +from pydeequ.analyzers import Distance, CategoricalDistanceMethod + +distance = Distance(spark) + +# Two distributions as {category: count} +reference = {"a": 10, "b": 20, "c": 30} +current = {"a": 11, "b": 20, "c": 29} + +# L-infinity distance (default) +linf = distance.categoricalDistance(reference, current) + +# Chi-squared distance +chi = distance.categoricalDistance( + reference, current, method=CategoricalDistanceMethod.Chisquare +) +``` + +Only the categorical path is supported. The numerical path +(`numericalDistance`) requires a JVM `QuantileNonSample[Double]` and is out of +scope (see issue #164). \ No newline at end of file diff --git a/pydeequ/analyzers.py b/pydeequ/analyzers.py index 3952c93..8d26a08 100644 --- a/pydeequ/analyzers.py +++ b/pydeequ/analyzers.py @@ -829,6 +829,165 @@ def _analyzer_jvm(self): self._jvm.scala.Option.apply(None) ) +class CategoricalDistanceMethod(Enum): + """ + Enum of the categorical distance methods supported by Deequ's + ``com.amazon.deequ.analyzers.Distance.categoricalDistance``. + + - ``LInfinity``: the L-infinity distance between the two distributions + (the maximum absolute difference between the per-category relative + frequencies). Optionally a Kolmogorov-Smirnov ``alpha`` can be supplied + to scale the result by the critical value. + - ``Chisquare``: the chi-squared distance between the two distributions, + with the standard Yates/Cochran corrections for low sample counts. + """ + + LInfinity = "LInfinity" + Chisquare = "Chisquare" + + +class Distance: + """ + Computes the distance (feature drift) between two categorical + distributions, mirroring Deequ's ``com.amazon.deequ.analyzers.Distance`` + object. + + In Deequ, ``Distance`` is a plain object exposing static-style methods + rather than an ``Analyzer`` subclass, so it is not added through + ``AnalysisRunBuilder.addAnalyzer(...)``. This class is therefore a thin, + faithful Python wrapper that bridges the two histograms to the JVM and + returns the numeric distance. + + The two input distributions are absolute category counts, e.g. as produced + by the :class:`Histogram` analyzer. Each is a ``dict`` mapping the category + value (``str``) to its count (``int``). + + Only the categorical path is wrapped. The numerical path + (``Distance.numericalDistance``) requires a JVM + ``QuantileNonSample[Double]`` instance which has no convenient Python + construction path and is intentionally left out of scope (see issue #164). + + :param SparkSession spark_session: SparkSession used to reach the JVM. + """ + + def __init__(self, spark_session: SparkSession): + self._spark_session = spark_session + self._jvm = spark_session._jvm + self._gateway = spark_session.sparkContext._gateway + + def _to_scala_mutable_long_map(self, distribution: dict): + """ + Build a ``scala.collection.mutable.Map[String, Long]`` from a Python + dict of ``{str: int}``, as required by ``Distance.categoricalDistance``. + + py4j auto-unboxes any individual ``java.lang.Long`` it returns to (or + receives from) Python into a Python ``int``, which then re-enters the + JVM as an ``Integer``. Building the map value-by-value therefore yields + an ``Integer``-typed map, and Deequ's ``e._2.toDouble`` throws a + ``ClassCastException``. To keep the values genuinely typed as ``Long`` + without firing a Spark job, we assign the counts into a JVM + ``java.lang.Long[]`` array (array element slots preserve the ``Long`` + boxing JVM-side), wrap both the key and value arrays as Scala + sequences, ``zip`` them, and materialize the result as a + ``mutable.HashMap[String, Long]``. No element is ever read back into + Python, so the ``Long`` typing survives end to end. + + These calls use only core Scala 2.12 stdlib APIs + (``Predef.genericWrapArray``, ``Seq.canBuildFrom``, ``Seq.zip``, + ``Seq.toMap``, ``mutable.HashMap``), which are present and identical + across every Spark/Deequ build PyDeequ supports (3.1-3.5, all Scala + 2.12). We do not rely on any ambient Java->Scala conversion implicits. + """ + items = list(distribution.items()) + size = len(items) + + keys = self._gateway.new_array(self._jvm.java.lang.String, size) + values = self._gateway.new_array(self._jvm.java.lang.Long, size) + for index, (key, count) in enumerate(items): + keys[index] = str(key) + # Assigning a Python int into a java.lang.Long[] slot stores a + # genuine java.lang.Long JVM-side (verified on Deequ 2.0.8). + values[index] = int(count) + + keys_seq = self._jvm.scala.Predef.genericWrapArray(keys) + values_seq = self._jvm.scala.Predef.genericWrapArray(values) + can_build_from = self._jvm.scala.collection.Seq.canBuildFrom() + zipped = keys_seq.zip(values_seq, can_build_from) + conforms = getattr(self._jvm.scala.Predef, "$conforms")() + immutable_map = zipped.toMap(conforms) + + # Copy the immutable Scala Map[String, Long] into a mutable.HashMap, + # which is the exact type categoricalDistance expects. + empty_mutable = self._jvm.scala.collection.mutable.HashMap() + return getattr(empty_mutable, "$plus$plus$eq")(immutable_map) + + def categoricalDistance( + self, + distribution1: dict, + distribution2: dict, + correctForLowNumberOfSamples: bool = False, + method: CategoricalDistanceMethod = CategoricalDistanceMethod.LInfinity, + alpha: float = None, + absThresholdYates: int = 5, + percThresholdYates: float = 0.2, + absThresholdCochran: int = 10, + ) -> float: + """ + Computes the categorical distance between two distributions. + + :param dict distribution1: First distribution as ``{category: count}``, + e.g. the histogram of a column on a reference dataset. + :param dict distribution2: Second distribution as ``{category: count}``, + e.g. the histogram of the same column on a new dataset. + :param bool correctForLowNumberOfSamples: If True, returns the raw + statistic (the unscaled L-infinity distance, or the chi-squared + statistic) instead of the normalized result (the + Kolmogorov-Smirnov-corrected L-infinity distance, or the + chi-squared p-value). For small samples the normalized result may + be 0.0, so set this to True when the sample count is low. Defaults + to False. + :param CategoricalDistanceMethod method: Distance method to use, + ``LInfinity`` (default) or ``Chisquare``. + :param float alpha: Only used for ``LInfinity``. Optional + Kolmogorov-Smirnov alpha used to scale the distance by the critical + value. Ignored for ``Chisquare``. + :param int absThresholdYates: Only used for ``Chisquare``. Absolute + threshold for the Yates correction. Defaults to 5. + :param float percThresholdYates: Only used for ``Chisquare``. + Percentage threshold for the Yates correction. Defaults to 0.2. + :param int absThresholdCochran: Only used for ``Chisquare``. Absolute + threshold for the Cochran correction. Defaults to 10. + :return float: The computed distance between the two distributions. + :raises ValueError: If either distribution is empty. + """ + if not distribution1 or not distribution2: + raise ValueError( + "Both distribution1 and distribution2 must be non-empty " + "dicts of {category: count}." + ) + + sample1 = self._to_scala_mutable_long_map(distribution1) + sample2 = self._to_scala_mutable_long_map(distribution2) + + # LInfinityMethod and ChisquareMethod are case classes nested inside the + # Deequ ``Distance`` object, so they are reached via Distance.. + _distance = self._jvm.com.amazon.deequ.analyzers.Distance + if method == CategoricalDistanceMethod.LInfinity: + jvm_method = _distance.LInfinityMethod(self._jvm.scala.Option.apply(alpha)) + elif method == CategoricalDistanceMethod.Chisquare: + jvm_method = _distance.ChisquareMethod( + int(absThresholdYates), + float(percThresholdYates), + int(absThresholdCochran), + ) + else: + raise ValueError(f"{method} is not a valid CategoricalDistanceMethod") + + return _distance.categoricalDistance( + sample1, sample2, correctForLowNumberOfSamples, jvm_method + ) + + class DataTypeInstances(Enum): """ An enum class that types columns to scala datatypes diff --git a/tests/test_analyzers.py b/tests/test_analyzers.py index 175e8ae..17b0e61 100644 --- a/tests/test_analyzers.py +++ b/tests/test_analyzers.py @@ -12,9 +12,11 @@ ApproxQuantiles, Completeness, Compliance, + CategoricalDistanceMethod, Correlation, CountDistinct, DataType, + Distance, Distinctness, Entropy, Histogram, @@ -561,5 +563,85 @@ def test_fail_UniqueValueRatio(self): self.assertEqual(self.UniqueValueRatio(["a", "a"]), []) + def test_Distance_categorical_LInfinity(self): + distance = Distance(self.spark) + # Identical distributions -> zero raw L-infinity distance. + dist = {"a": 10, "b": 20, "c": 30} + self.assertEqual( + distance.categoricalDistance(dist, dist, correctForLowNumberOfSamples=True), 0.0 + ) + + # Raw L-infinity is the max abs difference of relative frequencies. + # dist1: a=0.5, b=0.5 ; dist2: a=1.0 -> max diff = 0.5 + result = distance.categoricalDistance( + {"a": 5, "b": 5}, {"a": 10}, correctForLowNumberOfSamples=True + ) + self.assertAlmostEqual(result, 0.5) + self.assertTrue(0.0 <= result <= 1.0) + + # With large samples, the default KS-corrected L-infinity is positive. + corrected = distance.categoricalDistance( + {"a": 1000, "b": 1000}, {"a": 2000, "b": 10} + ) + self.assertGreater(corrected, 0.0) + + def test_Distance_categorical_LInfinity_alpha(self): + # Exercises the LInfinity ``alpha`` path -- the only branch that marshals an + # ``Option[Double]`` into the JVM (``LInfinityMethod(Option.apply(alpha))``), + # which no other test covers. A supplied alpha scales the distance by the + # Kolmogorov-Smirnov critical value at that significance level. + distance = Distance(self.spark) + d1 = {"a": 1000, "b": 1000} + d2 = {"a": 2000, "b": 10} + + with_alpha = distance.categoricalDistance( + d1, d2, method=CategoricalDistanceMethod.LInfinity, alpha=0.05 + ) + # The Option[Double] is genuinely consumed: a different significance level + # produces a different distance (it is ignored / unmarshalled only if broken). + other_alpha = distance.categoricalDistance( + d1, d2, method=CategoricalDistanceMethod.LInfinity, alpha=0.5 + ) + self.assertIsInstance(with_alpha, float) + self.assertIsInstance(other_alpha, float) + self.assertGreaterEqual(with_alpha, 0.0) + self.assertNotEqual(with_alpha, other_alpha) + + def test_Distance_categorical_Chisquare(self): + distance = Distance(self.spark) + result = distance.categoricalDistance( + {"a": 100, "b": 100, "c": 100}, + {"a": 80, "b": 100, "c": 120}, + method=CategoricalDistanceMethod.Chisquare, + ) + # A non-identical pair yields a strictly positive chi-squared distance. + self.assertGreater(result, 0.0) + self.assertIsInstance(result, float) + + def test_Distance_categorical_single_category(self): + distance = Distance(self.spark) + # Single-category distributions: each is 100% of its only key, so the + # relative frequencies are identical and the raw L-infinity is 0.0. + result = distance.categoricalDistance( + {"a": 7}, {"a": 42}, correctForLowNumberOfSamples=True + ) + self.assertEqual(result, 0.0) + + def test_Distance_categorical_empty_dict_raises(self): + distance = Distance(self.spark) + with self.assertRaises(ValueError): + distance.categoricalDistance({}, {"a": 1}) + with self.assertRaises(ValueError): + distance.categoricalDistance({"a": 1}, {}) + with self.assertRaises(ValueError): + distance.categoricalDistance({}, {}) + + def test_Distance_categorical_invalid_method_raises(self): + distance = Distance(self.spark) + # A non-CategoricalDistanceMethod value hits the explicit else branch. + with self.assertRaises(ValueError): + distance.categoricalDistance({"a": 1}, {"a": 1}, method="not_a_method") + + if __name__ == "__main__": unittest.main()