Repository: great-expectations/great_expectations

Issue:
Memory problems in map_condition_auxillary_methods._spark_map_condition_index() in 0.16.6
**Describe the bug**
I was upgrading our GX stack to the latest version (0.16.6) from 0.15.26. Some of our benchmarks on medium-sized data (~ 5 million rows) were running out of memory where they had run fine before. I did some debugging and it looks like there are two problems in `map_condition_auxillary_methods._spark_map_condition_index()`.

Apologies for jamming two problems into one issue, but they're literally in two consecutive lines in the code.

**Problem 1** 
The following code attempts to limit the number of results based on the unexpected count. However, it doesn't re-assign the the dataframe.
```
if result_format["result_format"] != "COMPLETE":
        filtered.limit(result_format["partial_unexpected_count"])
```

After changing it to following:

```
    if result_format["result_format"] != "COMPLETE":
        filtered = filtered.limit(result_format["partial_unexpected_count"])
```

The number of objects created by `run_checkpoint` dropped dramatically. On a pretty small dataset (500K rows), the number of python dictionaries created went from 

```
                                     types |   # objects |   total size
========================================== | =========== | ============
                                      dict |     1601406 |    354.50 MB
...
```

to

```
                                    types |   # objects |   total size
========================================== | =========== | ============
                                     dict |        5600 |      1.23 MB
...
```

The difference was even more pronounced on larger datasets.

**Problem 2**

In the original code, I noticed that this line was consuming the bulk of our runtime per the spark UI.

```
    if result_format["result_format"] != "COMPLETE":
        filtered.limit(result_format["partial_unexpected_count"])

    for row in filtered.collect():  <<<----- This was the culprit
        dict_to_add: dict = {}
        for col_name in columns_to_keep:
            dict_to_add[col_name] = row[col_name]
        unexpected_index_list.append(dict_to_add)

    return unexpected_index_list
```

After looking at this in the debugger, the problem was that the dataframe `filtered` had all 260 columns present even though we only cared about a few in `columns_to_keep`. The code inside the `collect()` for loop ignores anything that isn't in `columns_to_keep`.  Adding the following line dramatically reduced memory (I couldn't even run our 5 million row test set without kube's OOM killer stomping our container).

```
    if result_format["result_format"] != "COMPLETE":
        filtered.limit(result_format["partial_unexpected_count"])

    # Prune the dataframe down only the columns we care about
    filtered = filtered.select(columns_to_keep)  <<<------ This is the line I added

    for row in filtered.collect():
        dict_to_add: dict = {}
        for col_name in columns_to_keep:
            dict_to_add[col_name] = row[col_name]
        unexpected_index_list.append(dict_to_add)

    return unexpected_index_list
```

**To Reproduce**
Steps to reproduce the behavior:
1. You should be able to see a difference running sufficiently large data with some expectations that land in this part of the code. For our suite, a good 1/3 of the expectations landed here.

**Expected behavior**
Be as frugal as we can be with memory. There's still a gap in runtime between 0.15.26 and 0.16.6 that I'm hunting down. It's about 25% slower even with the above changes, but there could be some API changes that we're not utilizing properly. Still, I believe the above changes are good general changes to consider.

**Environment (please complete the following information):**
 - Operating System:  MacOS (running on an Ubuntu container)
 - Great Expectations Version: 0.16.6

**Additional context**
For now, I'll just build these in as overrides into our container an press on, but I thought it might be good to share these findings with the community for consideration.



Discussion:
Hey @brendan-bull! Thanks for raising this, and for all the context provided. We are reviewing now 🚀 
@austiezr - Thanks for looking! One clarification to this comment above "It's about 25% slower even with the above changes..." - After updating spark from 3.2.1 to 3.3.2, that gap has been closed and we can now replicate previous performance with the two code changes from above in place.

