Skip to content

Commit

Permalink
Adding Wilson Score Confidence Interval Strategy (awslabs#567)
Browse files Browse the repository at this point in the history
* Configurable RetainCompletenessRule

* Add doc string

* Add default completeness const

* Add ConfidenceIntervalStrategy

* Add Separate Wilson and Wald Interval Test

* Add License information, Fix formatting

* Add License information

* formatting fix

* Update documentation

* Make WaldInterval the default strategy for now

* Formatting import to per line

* Separate group import to per line import
  • Loading branch information
zeotuan authored and eycho-am committed Oct 9, 2024
1 parent 5f723d8 commit 2e0efde
Show file tree
Hide file tree
Showing 9 changed files with 299 additions and 54 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
package com.amazon.deequ.examples

import com.amazon.deequ.examples.ExampleUtils.withSpark
import com.amazon.deequ.suggestions.rules.RetainCompletenessRule
import com.amazon.deequ.suggestions.rules.interval.WilsonScoreIntervalStrategy
import com.amazon.deequ.suggestions.{ConstraintSuggestionRunner, Rules}

private[examples] object ConstraintSuggestionExample extends App {
Expand Down Expand Up @@ -51,6 +53,10 @@ private[examples] object ConstraintSuggestionExample extends App {
val suggestionResult = ConstraintSuggestionRunner()
.onData(data)
.addConstraintRules(Rules.EXTENDED)
// We can also add our own constraint and customize constraint parameters
.addConstraintRule(
RetainCompletenessRule(intervalStrategy = WilsonScoreIntervalStrategy())
)
.run()

// We can now investigate the constraints that deequ suggested. We get a textual description
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,17 @@ val suggestionResult = ConstraintSuggestionRunner()
.run()
```

Alternatively, we also support customizing and adding individual constraint rule using `addConstraintRule()`
```scala
val suggestionResult = ConstraintSuggestionRunner()
.onData(data)

.addConstraintRule(
RetainCompletenessRule(intervalStrategy = WilsonScoreIntervalStrategy())
)
.run()
```

We can now investigate the constraints that deequ suggested. We get a textual description and the corresponding scala code for each suggested constraint. Note that the constraint suggestion is based on heuristic rules and assumes that the data it is shown is 'static' and correct, which might often not be the case in the real world. Therefore the suggestions should always be manually reviewed before being applied in real deployments.
```scala
suggestionResult.constraintSuggestions.foreach { case (column, suggestions) =>
Expand Down Expand Up @@ -92,3 +103,5 @@ The corresponding scala code is .isContainedIn("status", Array("DELAYED", "UNKNO
Currently, we leave it up to the user to decide whether they want to apply the suggested constraints or not, and provide the corresponding Scala code for convenience. For larger datasets, it makes sense to evaluate the suggested constraints on some held-out portion of the data to see whether they hold or not. You can test this by adding an invocation of `.useTrainTestSplitWithTestsetRatio(0.1)` to the `ConstraintSuggestionRunner`. With this configuration, it would compute constraint suggestions on 90% of the data and evaluate the suggested constraints on the remaining 10%.

Finally, we would also like to note that the constraint suggestion code provides access to the underlying [column profiles](https://github.com/awslabs/deequ/blob/master/src/main/scala/com/amazon/deequ/examples/data_profiling_example.md) that it computed via `suggestionResult.columnProfiles`.

An [executable and extended version of this example](https://github.com/awslabs/deequ/blob/master/src/main/scala/com/amazon/deequ/examples/.scala) is part of our code base.
Original file line number Diff line number Diff line change
Expand Up @@ -23,16 +23,17 @@ import com.amazon.deequ.metrics.DistributionValue
import com.amazon.deequ.profiles.ColumnProfile
import com.amazon.deequ.suggestions.ConstraintSuggestion
import com.amazon.deequ.suggestions.ConstraintSuggestionWithValue
import com.amazon.deequ.suggestions.rules.interval.ConfidenceIntervalStrategy.defaultIntervalStrategy
import com.amazon.deequ.suggestions.rules.interval.ConfidenceIntervalStrategy
import org.apache.commons.lang3.StringEscapeUtils

import scala.math.BigDecimal.RoundingMode

/** If we see a categorical range for most values in a column, we suggest an IS IN (...)
* constraint that should hold for most values */
case class FractionalCategoricalRangeRule(
targetDataCoverageFraction: Double = 0.9,
categorySorter: Array[(String, DistributionValue)] => Array[(String, DistributionValue)] =
categories => categories.sortBy({ case (_, value) => value.absolute }).reverse
categories => categories.sortBy({ case (_, value) => value.absolute }).reverse,
intervalStrategy: ConfidenceIntervalStrategy = defaultIntervalStrategy
) extends ConstraintRule[ColumnProfile] {

override def shouldBeApplied(profile: ColumnProfile, numRecords: Long): Boolean = {
Expand Down Expand Up @@ -79,11 +80,8 @@ case class FractionalCategoricalRangeRule(

val p = ratioSums
val n = numRecords
val z = 1.96

// TODO this needs to be more robust for p's close to 0 or 1
val targetCompliance = BigDecimal(p - z * math.sqrt(p * (1 - p) / n))
.setScale(2, RoundingMode.DOWN).toDouble
val targetCompliance = intervalStrategy.calculateTargetConfidenceInterval(p, n).lowerBound

val description = s"'${profile.column}' has value range $categoriesSql for at least " +
s"${targetCompliance * 100}% of values"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,8 @@ import com.amazon.deequ.profiles.ColumnProfile
import com.amazon.deequ.suggestions.CommonConstraintSuggestion
import com.amazon.deequ.suggestions.ConstraintSuggestion
import com.amazon.deequ.suggestions.rules.RetainCompletenessRule._

import scala.math.BigDecimal.RoundingMode
import com.amazon.deequ.suggestions.rules.interval.ConfidenceIntervalStrategy.defaultIntervalStrategy
import com.amazon.deequ.suggestions.rules.interval.ConfidenceIntervalStrategy

/**
* If a column is incomplete in the sample, we model its completeness as a binomial variable,
Expand All @@ -33,21 +33,18 @@ import scala.math.BigDecimal.RoundingMode
*/
case class RetainCompletenessRule(
minCompleteness: Double = defaultMinCompleteness,
maxCompleteness: Double = defaultMaxCompleteness
maxCompleteness: Double = defaultMaxCompleteness,
intervalStrategy: ConfidenceIntervalStrategy = defaultIntervalStrategy
) extends ConstraintRule[ColumnProfile] {
override def shouldBeApplied(profile: ColumnProfile, numRecords: Long): Boolean = {
profile.completeness > minCompleteness && profile.completeness < maxCompleteness
}

override def candidate(profile: ColumnProfile, numRecords: Long): ConstraintSuggestion = {

val p = profile.completeness
val n = numRecords
val z = 1.96

// TODO this needs to be more robust for p's close to 0 or 1
val targetCompleteness = BigDecimal(p - z * math.sqrt(p * (1 - p) / n))
.setScale(2, RoundingMode.DOWN).toDouble
val targetCompleteness = intervalStrategy.calculateTargetConfidenceInterval(
profile.completeness,
numRecords
).lowerBound

val constraint = completenessConstraint(profile.column, _ >= targetCompleteness)

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
/**
* Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License"). You may not
* use this file except in compliance with the License. A copy of the License
* is located at
*
* http://aws.amazon.com/apache2.0/
*
* or in the "license" file accompanying this file. This file is distributed on
* an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*
*/

package com.amazon.deequ.suggestions.rules.interval

import breeze.stats.distributions.{Gaussian, Rand}
import com.amazon.deequ.suggestions.rules.interval.ConfidenceIntervalStrategy._

/**
* Strategy for calculate confidence interval
* */
trait ConfidenceIntervalStrategy {

/**
* Generated confidence interval interval
* @param pHat sample of the population that share a trait
* @param numRecords overall number of records
* @param confidence confidence level of method used to estimate the interval.
* @return
*/
def calculateTargetConfidenceInterval(
pHat: Double,
numRecords: Long,
confidence: Double = defaultConfidence
): ConfidenceInterval

def validateInput(pHat: Double, confidence: Double): Unit = {
require(0.0 <= pHat && pHat <= 1.0, "pHat must be between 0.0 and 1.0")
require(0.0 <= confidence && confidence <= 1.0, "confidence must be between 0.0 and 1.0")
}

def calculateZScore(confidence: Double): Double = Gaussian(0, 1)(Rand).inverseCdf(1 - ((1.0 - confidence)/ 2.0))
}

object ConfidenceIntervalStrategy {
val defaultConfidence = 0.95
val defaultIntervalStrategy: ConfidenceIntervalStrategy = WaldIntervalStrategy()

case class ConfidenceInterval(lowerBound: Double, upperBound: Double)
}


Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
/**
* Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License"). You may not
* use this file except in compliance with the License. A copy of the License
* is located at
*
* http://aws.amazon.com/apache2.0/
*
* or in the "license" file accompanying this file. This file is distributed on
* an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*
*/

package com.amazon.deequ.suggestions.rules.interval

import com.amazon.deequ.suggestions.rules.interval.ConfidenceIntervalStrategy.ConfidenceInterval
import com.amazon.deequ.suggestions.rules.interval.ConfidenceIntervalStrategy.defaultConfidence

import scala.math.BigDecimal.RoundingMode

/**
* Implements the Wald Interval method for creating a binomial proportion confidence interval. Provided for backwards
* compatibility. using [[WaldIntervalStrategy]] for calculating confidence interval can be problematic when dealing
* with small sample sizes or proportions close to 0 or 1. It also have poorer coverage and might produce confidence
* limit outside the range of [0,1]
* @see <a
* href="http://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval#Normal_approximation_interval">
* Normal approximation interval (Wikipedia)</a>
*/
@deprecated("WilsonScoreIntervalStrategy is recommended for calculating confidence interval")
case class WaldIntervalStrategy() extends ConfidenceIntervalStrategy {
def calculateTargetConfidenceInterval(
pHat: Double,
numRecords: Long,
confidence: Double = defaultConfidence
): ConfidenceInterval = {
validateInput(pHat, confidence)
val successRatio = BigDecimal(pHat)
val marginOfError = BigDecimal(calculateZScore(confidence) * math.sqrt(pHat * (1 - pHat) / numRecords))
val lowerBound = (successRatio - marginOfError).setScale(2, RoundingMode.DOWN).toDouble
val upperBound = (successRatio + marginOfError).setScale(2, RoundingMode.UP).toDouble
ConfidenceInterval(lowerBound, upperBound)
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
/**
* Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License"). You may not
* use this file except in compliance with the License. A copy of the License
* is located at
*
* http://aws.amazon.com/apache2.0/
*
* or in the "license" file accompanying this file. This file is distributed on
* an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*
*/

package com.amazon.deequ.suggestions.rules.interval

import com.amazon.deequ.suggestions.rules.interval.ConfidenceIntervalStrategy.ConfidenceInterval
import com.amazon.deequ.suggestions.rules.interval.ConfidenceIntervalStrategy.defaultConfidence

import scala.math.BigDecimal.RoundingMode

/**
* Using Wilson score method for creating a binomial proportion confidence interval.
*
* @see <a
* href="http://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval#Wilson_score_interval">
* Wilson score interval (Wikipedia)</a>
*/
case class WilsonScoreIntervalStrategy() extends ConfidenceIntervalStrategy {

def calculateTargetConfidenceInterval(
pHat: Double, numRecords: Long,
confidence: Double = defaultConfidence
): ConfidenceInterval = {
validateInput(pHat, confidence)
val zScore = calculateZScore(confidence)
val zSquareOverN = math.pow(zScore, 2) / numRecords
val factor = 1.0 / (1 + zSquareOverN)
val adjustedSuccessRatio = pHat + zSquareOverN/2
val marginOfError = zScore * math.sqrt(pHat * (1 - pHat)/numRecords + zSquareOverN/(4 * numRecords))
val lowerBound = BigDecimal(factor * (adjustedSuccessRatio - marginOfError)).setScale(2, RoundingMode.DOWN).toDouble
val upperBound = BigDecimal(factor * (adjustedSuccessRatio + marginOfError)).setScale(2, RoundingMode.UP).toDouble
ConfidenceInterval(lowerBound, upperBound)
}
}
Loading

0 comments on commit 2e0efde

Please sign in to comment.