Repository: great-expectations/great_expectations

Issue:
New approach for computating expect_column_values_to_be_in_set in spark environment
**Background**
We are trying to apply Great Expectations to check the data quality.
A foreign key checking between a 200k rows value_set against a 60k rows checking_set (60k rows) requires 1000 second to complete. This is much longer than our expectation in a spark environment. After looking into the source code of .expect_column_values_to_be_in_set() under a spark dataframe, we found that _"column[0].isin(value_set)"_ is the root cause of such behavior (see the first screen cap below)

**Is your feature request related to a problem? Please describe.**
Performance of the current solution which uses a python list as input of value_set deteriorates when a big dataset is involved.

Our input is a spark dataframe, it is casted into a list before it is thrown into the function.
After a series of testing, the runtime of the current solution grows linearly with len(value_set). This implies a large table may slow down the process. By testing on a ~10m rows value_set, the process cannot be completed after 3 days (we killed it at the end).

**Describe the solution you'd like**
In spark environment, we would like to use pyspark function which make use of the calculating power of spark to do the matching instead of using a python list
_i.e. use join() to replace isin() in the following module_ in _great_expectations/dataset/sparkdf_dataset.py_
![image](https://user-images.githubusercontent.com/90680179/142392615-06e8e162-1e20-4fe1-9b23-ee515e18c465.png)

**Testing result**
![image](https://user-images.githubusercontent.com/90680179/142389912-57007ad0-764a-43f3-ab91-662e92d884bd.png)
Join brings a tremendous improvement in runtime, which turns our big foreign key mapping from something like more than 3 days to just 10s under the Databricks environment
Therefore, joining is strongly recommended when the dataframe is a spark dataframe.

**Appendix**
Attached the codes for your further studies, thank you very much !!!! :)

```
import time
import numpy as np
import pandas as pd
from random import shuffle
from pyspark.sql import SparkSession
from pyspark.sql.functions import col,when
spark = SparkSession.builder.getOrCreate()

expect_column_values = spark.sql("select distinct member_key as column from rdm.f_transaction where business_date_key = 20211111")
to_be_in_set = spark.sql("select distinct member_key as column from rdm.d_member order by 1")

print(expect_column_values.count())
print(to_be_in_set.count())

df = expect_column_values.join(to_be_in_set,expect_column_values["column"] ==  to_be_in_set["column"] ,"left")
df1=df.withColumn("__success",when(to_be_in_set["column"] == None,False).otherwise(True))
df2 = df1.select(expect_column_values["column"],"__success").show(5,False)
```


Discussion:
Hey @derek-hk! Thanks so much for opening up this issue after our discussion in Slack.

We'll discuss internally and be in touch!
Hey @derek-hk! We would be able to accept a [contrib expectation](https://docs.greatexpectations.io/docs/guides/expectations/contributing/how_to_contribute_a_new_expectation_to_great_expectations) with this functionality in the short term; before accepting this as a change to a core expectation we would need testing to confirm that this approach is performant in both low-cardinality and high-cardinality situations.

Is this something you'd be willing to submit as a contribution? We're happy to provide guidance and review on that process.
Hi @austiezr 
Does it mean that it is preferred to exist as a new expectation instead of expect_column_values_to_be_in_set?
Hey @derek-hk! Yes, in the short term we would accept this as a new expectation.
Hi @austiezr 
I would like to help in the testing of the proposed approach.
Please let me know if anything I can help!
Hey @derek-hk! 

A couple of things you could do if interested:
- Submit a PR for this expectation as an [experimental Custom Expectation](https://docs.greatexpectations.io/docs/guides/expectations/contributing/how_to_contribute_a_new_expectation_to_great_expectations)
- Benchmark the performance of your new approach vs the current core approach on data of various size and complexity and report back here

Either / Both of these would help us drive the conversation forward around a change like this.
Hi @austiezr 
As far as I know, passing a spark dataframe into value_set and edit the return directly doesn't work.
So let me first share the performance of my approach vs the current approach.

**Methodology:**

1. Declare a dataframe called "column", which contains data with 248678 rows to be validated.
2. Declare the value_set for the testing, we have product_key, store_key, card_key and member_key as our testing case.
Note the following row count of different testing set: 
product_key: 199210
store_key: 728
card_key: 9874598
member_key: 9871723
2. Use `column.withColumn("__success", column[0].isin(value_set))`, where value_set is a list, to produce the `__success` column as follow.
![image](https://user-images.githubusercontent.com/90680179/152509358-29e4c6ed-902f-45c4-be0f-e28809a30468.png)

3. Use `column.join(value_set,column[0] == value_set[0] ,"left").withColumn("__success",when(value_set[0].isNull(),False).otherwise(True)).select(column[0],"__success")`, where value_set is a spark dataframe to produce the `__success` column as follow.
![image](https://user-images.githubusercontent.com/90680179/152509420-fd48f371-e3bb-48a4-80d7-87ee08b2ad18.png)

4. Compare and record the performance.

5. Result of the tests:
![image](https://user-images.githubusercontent.com/90680179/152510246-797b5e7d-c7af-47a8-af8c-11c2e6ce8cbf.png)

6. My conclusion: Joining approach is more efficient in both big value_set and small value_set

Resolving patch:
diff --git a/contrib/experimental/great_expectations_experimental/expectations/expect_column_values_to_be_in_set_spark_optimized.py b/contrib/experimental/great_expectations_experimental/expectations/expect_column_values_to_be_in_set_spark_optimized.py
new file mode 100644
--- /dev/null
+++ b/contrib/experimental/great_expectations_experimental/expectations/expect_column_values_to_be_in_set_spark_optimized.py
@@ -0,0 +1,231 @@
+import warnings
+from typing import Dict, Optional
+
+import numpy as np
+
+from great_expectations.core.expectation_configuration import ExpectationConfiguration
+from great_expectations.exceptions import InvalidExpectationConfigurationError
+from great_expectations.exceptions.exceptions import MetricResolutionError
+from great_expectations.execution_engine import ExecutionEngine, SparkDFExecutionEngine
+from great_expectations.expectations.expectation import ColumnExpectation
+from great_expectations.expectations.metrics import (
+    ColumnAggregateMetricProvider,
+    metric_value,
+)
+from great_expectations.expectations.metrics.import_manager import F, sparktypes
+from great_expectations.expectations.metrics.map_metric_provider import (
+    ColumnMapMetricProvider,
+    MetricDomainTypes,
+    MetricPartialFunctionTypes,
+    column_condition_partial,
+    metric_partial,
+)
+
+
+# This class defines a Metric to support your Expectation.
+# For most ColumnExpectations, the main business logic for calculation will live in this class.
+class ColumnValuesInSetSparkOptimized(ColumnAggregateMetricProvider):
+
+    metric_name = "column_values.in_set.spark_optimized"
+    value_keys = (
+        "column",
+        "value_set",
+    )
+
+    @metric_value(
+        engine=SparkDFExecutionEngine,
+        partial_fn_type=MetricPartialFunctionTypes.MAP_CONDITION_FN,
+        domain_type=MetricDomainTypes.COLUMN,
+    )
+    def _spark(
+        cls,
+        execution_engine: SparkDFExecutionEngine,
+        metric_domain_kwargs,
+        metric_value_kwargs,
+        metrics,
+        runtime_configuration,
+        **kwargs,
+    ):
+        column_name = metric_domain_kwargs.get("column")
+        value_set = metric_value_kwargs.get("value_set")
+        (
+            df,
+            compute_domain_kwargs,
+            accessor_domain_kwargs,
+        ) = execution_engine.get_compute_domain(
+            domain_kwargs=metric_domain_kwargs, domain_type=MetricDomainTypes.TABLE
+        )
+
+        try:
+            value_set = execution_engine.spark.createDataFrame(
+                data=value_set, schema=df.schema[column_name].dataType
+            )
+        except TypeError as e:
+            raise InvalidExpectationConfigurationError(
+                f"`value_set` must be of same type as `column` : {e}"
+            )
+
+        joined = df.join(value_set, df[column_name] == value_set["value"], "left")
+        success = joined.withColumn(
+            "__success", F.when(joined["value"].isNull(), False).otherwise(True)
+        )
+
+        return success.select(column_name, "__success").collect()
+
+
+# This class defines the Expectation itself
+class ExpectColumnValuesToBeInSetSparkOptimized(ColumnExpectation):
+    """Expect each column value to be in a given set; optimized using `join` for spark backends.
+
+    Args:
+        column (str): \
+            The column name.
+        value_set (set-like): \
+            A set of objects used for comparison.
+
+    Keyword Args:
+        mostly (None or a float between 0 and 1): \
+            Return `"success": True` if at least mostly fraction of values match the expectation. \
+            For more detail, see :ref:`mostly`.
+        strict (boolean or None) : If True, percentage of values in set must exceed mostly.
+
+    """
+
+    # This is a tuple consisting of all Metrics necessary to evaluate the Expectation.
+    metric_dependencies = ("column_values.in_set.spark_optimized",)
+
+    # This a tuple of parameter names that can affect whether the Expectation evaluates to True or False.
+    success_keys = ("column", "value_set", "mostly", "strict")
+
+    # This dictionary contains default values for any parameters that should have default values.
+    default_kwarg_values = {"mostly": 1, "strict": True, "value_set": []}
+
+    def validate_configuration(
+        self, configuration: Optional[ExpectationConfiguration]
+    ) -> None:
+        """
+        Validates that a configuration has been set, and sets a configuration if it has yet to be set. Ensures that
+        necessary configuration arguments have been provided for the validation of the expectation.
+
+        Args:
+            configuration (OPTIONAL[ExpectationConfiguration]): \
+                An optional Expectation Configuration entry that will be used to configure the expectation
+        Returns:
+            None. Raises InvalidExpectationConfigurationError if the config is not validated successfully
+        """
+
+        super().validate_configuration(configuration)
+        if configuration is None:
+            configuration = self.configuration
+        value_set = configuration.kwargs.get(
+            "value_set"
+        ) or self.default_kwarg_values.get("value_set")
+        column = configuration.kwargs.get("column")
+
+        try:
+            assert column is not None, "`column` must be specified"
+            assert (
+                "value_set" in configuration.kwargs or value_set
+            ), "value_set is required"
+            assert isinstance(
+                value_set, (list, set, dict)
+            ), "value_set must be a list, set, or dict"
+            if isinstance(value_set, dict):
+                assert (
+                    "$PARAMETER" in value_set
+                ), 'Evaluation Parameter dict for value_set kwarg must have "$PARAMETER" key.'
+        except AssertionError as e:
+            raise InvalidExpectationConfigurationError(str(e))
+
+    # This method performs a validation of your metrics against your success keys, returning a dict indicating the success or failure of the Expectation.
+    def _validate(
+        self,
+        configuration: ExpectationConfiguration,
+        metrics: Dict,
+        runtime_configuration: dict = None,
+        execution_engine: ExecutionEngine = None,
+    ):
+        mostly = configuration["kwargs"].get("mostly")
+        strict = configuration["kwargs"].get("strict")
+        result = metrics.get("column_values.in_set.spark_optimized")
+        result = dict(result)
+
+        if strict is True and mostly < 1:
+            success = (sum(result.values()) / len(result.values())) > mostly
+        else:
+            success = (sum(result.values()) / len(result.values())) >= mostly
+
+        return {"success": success, "result": {"observed_values": result}}
+
+    examples = [
+        {
+            "data": {
+                "col1": ["ES", "BE", "FR", "DE", "CH"],
+                "col2": [1, 2, 3, 5, 8],
+            },
+            "tests": [
+                {
+                    "title": "positive_test",
+                    "exact_match_out": False,
+                    "include_in_gallery": True,
+                    "in": {
+                        "column": "col1",
+                        "value_set": ["ES", "BE", "UK"],
+                        "mostly": 0.4,
+                        "strict": False,
+                    },
+                    "out": {
+                        "success": True,
+                    },
+                },
+                {
+                    "title": "negative_test",
+                    "exact_match_out": False,
+                    "include_in_gallery": True,
+                    "in": {
+                        "column": "col2",
+                        "value_set": [3, 8, 1, 22, 74],
+                        "mostly": 0.6,
+                        "strict": True,
+                    },
+                    "out": {
+                        "success": False,
+                    },
+                },
+                {
+                    "title": "failing_test",
+                    "exact_match_out": False,
+                    "include_in_gallery": True,
+                    "in": {
+                        "column": "col2",
+                        "value_set": ["ES", "BE", "UK"],
+                        "mostly": 0.6,
+                        "strict": True,
+                        "catch_exceptions": True,
+                    },
+                    "out": {},
+                    "error": {
+                        "traceback_substring": "`value_set` must be of same type as `column`",
+                    },
+                },
+            ],
+            "only_for": ["spark"],
+        }
+    ]
+
+    # This object contains metadata for display in the public Gallery
+    library_metadata = {
+        "tags": [
+            "column values in set",
+            "experimental",
+            "spark optimized",
+        ],  # Tags for this Expectation in the Gallery
+        "contributors": [  # Github handles for all contributors to this Expectation.
+            "@austiezr",
+            "@derek-hk",  # Don't forget to add your github handle here!
+        ],
+    }
+
+
+if __name__ == "__main__":
+    ExpectColumnValuesToBeInSetSparkOptimized().print_diagnostic_checklist()


Test patch: