We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Not able to set persist to False using Spark Execution Engine 
Hey Guys, I was trying to migrate GX from 0.16.5 and am not being able to, cause apparently since GX 0.16.12 supposedly there was a fix for the persist parameter to work. The thing is that I wanted it to be False, but it seems the parameter is not taking effect.
Anyone facing similar problems?
I was following this guide [How to connect to in-memory data in a Spark dataframe | Great Expectations](https://docs.greatexpectations.io/docs/0.15.50/guides/connecting_to_your_data/in_memory/spark/), which seems to have been removed. Also tried with fluent approach and still failing.
Is the property persist being considered/passed? Shouldn’t it be one of the parameters in add_or_update_spark ?

More information: https://discourse.greatexpectations.io/t/not-able-to-set-persist-to-false-using-spark-execution-engine/1320?u=jose.correia

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `great_expectations/datasource/fluent/spark_datasource.py`
Content:
```
1 from __future__ import annotations
2 
3 import logging
4 from pprint import pformat as pf
5 from typing import (
6     TYPE_CHECKING,
7     ClassVar,
8     Dict,
9     Generic,
10     List,
11     Literal,
12     Optional,
13     Type,
14     TypeVar,
15     Union,
16 )
17 
18 import pydantic
19 from pydantic import StrictBool, StrictFloat, StrictInt, StrictStr
20 
21 import great_expectations.exceptions as gx_exceptions
22 from great_expectations.compatibility.pyspark import DataFrame, pyspark
23 from great_expectations.core._docs_decorators import (
24     deprecated_argument,
25     new_argument,
26     public_api,
27 )
28 from great_expectations.core.batch_spec import RuntimeDataBatchSpec
29 from great_expectations.datasource.fluent import BatchRequest
30 from great_expectations.datasource.fluent.constants import (
31     _DATA_CONNECTOR_NAME,
32 )
33 from great_expectations.datasource.fluent.interfaces import (
34     Batch,
35     DataAsset,
36     Datasource,
37     _DataAssetT,
38 )
39 
40 if TYPE_CHECKING:
41     from typing_extensions import TypeAlias
42 
43     from great_expectations.datasource.fluent.interfaces import BatchMetadata
44     from great_expectations.execution_engine import SparkDFExecutionEngine
45 
46 
47 logger = logging.getLogger(__name__)
48 
49 
50 # this enables us to include dataframe in the json schema
51 _SparkDataFrameT = TypeVar("_SparkDataFrameT")
52 
53 SparkConfig: TypeAlias = Dict[
54     StrictStr, Union[StrictStr, StrictInt, StrictFloat, StrictBool]
55 ]
56 
57 
58 class SparkDatasourceError(Exception):
59     pass
60 
61 
62 class _SparkDatasource(Datasource):
63     # instance attributes
64     spark_config: Union[SparkConfig, None] = None
65     force_reuse_spark_context: bool = True
66 
67     @staticmethod
68     def _update_asset_forward_refs(asset_type: Type[_DataAssetT]) -> None:
69         # Only update forward refs if pyspark types are available.
70         if pyspark:
71             asset_type.update_forward_refs()
72 
73     # Abstract Methods
74     @property
75     def execution_engine_type(self) -> Type[SparkDFExecutionEngine]:
76         """Return the SparkDFExecutionEngine unless the override is set"""
77         from great_expectations.execution_engine.sparkdf_execution_engine import (
78             SparkDFExecutionEngine,
79         )
80 
81         return SparkDFExecutionEngine
82 
83     def test_connection(self, test_assets: bool = True) -> None:
84         """Test the connection for the _SparkDatasource.
85 
86         Args:
87             test_assets: If assets have been passed to the _SparkDatasource,
88                          an attempt can be made to test them as well.
89 
90         Raises:
91             TestConnectionError: If the connection test fails.
92         """
93         raise NotImplementedError(
94             """One needs to implement "test_connection" on a _SparkDatasource subclass."""
95         )
96 
97     # End Abstract Methods
98 
99 
100 class DataFrameAsset(DataAsset, Generic[_SparkDataFrameT]):
101     # instance attributes
102     type: Literal["dataframe"] = "dataframe"
103     # TODO: <Alex>05/31/2023: Upon removal of deprecated "dataframe" argument to "PandasDatasource.add_dataframe_asset()", default can be deleted.</Alex>
104     dataframe: Optional[_SparkDataFrameT] = pydantic.Field(
105         default=None, exclude=True, repr=False
106     )
107 
108     class Config:
109         extra = pydantic.Extra.forbid
110 
111     @pydantic.validator("dataframe")
112     def _validate_dataframe(cls, dataframe: DataFrame) -> DataFrame:
113         if not (DataFrame and isinstance(dataframe, DataFrame)):  # type: ignore[truthy-function]
114             raise ValueError("dataframe must be of type pyspark.sql.DataFrame")
115 
116         return dataframe
117 
118     def test_connection(self) -> None:
119         ...
120 
121     @property
122     def batch_request_options(self) -> tuple[str, ...]:
123         return tuple()
124 
125     def _get_reader_method(self) -> str:
126         raise NotImplementedError(
127             """Spark DataFrameAsset does not implement "_get_reader_method()" method, because DataFrame is already available."""
128         )
129 
130     def _get_reader_options_include(self) -> set[str]:
131         raise NotImplementedError(
132             """Spark DataFrameAsset does not implement "_get_reader_options_include()" method, because DataFrame is already available."""
133         )
134 
135     @public_api
136     # TODO: <Alex>05/31/2023: Upon removal of deprecated "dataframe" argument to "PandasDatasource.add_dataframe_asset()", its validation code must be deleted.</Alex>
137     @new_argument(
138         argument_name="dataframe",
139         message='The "dataframe" argument is no longer part of "PandasDatasource.add_dataframe_asset()" method call; instead, "dataframe" is the required argument to "DataFrameAsset.build_batch_request()" method.',
140         version="0.16.15",
141     )
142     def build_batch_request(
143         self, dataframe: Optional[_SparkDataFrameT] = None
144     ) -> BatchRequest:
145         """A batch request that can be used to obtain batches for this DataAsset.
146 
147         Args:
148             dataframe: The Spark Dataframe containing the data for this DataFrame data asset.
149 
150         Returns:
151             A BatchRequest object that can be used to obtain a batch list from a Datasource by calling the
152             get_batch_list_from_batch_request method.
153         """
154         if dataframe is None:
155             df = self.dataframe
156         else:
157             df = dataframe
158 
159         if df is None:
160             raise ValueError(
161                 "Cannot build batch request for dataframe asset without a dataframe"
162             )
163 
164         self.dataframe = df
165 
166         return BatchRequest(
167             datasource_name=self.datasource.name,
168             data_asset_name=self.name,
169             options={},
170         )
171 
172     def _validate_batch_request(self, batch_request: BatchRequest) -> None:
173         """Validates the batch_request has the correct form.
174 
175         Args:
176             batch_request: A batch request object to be validated.
177         """
178         if not (
179             batch_request.datasource_name == self.datasource.name
180             and batch_request.data_asset_name == self.name
181             and not batch_request.options
182         ):
183             expect_batch_request_form = BatchRequest(
184                 datasource_name=self.datasource.name,
185                 data_asset_name=self.name,
186                 options={},
187                 batch_slice=batch_request._batch_slice_input,  # type: ignore[attr-defined]
188             )
189             raise gx_exceptions.InvalidBatchRequestError(
190                 "BatchRequest should have form:\n"
191                 f"{pf(expect_batch_request_form.dict())}\n"
192                 f"but actually has form:\n{pf(batch_request.dict())}\n"
193             )
194 
195     def get_batch_list_from_batch_request(
196         self, batch_request: BatchRequest
197     ) -> list[Batch]:
198         self._validate_batch_request(batch_request)
199 
200         batch_spec = RuntimeDataBatchSpec(batch_data=self.dataframe)
201         execution_engine: SparkDFExecutionEngine = (
202             self.datasource.get_execution_engine()
203         )
204         data, markers = execution_engine.get_batch_data_and_markers(
205             batch_spec=batch_spec
206         )
207 
208         # batch_definition (along with batch_spec and markers) is only here to satisfy a
209         # legacy constraint when computing usage statistics in a validator. We hope to remove
210         # it in the future.
211         # imports are done inline to prevent a circular dependency with core/batch.py
212         from great_expectations.core import IDDict
213         from great_expectations.core.batch import BatchDefinition
214 
215         batch_definition = BatchDefinition(
216             datasource_name=self.datasource.name,
217             data_connector_name=_DATA_CONNECTOR_NAME,
218             data_asset_name=self.name,
219             batch_identifiers=IDDict(batch_request.options),
220             batch_spec_passthrough=None,
221         )
222 
223         batch_metadata: BatchMetadata = self._get_batch_metadata_from_batch_request(
224             batch_request=batch_request
225         )
226 
227         # Some pydantic annotations are postponed due to circular imports.
228         # Batch.update_forward_refs() will set the annotations before we
229         # instantiate the Batch class since we can import them in this scope.
230         Batch.update_forward_refs()
231 
232         return [
233             Batch(
234                 datasource=self.datasource,
235                 data_asset=self,
236                 batch_request=batch_request,
237                 data=data,
238                 metadata=batch_metadata,
239                 legacy_batch_markers=markers,
240                 legacy_batch_spec=batch_spec,
241                 legacy_batch_definition=batch_definition,
242             )
243         ]
244 
245 
246 @public_api
247 class SparkDatasource(_SparkDatasource):
248     # class attributes
249     asset_types: ClassVar[List[Type[DataAsset]]] = [DataFrameAsset]
250 
251     # instance attributes
252     type: Literal["spark"] = "spark"
253 
254     assets: List[DataFrameAsset] = []  # type: ignore[assignment]
255 
256     def test_connection(self, test_assets: bool = True) -> None:
257         ...
258 
259     @public_api
260     @deprecated_argument(
261         argument_name="dataframe",
262         message='The "dataframe" argument is no longer part of "PandasDatasource.add_dataframe_asset()" method call; instead, "dataframe" is the required argument to "DataFrameAsset.build_batch_request()" method.',
263         version="0.16.15",
264     )
265     def add_dataframe_asset(
266         self,
267         name: str,
268         dataframe: Optional[_SparkDataFrameT] = None,
269         batch_metadata: Optional[BatchMetadata] = None,
270     ) -> DataFrameAsset:
271         """Adds a Dataframe DataAsset to this SparkDatasource object.
272 
273         Args:
274             name: The name of the DataFrame asset. This can be any arbitrary string.
275             dataframe: The Spark Dataframe containing the data for this DataFrame data asset.
276             batch_metadata: An arbitrary user defined dictionary with string keys which will get inherited by any
277                             batches created from the asset.
278 
279         Returns:
280             The DataFameAsset that has been added to this datasource.
281         """
282         asset: DataFrameAsset = DataFrameAsset(
283             name=name,
284             batch_metadata=batch_metadata or {},
285         )
286         asset.dataframe = dataframe
287         return self._add_asset(asset=asset)
288
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/great_expectations/datasource/fluent/spark_datasource.py b/great_expectations/datasource/fluent/spark_datasource.py
--- a/great_expectations/datasource/fluent/spark_datasource.py
+++ b/great_expectations/datasource/fluent/spark_datasource.py
@@ -63,6 +63,7 @@
     # instance attributes
     spark_config: Union[SparkConfig, None] = None
     force_reuse_spark_context: bool = True
+    persist: bool = True
 
     @staticmethod
     def _update_asset_forward_refs(asset_type: Type[_DataAssetT]) -> None:

{"golden_diff": "diff --git a/great_expectations/datasource/fluent/spark_datasource.py b/great_expectations/datasource/fluent/spark_datasource.py\n--- a/great_expectations/datasource/fluent/spark_datasource.py\n+++ b/great_expectations/datasource/fluent/spark_datasource.py\n@@ -63,6 +63,7 @@\n     # instance attributes\n     spark_config: Union[SparkConfig, None] = None\n     force_reuse_spark_context: bool = True\n+    persist: bool = True\n \n     @staticmethod\n     def _update_asset_forward_refs(asset_type: Type[_DataAssetT]) -> None:\n", "issue": "Not able to set persist to False using Spark Execution Engine \nHey Guys, I was trying to migrate GX from 0.16.5 and am not being able to, cause apparently since GX 0.16.12 supposedly there was a fix for the persist parameter to work. The thing is that I wanted it to be False, but it seems the parameter is not taking effect.\r\nAnyone facing similar problems?\r\nI was following this guide [How to connect to in-memory data in a Spark dataframe | Great Expectations](https://docs.greatexpectations.io/docs/0.15.50/guides/connecting_to_your_data/in_memory/spark/), which seems to have been removed. Also tried with fluent approach and still failing.\r\nIs the property persist being considered/passed? Shouldn\u2019t it be one of the parameters in add_or_update_spark ?\r\n\r\nMore information: https://discourse.greatexpectations.io/t/not-able-to-set-persist-to-false-using-spark-execution-engine/1320?u=jose.correia\n", "before_files": [{"content": "from __future__ import annotations\n\nimport logging\nfrom pprint import pformat as pf\nfrom typing import (\n    TYPE_CHECKING,\n    ClassVar,\n    Dict,\n    Generic,\n    List,\n    Literal,\n    Optional,\n    Type,\n    TypeVar,\n    Union,\n)\n\nimport pydantic\nfrom pydantic import StrictBool, StrictFloat, StrictInt, StrictStr\n\nimport great_expectations.exceptions as gx_exceptions\nfrom great_expectations.compatibility.pyspark import DataFrame, pyspark\nfrom great_expectations.core._docs_decorators import (\n    deprecated_argument,\n    new_argument,\n    public_api,\n)\nfrom great_expectations.core.batch_spec import RuntimeDataBatchSpec\nfrom great_expectations.datasource.fluent import BatchRequest\nfrom great_expectations.datasource.fluent.constants import (\n    _DATA_CONNECTOR_NAME,\n)\nfrom great_expectations.datasource.fluent.interfaces import (\n    Batch,\n    DataAsset,\n    Datasource,\n    _DataAssetT,\n)\n\nif TYPE_CHECKING:\n    from typing_extensions import TypeAlias\n\n    from great_expectations.datasource.fluent.interfaces import BatchMetadata\n    from great_expectations.execution_engine import SparkDFExecutionEngine\n\n\nlogger = logging.getLogger(__name__)\n\n\n# this enables us to include dataframe in the json schema\n_SparkDataFrameT = TypeVar(\"_SparkDataFrameT\")\n\nSparkConfig: TypeAlias = Dict[\n    StrictStr, Union[StrictStr, StrictInt, StrictFloat, StrictBool]\n]\n\n\nclass SparkDatasourceError(Exception):\n    pass\n\n\nclass _SparkDatasource(Datasource):\n    # instance attributes\n    spark_config: Union[SparkConfig, None] = None\n    force_reuse_spark_context: bool = True\n\n    @staticmethod\n    def _update_asset_forward_refs(asset_type: Type[_DataAssetT]) -> None:\n        # Only update forward refs if pyspark types are available.\n        if pyspark:\n            asset_type.update_forward_refs()\n\n    # Abstract Methods\n    @property\n    def execution_engine_type(self) -> Type[SparkDFExecutionEngine]:\n        \"\"\"Return the SparkDFExecutionEngine unless the override is set\"\"\"\n        from great_expectations.execution_engine.sparkdf_execution_engine import (\n            SparkDFExecutionEngine,\n        )\n\n        return SparkDFExecutionEngine\n\n    def test_connection(self, test_assets: bool = True) -> None:\n        \"\"\"Test the connection for the _SparkDatasource.\n\n        Args:\n            test_assets: If assets have been passed to the _SparkDatasource,\n                         an attempt can be made to test them as well.\n\n        Raises:\n            TestConnectionError: If the connection test fails.\n        \"\"\"\n        raise NotImplementedError(\n            \"\"\"One needs to implement \"test_connection\" on a _SparkDatasource subclass.\"\"\"\n        )\n\n    # End Abstract Methods\n\n\nclass DataFrameAsset(DataAsset, Generic[_SparkDataFrameT]):\n    # instance attributes\n    type: Literal[\"dataframe\"] = \"dataframe\"\n    # TODO: <Alex>05/31/2023: Upon removal of deprecated \"dataframe\" argument to \"PandasDatasource.add_dataframe_asset()\", default can be deleted.</Alex>\n    dataframe: Optional[_SparkDataFrameT] = pydantic.Field(\n        default=None, exclude=True, repr=False\n    )\n\n    class Config:\n        extra = pydantic.Extra.forbid\n\n    @pydantic.validator(\"dataframe\")\n    def _validate_dataframe(cls, dataframe: DataFrame) -> DataFrame:\n        if not (DataFrame and isinstance(dataframe, DataFrame)):  # type: ignore[truthy-function]\n            raise ValueError(\"dataframe must be of type pyspark.sql.DataFrame\")\n\n        return dataframe\n\n    def test_connection(self) -> None:\n        ...\n\n    @property\n    def batch_request_options(self) -> tuple[str, ...]:\n        return tuple()\n\n    def _get_reader_method(self) -> str:\n        raise NotImplementedError(\n            \"\"\"Spark DataFrameAsset does not implement \"_get_reader_method()\" method, because DataFrame is already available.\"\"\"\n        )\n\n    def _get_reader_options_include(self) -> set[str]:\n        raise NotImplementedError(\n            \"\"\"Spark DataFrameAsset does not implement \"_get_reader_options_include()\" method, because DataFrame is already available.\"\"\"\n        )\n\n    @public_api\n    # TODO: <Alex>05/31/2023: Upon removal of deprecated \"dataframe\" argument to \"PandasDatasource.add_dataframe_asset()\", its validation code must be deleted.</Alex>\n    @new_argument(\n        argument_name=\"dataframe\",\n        message='The \"dataframe\" argument is no longer part of \"PandasDatasource.add_dataframe_asset()\" method call; instead, \"dataframe\" is the required argument to \"DataFrameAsset.build_batch_request()\" method.',\n        version=\"0.16.15\",\n    )\n    def build_batch_request(\n        self, dataframe: Optional[_SparkDataFrameT] = None\n    ) -> BatchRequest:\n        \"\"\"A batch request that can be used to obtain batches for this DataAsset.\n\n        Args:\n            dataframe: The Spark Dataframe containing the data for this DataFrame data asset.\n\n        Returns:\n            A BatchRequest object that can be used to obtain a batch list from a Datasource by calling the\n            get_batch_list_from_batch_request method.\n        \"\"\"\n        if dataframe is None:\n            df = self.dataframe\n        else:\n            df = dataframe\n\n        if df is None:\n            raise ValueError(\n                \"Cannot build batch request for dataframe asset without a dataframe\"\n            )\n\n        self.dataframe = df\n\n        return BatchRequest(\n            datasource_name=self.datasource.name,\n            data_asset_name=self.name,\n            options={},\n        )\n\n    def _validate_batch_request(self, batch_request: BatchRequest) -> None:\n        \"\"\"Validates the batch_request has the correct form.\n\n        Args:\n            batch_request: A batch request object to be validated.\n        \"\"\"\n        if not (\n            batch_request.datasource_name == self.datasource.name\n            and batch_request.data_asset_name == self.name\n            and not batch_request.options\n        ):\n            expect_batch_request_form = BatchRequest(\n                datasource_name=self.datasource.name,\n                data_asset_name=self.name,\n                options={},\n                batch_slice=batch_request._batch_slice_input,  # type: ignore[attr-defined]\n            )\n            raise gx_exceptions.InvalidBatchRequestError(\n                \"BatchRequest should have form:\\n\"\n                f\"{pf(expect_batch_request_form.dict())}\\n\"\n                f\"but actually has form:\\n{pf(batch_request.dict())}\\n\"\n            )\n\n    def get_batch_list_from_batch_request(\n        self, batch_request: BatchRequest\n    ) -> list[Batch]:\n        self._validate_batch_request(batch_request)\n\n        batch_spec = RuntimeDataBatchSpec(batch_data=self.dataframe)\n        execution_engine: SparkDFExecutionEngine = (\n            self.datasource.get_execution_engine()\n        )\n        data, markers = execution_engine.get_batch_data_and_markers(\n            batch_spec=batch_spec\n        )\n\n        # batch_definition (along with batch_spec and markers) is only here to satisfy a\n        # legacy constraint when computing usage statistics in a validator. We hope to remove\n        # it in the future.\n        # imports are done inline to prevent a circular dependency with core/batch.py\n        from great_expectations.core import IDDict\n        from great_expectations.core.batch import BatchDefinition\n\n        batch_definition = BatchDefinition(\n            datasource_name=self.datasource.name,\n            data_connector_name=_DATA_CONNECTOR_NAME,\n            data_asset_name=self.name,\n            batch_identifiers=IDDict(batch_request.options),\n            batch_spec_passthrough=None,\n        )\n\n        batch_metadata: BatchMetadata = self._get_batch_metadata_from_batch_request(\n            batch_request=batch_request\n        )\n\n        # Some pydantic annotations are postponed due to circular imports.\n        # Batch.update_forward_refs() will set the annotations before we\n        # instantiate the Batch class since we can import them in this scope.\n        Batch.update_forward_refs()\n\n        return [\n            Batch(\n                datasource=self.datasource,\n                data_asset=self,\n                batch_request=batch_request,\n                data=data,\n                metadata=batch_metadata,\n                legacy_batch_markers=markers,\n                legacy_batch_spec=batch_spec,\n                legacy_batch_definition=batch_definition,\n            )\n        ]\n\n\n@public_api\nclass SparkDatasource(_SparkDatasource):\n    # class attributes\n    asset_types: ClassVar[List[Type[DataAsset]]] = [DataFrameAsset]\n\n    # instance attributes\n    type: Literal[\"spark\"] = \"spark\"\n\n    assets: List[DataFrameAsset] = []  # type: ignore[assignment]\n\n    def test_connection(self, test_assets: bool = True) -> None:\n        ...\n\n    @public_api\n    @deprecated_argument(\n        argument_name=\"dataframe\",\n        message='The \"dataframe\" argument is no longer part of \"PandasDatasource.add_dataframe_asset()\" method call; instead, \"dataframe\" is the required argument to \"DataFrameAsset.build_batch_request()\" method.',\n        version=\"0.16.15\",\n    )\n    def add_dataframe_asset(\n        self,\n        name: str,\n        dataframe: Optional[_SparkDataFrameT] = None,\n        batch_metadata: Optional[BatchMetadata] = None,\n    ) -> DataFrameAsset:\n        \"\"\"Adds a Dataframe DataAsset to this SparkDatasource object.\n\n        Args:\n            name: The name of the DataFrame asset. This can be any arbitrary string.\n            dataframe: The Spark Dataframe containing the data for this DataFrame data asset.\n            batch_metadata: An arbitrary user defined dictionary with string keys which will get inherited by any\n                            batches created from the asset.\n\n        Returns:\n            The DataFameAsset that has been added to this datasource.\n        \"\"\"\n        asset: DataFrameAsset = DataFrameAsset(\n            name=name,\n            batch_metadata=batch_metadata or {},\n        )\n        asset.dataframe = dataframe\n        return self._add_asset(asset=asset)\n", "path": "great_expectations/datasource/fluent/spark_datasource.py"}], "after_files": [{"content": "from __future__ import annotations\n\nimport logging\nfrom pprint import pformat as pf\nfrom typing import (\n    TYPE_CHECKING,\n    ClassVar,\n    Dict,\n    Generic,\n    List,\n    Literal,\n    Optional,\n    Type,\n    TypeVar,\n    Union,\n)\n\nimport pydantic\nfrom pydantic import StrictBool, StrictFloat, StrictInt, StrictStr\n\nimport great_expectations.exceptions as gx_exceptions\nfrom great_expectations.compatibility.pyspark import DataFrame, pyspark\nfrom great_expectations.core._docs_decorators import (\n    deprecated_argument,\n    new_argument,\n    public_api,\n)\nfrom great_expectations.core.batch_spec import RuntimeDataBatchSpec\nfrom great_expectations.datasource.fluent import BatchRequest\nfrom great_expectations.datasource.fluent.constants import (\n    _DATA_CONNECTOR_NAME,\n)\nfrom great_expectations.datasource.fluent.interfaces import (\n    Batch,\n    DataAsset,\n    Datasource,\n    _DataAssetT,\n)\n\nif TYPE_CHECKING:\n    from typing_extensions import TypeAlias\n\n    from great_expectations.datasource.fluent.interfaces import BatchMetadata\n    from great_expectations.execution_engine import SparkDFExecutionEngine\n\n\nlogger = logging.getLogger(__name__)\n\n\n# this enables us to include dataframe in the json schema\n_SparkDataFrameT = TypeVar(\"_SparkDataFrameT\")\n\nSparkConfig: TypeAlias = Dict[\n    StrictStr, Union[StrictStr, StrictInt, StrictFloat, StrictBool]\n]\n\n\nclass SparkDatasourceError(Exception):\n    pass\n\n\nclass _SparkDatasource(Datasource):\n    # instance attributes\n    spark_config: Union[SparkConfig, None] = None\n    force_reuse_spark_context: bool = True\n    persist: bool = True\n\n    @staticmethod\n    def _update_asset_forward_refs(asset_type: Type[_DataAssetT]) -> None:\n        # Only update forward refs if pyspark types are available.\n        if pyspark:\n            asset_type.update_forward_refs()\n\n    # Abstract Methods\n    @property\n    def execution_engine_type(self) -> Type[SparkDFExecutionEngine]:\n        \"\"\"Return the SparkDFExecutionEngine unless the override is set\"\"\"\n        from great_expectations.execution_engine.sparkdf_execution_engine import (\n            SparkDFExecutionEngine,\n        )\n\n        return SparkDFExecutionEngine\n\n    def test_connection(self, test_assets: bool = True) -> None:\n        \"\"\"Test the connection for the _SparkDatasource.\n\n        Args:\n            test_assets: If assets have been passed to the _SparkDatasource,\n                         an attempt can be made to test them as well.\n\n        Raises:\n            TestConnectionError: If the connection test fails.\n        \"\"\"\n        raise NotImplementedError(\n            \"\"\"One needs to implement \"test_connection\" on a _SparkDatasource subclass.\"\"\"\n        )\n\n    # End Abstract Methods\n\n\nclass DataFrameAsset(DataAsset, Generic[_SparkDataFrameT]):\n    # instance attributes\n    type: Literal[\"dataframe\"] = \"dataframe\"\n    # TODO: <Alex>05/31/2023: Upon removal of deprecated \"dataframe\" argument to \"PandasDatasource.add_dataframe_asset()\", default can be deleted.</Alex>\n    dataframe: Optional[_SparkDataFrameT] = pydantic.Field(\n        default=None, exclude=True, repr=False\n    )\n\n    class Config:\n        extra = pydantic.Extra.forbid\n\n    @pydantic.validator(\"dataframe\")\n    def _validate_dataframe(cls, dataframe: DataFrame) -> DataFrame:\n        if not (DataFrame and isinstance(dataframe, DataFrame)):  # type: ignore[truthy-function]\n            raise ValueError(\"dataframe must be of type pyspark.sql.DataFrame\")\n\n        return dataframe\n\n    def test_connection(self) -> None:\n        ...\n\n    @property\n    def batch_request_options(self) -> tuple[str, ...]:\n        return tuple()\n\n    def _get_reader_method(self) -> str:\n        raise NotImplementedError(\n            \"\"\"Spark DataFrameAsset does not implement \"_get_reader_method()\" method, because DataFrame is already available.\"\"\"\n        )\n\n    def _get_reader_options_include(self) -> set[str]:\n        raise NotImplementedError(\n            \"\"\"Spark DataFrameAsset does not implement \"_get_reader_options_include()\" method, because DataFrame is already available.\"\"\"\n        )\n\n    @public_api\n    # TODO: <Alex>05/31/2023: Upon removal of deprecated \"dataframe\" argument to \"PandasDatasource.add_dataframe_asset()\", its validation code must be deleted.</Alex>\n    @new_argument(\n        argument_name=\"dataframe\",\n        message='The \"dataframe\" argument is no longer part of \"PandasDatasource.add_dataframe_asset()\" method call; instead, \"dataframe\" is the required argument to \"DataFrameAsset.build_batch_request()\" method.',\n        version=\"0.16.15\",\n    )\n    def build_batch_request(\n        self, dataframe: Optional[_SparkDataFrameT] = None\n    ) -> BatchRequest:\n        \"\"\"A batch request that can be used to obtain batches for this DataAsset.\n\n        Args:\n            dataframe: The Spark Dataframe containing the data for this DataFrame data asset.\n\n        Returns:\n            A BatchRequest object that can be used to obtain a batch list from a Datasource by calling the\n            get_batch_list_from_batch_request method.\n        \"\"\"\n        if dataframe is None:\n            df = self.dataframe\n        else:\n            df = dataframe\n\n        if df is None:\n            raise ValueError(\n                \"Cannot build batch request for dataframe asset without a dataframe\"\n            )\n\n        self.dataframe = df\n\n        return BatchRequest(\n            datasource_name=self.datasource.name,\n            data_asset_name=self.name,\n            options={},\n        )\n\n    def _validate_batch_request(self, batch_request: BatchRequest) -> None:\n        \"\"\"Validates the batch_request has the correct form.\n\n        Args:\n            batch_request: A batch request object to be validated.\n        \"\"\"\n        if not (\n            batch_request.datasource_name == self.datasource.name\n            and batch_request.data_asset_name == self.name\n            and not batch_request.options\n        ):\n            expect_batch_request_form = BatchRequest(\n                datasource_name=self.datasource.name,\n                data_asset_name=self.name,\n                options={},\n                batch_slice=batch_request._batch_slice_input,  # type: ignore[attr-defined]\n            )\n            raise gx_exceptions.InvalidBatchRequestError(\n                \"BatchRequest should have form:\\n\"\n                f\"{pf(expect_batch_request_form.dict())}\\n\"\n                f\"but actually has form:\\n{pf(batch_request.dict())}\\n\"\n            )\n\n    def get_batch_list_from_batch_request(\n        self, batch_request: BatchRequest\n    ) -> list[Batch]:\n        self._validate_batch_request(batch_request)\n\n        batch_spec = RuntimeDataBatchSpec(batch_data=self.dataframe)\n        execution_engine: SparkDFExecutionEngine = (\n            self.datasource.get_execution_engine()\n        )\n        data, markers = execution_engine.get_batch_data_and_markers(\n            batch_spec=batch_spec\n        )\n\n        # batch_definition (along with batch_spec and markers) is only here to satisfy a\n        # legacy constraint when computing usage statistics in a validator. We hope to remove\n        # it in the future.\n        # imports are done inline to prevent a circular dependency with core/batch.py\n        from great_expectations.core import IDDict\n        from great_expectations.core.batch import BatchDefinition\n\n        batch_definition = BatchDefinition(\n            datasource_name=self.datasource.name,\n            data_connector_name=_DATA_CONNECTOR_NAME,\n            data_asset_name=self.name,\n            batch_identifiers=IDDict(batch_request.options),\n            batch_spec_passthrough=None,\n        )\n\n        batch_metadata: BatchMetadata = self._get_batch_metadata_from_batch_request(\n            batch_request=batch_request\n        )\n\n        # Some pydantic annotations are postponed due to circular imports.\n        # Batch.update_forward_refs() will set the annotations before we\n        # instantiate the Batch class since we can import them in this scope.\n        Batch.update_forward_refs()\n\n        return [\n            Batch(\n                datasource=self.datasource,\n                data_asset=self,\n                batch_request=batch_request,\n                data=data,\n                metadata=batch_metadata,\n                legacy_batch_markers=markers,\n                legacy_batch_spec=batch_spec,\n                legacy_batch_definition=batch_definition,\n            )\n        ]\n\n\n@public_api\nclass SparkDatasource(_SparkDatasource):\n    # class attributes\n    asset_types: ClassVar[List[Type[DataAsset]]] = [DataFrameAsset]\n\n    # instance attributes\n    type: Literal[\"spark\"] = \"spark\"\n\n    assets: List[DataFrameAsset] = []  # type: ignore[assignment]\n\n    def test_connection(self, test_assets: bool = True) -> None:\n        ...\n\n    @public_api\n    @deprecated_argument(\n        argument_name=\"dataframe\",\n        message='The \"dataframe\" argument is no longer part of \"PandasDatasource.add_dataframe_asset()\" method call; instead, \"dataframe\" is the required argument to \"DataFrameAsset.build_batch_request()\" method.',\n        version=\"0.16.15\",\n    )\n    def add_dataframe_asset(\n        self,\n        name: str,\n        dataframe: Optional[_SparkDataFrameT] = None,\n        batch_metadata: Optional[BatchMetadata] = None,\n    ) -> DataFrameAsset:\n        \"\"\"Adds a Dataframe DataAsset to this SparkDatasource object.\n\n        Args:\n            name: The name of the DataFrame asset. This can be any arbitrary string.\n            dataframe: The Spark Dataframe containing the data for this DataFrame data asset.\n            batch_metadata: An arbitrary user defined dictionary with string keys which will get inherited by any\n                            batches created from the asset.\n\n        Returns:\n            The DataFameAsset that has been added to this datasource.\n        \"\"\"\n        asset: DataFrameAsset = DataFrameAsset(\n            name=name,\n            batch_metadata=batch_metadata or {},\n        )\n        asset.dataframe = dataframe\n        return self._add_asset(asset=asset)\n", "path": "great_expectations/datasource/fluent/spark_datasource.py"}]}