Resolving patch:
diff --git a/great_expectations/expectations/metrics/map_metric_provider/map_condition_auxilliary_methods.py b/great_expectations/expectations/metrics/map_metric_provider/map_condition_auxilliary_methods.py
--- a/great_expectations/expectations/metrics/map_metric_provider/map_condition_auxilliary_methods.py
+++ b/great_expectations/expectations/metrics/map_metric_provider/map_condition_auxilliary_methods.py
@@ -772,7 +772,10 @@ def _spark_map_condition_index(
             )
 
     if result_format["result_format"] != "COMPLETE":
-        filtered.limit(result_format["partial_unexpected_count"])
+        filtered = filtered.limit(result_format["partial_unexpected_count"])
+
+    # Prune the dataframe down only the columns we care about
+    filtered = filtered.select(columns_to_keep)
 
     for row in filtered.collect():
         dict_to_add: dict = {}


Test patch:
diff --git a/tests/checkpoint/test_checkpoint_result_format.py b/tests/checkpoint/test_checkpoint_result_format.py
--- a/tests/checkpoint/test_checkpoint_result_format.py
+++ b/tests/checkpoint/test_checkpoint_result_format.py
@@ -2085,6 +2085,63 @@ def test_spark_result_format_in_checkpoint_pk_defined_one_expectation_summary_ou
     assert evrs[0]["results"][0]["result"].get("unexpected_index_query") is None
 
 
+@pytest.mark.integration
+def test_spark_result_format_in_checkpoint_pk_defined_one_expectation_summary_output_limit_1(
+    in_memory_runtime_context: AbstractDataContext,
+    batch_request_for_spark_unexpected_rows_and_index: dict,
+    reference_checkpoint_config_for_unexpected_column_names: dict,
+    expectation_config_expect_column_values_to_be_in_set: ExpectationConfiguration,
+    expected_unexpected_indices_output: list[dict[str, str | int]],
+):
+    """
+    What does this test?
+        - unexpected_index_column defined in Checkpoint only.
+        - SUMMARY output, which means we have `partial_unexpected_index_list` only
+        - 1 Expectations added to suite
+        - limit is 1 so we only get 1 output in the `partial_unexpected_index_list`
+    """
+    dict_to_update_checkpoint: dict = {
+        "result_format": {
+            "result_format": "SUMMARY",
+            "partial_unexpected_count": 1,
+            "unexpected_index_column_names": ["pk_1"],
+        }
+    }
+    context: DataContext = _add_expectations_and_checkpoint(
+        data_context=in_memory_runtime_context,
+        checkpoint_config=reference_checkpoint_config_for_unexpected_column_names,
+        expectations_list=[expectation_config_expect_column_values_to_be_in_set],
+        dict_to_update_checkpoint=dict_to_update_checkpoint,
+    )
+
+    result: CheckpointResult = context.run_checkpoint(
+        checkpoint_name="my_checkpoint",
+        expectation_suite_name="metrics_exp",
+        batch_request=batch_request_for_spark_unexpected_rows_and_index,
+    )
+    evrs: List[ExpectationSuiteValidationResult] = result.list_validation_results()
+
+    index_column_names: List[str] = evrs[0]["results"][0]["result"][
+        "unexpected_index_column_names"
+    ]
+    assert index_column_names == ["pk_1"]
+
+    first_result_full_list: List[Dict[str, Any]] = evrs[0]["results"][0]["result"].get(
+        "unexpected_index_list"
+    )
+    assert not first_result_full_list
+    first_result_partial_list: List[Dict[str, Any]] = evrs[0]["results"][0]["result"][
+        "partial_unexpected_index_list"
+    ]
+    assert first_result_partial_list == [
+        {
+            "animals": "giraffe",
+            "pk_1": 3,
+        }
+    ]
+    assert evrs[0]["results"][0]["result"].get("unexpected_index_query") is None
+
+
 @pytest.mark.integration
 def test_spark_result_format_in_checkpoint_pk_defined_one_expectation_basic_output(
     in_memory_runtime_context: AbstractDataContext,