from __future__ import annotations

import logging
from pprint import pformat as pf
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Dict,
    Generic,
    List,
    Literal,
    Optional,
    Type,
    TypeVar,
    Union,
)

import pydantic
from pydantic import StrictBool, StrictFloat, StrictInt, StrictStr

import great_expectations.exceptions as gx_exceptions
from great_expectations.compatibility.pyspark import DataFrame, pyspark
from great_expectations.core._docs_decorators import (
    deprecated_argument,
    new_argument,
    public_api,
)
from great_expectations.core.batch_spec import RuntimeDataBatchSpec
from great_expectations.datasource.fluent import BatchRequest
from great_expectations.datasource.fluent.constants import (
    _DATA_CONNECTOR_NAME,
)
from great_expectations.datasource.fluent.interfaces import (
    Batch,
    DataAsset,
    Datasource,
    _DataAssetT,
)

if TYPE_CHECKING:
    from typing_extensions import TypeAlias

    from great_expectations.datasource.fluent.interfaces import BatchMetadata
    from great_expectations.execution_engine import SparkDFExecutionEngine


logger = logging.getLogger(__name__)


# this enables us to include dataframe in the json schema
_SparkDataFrameT = TypeVar("_SparkDataFrameT")

SparkConfig: TypeAlias = Dict[
    StrictStr, Union[StrictStr, StrictInt, StrictFloat, StrictBool]
]


class SparkDatasourceError(Exception):
    pass


class _SparkDatasource(Datasource):
    # instance attributes
    spark_config: Union[SparkConfig, None] = None
    force_reuse_spark_context: bool = True

    @staticmethod
    def _update_asset_forward_refs(asset_type: Type[_DataAssetT]) -> None:
        # Only update forward refs if pyspark types are available.
        if pyspark:
            asset_type.update_forward_refs()

    # Abstract Methods
    @property
    def execution_engine_type(self) -> Type[SparkDFExecutionEngine]:
        """Return the SparkDFExecutionEngine unless the override is set"""
        from great_expectations.execution_engine.sparkdf_execution_engine import (
            SparkDFExecutionEngine,
        )

        return SparkDFExecutionEngine

    def test_connection(self, test_assets: bool = True) -> None:
        """Test the connection for the _SparkDatasource.

        Args:
            test_assets: If assets have been passed to the _SparkDatasource,
                         an attempt can be made to test them as well.

        Raises:
            TestConnectionError: If the connection test fails.
        """
        raise NotImplementedError(
            """One needs to implement "test_connection" on a _SparkDatasource subclass."""
        )

    # End Abstract Methods


class DataFrameAsset(DataAsset, Generic[_SparkDataFrameT]):
    # instance attributes
    type: Literal["dataframe"] = "dataframe"
    # TODO: <Alex>05/31/2023: Upon removal of deprecated "dataframe" argument to "PandasDatasource.add_dataframe_asset()", default can be deleted.</Alex>
    dataframe: Optional[_SparkDataFrameT] = pydantic.Field(
        default=None, exclude=True, repr=False
    )

    class Config:
        extra = pydantic.Extra.forbid

    @pydantic.validator("dataframe")
    def _validate_dataframe(cls, dataframe: DataFrame) -> DataFrame:
        if not (DataFrame and isinstance(dataframe, DataFrame)):  # type: ignore[truthy-function]
            raise ValueError("dataframe must be of type pyspark.sql.DataFrame")

        return dataframe

    def test_connection(self) -> None:
        ...

    @property
    def batch_request_options(self) -> tuple[str, ...]:
        return tuple()

    def _get_reader_method(self) -> str:
        raise NotImplementedError(
            """Spark DataFrameAsset does not implement "_get_reader_method()" method, because DataFrame is already available."""
        )

    def _get_reader_options_include(self) -> set[str]:
        raise NotImplementedError(
            """Spark DataFrameAsset does not implement "_get_reader_options_include()" method, because DataFrame is already available."""
        )

    @public_api
    # TODO: <Alex>05/31/2023: Upon removal of deprecated "dataframe" argument to "PandasDatasource.add_dataframe_asset()", its validation code must be deleted.</Alex>
    @new_argument(
        argument_name="dataframe",
        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.',
        version="0.16.15",
    )
    def build_batch_request(
        self, dataframe: Optional[_SparkDataFrameT] = None
    ) -> BatchRequest:
        """A batch request that can be used to obtain batches for this DataAsset.

        Args:
            dataframe: The Spark Dataframe containing the data for this DataFrame data asset.

        Returns:
            A BatchRequest object that can be used to obtain a batch list from a Datasource by calling the
            get_batch_list_from_batch_request method.
        """
        if dataframe is None:
            df = self.dataframe
        else:
            df = dataframe

        if df is None:
            raise ValueError(
                "Cannot build batch request for dataframe asset without a dataframe"
            )

        self.dataframe = df

        return BatchRequest(
            datasource_name=self.datasource.name,
            data_asset_name=self.name,
            options={},
        )

    def _validate_batch_request(self, batch_request: BatchRequest) -> None:
        """Validates the batch_request has the correct form.

        Args:
            batch_request: A batch request object to be validated.
        """
        if not (
            batch_request.datasource_name == self.datasource.name
            and batch_request.data_asset_name == self.name
            and not batch_request.options
        ):
            expect_batch_request_form = BatchRequest(
                datasource_name=self.datasource.name,
                data_asset_name=self.name,
                options={},
                batch_slice=batch_request._batch_slice_input,  # type: ignore[attr-defined]
            )
            raise gx_exceptions.InvalidBatchRequestError(
                "BatchRequest should have form:\n"
                f"{pf(expect_batch_request_form.dict())}\n"
                f"but actually has form:\n{pf(batch_request.dict())}\n"
            )

    def get_batch_list_from_batch_request(
        self, batch_request: BatchRequest
    ) -> list[Batch]:
        self._validate_batch_request(batch_request)

        batch_spec = RuntimeDataBatchSpec(batch_data=self.dataframe)
        execution_engine: SparkDFExecutionEngine = (
            self.datasource.get_execution_engine()
        )
        data, markers = execution_engine.get_batch_data_and_markers(
            batch_spec=batch_spec
        )

        # batch_definition (along with batch_spec and markers) is only here to satisfy a
        # legacy constraint when computing usage statistics in a validator. We hope to remove
        # it in the future.
        # imports are done inline to prevent a circular dependency with core/batch.py
        from great_expectations.core import IDDict
        from great_expectations.core.batch import BatchDefinition

        batch_definition = BatchDefinition(
            datasource_name=self.datasource.name,
            data_connector_name=_DATA_CONNECTOR_NAME,
            data_asset_name=self.name,
            batch_identifiers=IDDict(batch_request.options),
            batch_spec_passthrough=None,
        )

        batch_metadata: BatchMetadata = self._get_batch_metadata_from_batch_request(
            batch_request=batch_request
        )

        # Some pydantic annotations are postponed due to circular imports.
        # Batch.update_forward_refs() will set the annotations before we
        # instantiate the Batch class since we can import them in this scope.
        Batch.update_forward_refs()

        return [
            Batch(
                datasource=self.datasource,
                data_asset=self,
                batch_request=batch_request,
                data=data,
                metadata=batch_metadata,
                legacy_batch_markers=markers,
                legacy_batch_spec=batch_spec,
                legacy_batch_definition=batch_definition,
            )
        ]


@public_api
class SparkDatasource(_SparkDatasource):
    # class attributes
    asset_types: ClassVar[List[Type[DataAsset]]] = [DataFrameAsset]

    # instance attributes
    type: Literal["spark"] = "spark"

    assets: List[DataFrameAsset] = []  # type: ignore[assignment]

    def test_connection(self, test_assets: bool = True) -> None:
        ...

    @public_api
    @deprecated_argument(
        argument_name="dataframe",
        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.',
        version="0.16.15",
    )
    def add_dataframe_asset(
        self,
        name: str,
        dataframe: Optional[_SparkDataFrameT] = None,
        batch_metadata: Optional[BatchMetadata] = None,
    ) -> DataFrameAsset:
        """Adds a Dataframe DataAsset to this SparkDatasource object.

        Args:
            name: The name of the DataFrame asset. This can be any arbitrary string.
            dataframe: The Spark Dataframe containing the data for this DataFrame data asset.
            batch_metadata: An arbitrary user defined dictionary with string keys which will get inherited by any
                            batches created from the asset.

        Returns:
            The DataFameAsset that has been added to this datasource.
        """
        asset: DataFrameAsset = DataFrameAsset(
            name=name,
            batch_metadata=batch_metadata or {},
        )
        asset.dataframe = dataframe
        return self._add_asset(asset=asset)

great_expectations/datasource/fluent/spark_datasource.py
from __future__ import annotations

import logging
from pprint import pformat as pf
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Dict,
    Generic,
    List,
    Literal,
    Optional,
    Type,
    TypeVar,
    Union,
)

import pydantic
from pydantic import StrictBool, StrictFloat, StrictInt, StrictStr

import great_expectations.exceptions as gx_exceptions
from great_expectations.compatibility.pyspark import DataFrame, pyspark
from great_expectations.core._docs_decorators import (
    deprecated_argument,
    new_argument,
    public_api,
)
from great_expectations.core.batch_spec import RuntimeDataBatchSpec
from great_expectations.datasource.fluent import BatchRequest
from great_expectations.datasource.fluent.constants import (
    _DATA_CONNECTOR_NAME,
)
from great_expectations.datasource.fluent.interfaces import (
    Batch,
    DataAsset,
    Datasource,
    _DataAssetT,
)

if TYPE_CHECKING:
    from typing_extensions import TypeAlias

    from great_expectations.datasource.fluent.interfaces import BatchMetadata
    from great_expectations.execution_engine import SparkDFExecutionEngine


logger = logging.getLogger(__name__)


# this enables us to include dataframe in the json schema
_SparkDataFrameT = TypeVar("_SparkDataFrameT")

SparkConfig: TypeAlias = Dict[
    StrictStr, Union[StrictStr, StrictInt, StrictFloat, StrictBool]
]


class SparkDatasourceError(Exception):
    pass


class _SparkDatasource(Datasource):
    # 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:
        # Only update forward refs if pyspark types are available.
        if pyspark:
            asset_type.update_forward_refs()

    # Abstract Methods
    @property
    def execution_engine_type(self) -> Type[SparkDFExecutionEngine]:
        """Return the SparkDFExecutionEngine unless the override is set"""
        from great_expectations.execution_engine.sparkdf_execution_engine import (
            SparkDFExecutionEngine,
        )

        return SparkDFExecutionEngine

    def test_connection(self, test_assets: bool = True) -> None:
        """Test the connection for the _SparkDatasource.

        Args:
            test_assets: If assets have been passed to the _SparkDatasource,
                         an attempt can be made to test them as well.

        Raises:
            TestConnectionError: If the connection test fails.
        """
        raise NotImplementedError(
            """One needs to implement "test_connection" on a _SparkDatasource subclass."""
        )

    # End Abstract Methods


class DataFrameAsset(DataAsset, Generic[_SparkDataFrameT]):
    # instance attributes
    type: Literal["dataframe"] = "dataframe"
    # TODO: <Alex>05/31/2023: Upon removal of deprecated "dataframe" argument to "PandasDatasource.add_dataframe_asset()", default can be deleted.</Alex>
    dataframe: Optional[_SparkDataFrameT] = pydantic.Field(
        default=None, exclude=True, repr=False
    )

    class Config:
        extra = pydantic.Extra.forbid

    @pydantic.validator("dataframe")
    def _validate_dataframe(cls, dataframe: DataFrame) -> DataFrame:
        if not (DataFrame and isinstance(dataframe, DataFrame)):  # type: ignore[truthy-function]
            raise ValueError("dataframe must be of type pyspark.sql.DataFrame")

        return dataframe

    def test_connection(self) -> None:
        ...

    @property
    def batch_request_options(self) -> tuple[str, ...]:
        return tuple()

    def _get_reader_method(self) -> str:
        raise NotImplementedError(
            """Spark DataFrameAsset does not implement "_get_reader_method()" method, because DataFrame is already available."""
        )

    def _get_reader_options_include(self) -> set[str]:
        raise NotImplementedError(
            """Spark DataFrameAsset does not implement "_get_reader_options_include()" method, because DataFrame is already available."""
        )

    @public_api
    # TODO: <Alex>05/31/2023: Upon removal of deprecated "dataframe" argument to "PandasDatasource.add_dataframe_asset()", its validation code must be deleted.</Alex>
    @new_argument(
        argument_name="dataframe",
        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.',
        version="0.16.15",
    )
    def build_batch_request(
        self, dataframe: Optional[_SparkDataFrameT] = None
    ) -> BatchRequest:
        """A batch request that can be used to obtain batches for this DataAsset.

        Args:
            dataframe: The Spark Dataframe containing the data for this DataFrame data asset.

        Returns:
            A BatchRequest object that can be used to obtain a batch list from a Datasource by calling the
            get_batch_list_from_batch_request method.
        """
        if dataframe is None:
            df = self.dataframe
        else:
            df = dataframe

        if df is None:
            raise ValueError(
                "Cannot build batch request for dataframe asset without a dataframe"
            )

        self.dataframe = df

        return BatchRequest(
            datasource_name=self.datasource.name,
            data_asset_name=self.name,
            options={},
        )

    def _validate_batch_request(self, batch_request: BatchRequest) -> None:
        """Validates the batch_request has the correct form.

        Args:
            batch_request: A batch request object to be validated.
        """
        if not (
            batch_request.datasource_name == self.datasource.name
            and batch_request.data_asset_name == self.name
            and not batch_request.options
        ):
            expect_batch_request_form = BatchRequest(
                datasource_name=self.datasource.name,
                data_asset_name=self.name,
                options={},
                batch_slice=batch_request._batch_slice_input,  # type: ignore[attr-defined]
            )
            raise gx_exceptions.InvalidBatchRequestError(
                "BatchRequest should have form:\n"
                f"{pf(expect_batch_request_form.dict())}\n"
                f"but actually has form:\n{pf(batch_request.dict())}\n"
            )

    def get_batch_list_from_batch_request(
        self, batch_request: BatchRequest
    ) -> list[Batch]:
        self._validate_batch_request(batch_request)

        batch_spec = RuntimeDataBatchSpec(batch_data=self.dataframe)
        execution_engine: SparkDFExecutionEngine = (
            self.datasource.get_execution_engine()
        )
        data, markers = execution_engine.get_batch_data_and_markers(
            batch_spec=batch_spec
        )

        # batch_definition (along with batch_spec and markers) is only here to satisfy a
        # legacy constraint when computing usage statistics in a validator. We hope to remove
        # it in the future.
        # imports are done inline to prevent a circular dependency with core/batch.py
        from great_expectations.core import IDDict
        from great_expectations.core.batch import BatchDefinition

        batch_definition = BatchDefinition(
            datasource_name=self.datasource.name,
            data_connector_name=_DATA_CONNECTOR_NAME,
            data_asset_name=self.name,
            batch_identifiers=IDDict(batch_request.options),
            batch_spec_passthrough=None,
        )

        batch_metadata: BatchMetadata = self._get_batch_metadata_from_batch_request(
            batch_request=batch_request
        )

        # Some pydantic annotations are postponed due to circular imports.
        # Batch.update_forward_refs() will set the annotations before we
        # instantiate the Batch class since we can import them in this scope.
        Batch.update_forward_refs()

        return [
            Batch(
                datasource=self.datasource,
                data_asset=self,
                batch_request=batch_request,
                data=data,
                metadata=batch_metadata,
                legacy_batch_markers=markers,
                legacy_batch_spec=batch_spec,
                legacy_batch_definition=batch_definition,
            )
        ]


@public_api
class SparkDatasource(_SparkDatasource):
    # class attributes
    asset_types: ClassVar[List[Type[DataAsset]]] = [DataFrameAsset]

    # instance attributes
    type: Literal["spark"] = "spark"

    assets: List[DataFrameAsset] = []  # type: ignore[assignment]

    def test_connection(self, test_assets: bool = True) -> None:
        ...

    @public_api
    @deprecated_argument(
        argument_name="dataframe",
        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.',
        version="0.16.15",
    )
    def add_dataframe_asset(
        self,
        name: str,
        dataframe: Optional[_SparkDataFrameT] = None,
        batch_metadata: Optional[BatchMetadata] = None,
    ) -> DataFrameAsset:
        """Adds a Dataframe DataAsset to this SparkDatasource object.

        Args:
            name: The name of the DataFrame asset. This can be any arbitrary string.
            dataframe: The Spark Dataframe containing the data for this DataFrame data asset.
            batch_metadata: An arbitrary user defined dictionary with string keys which will get inherited by any
                            batches created from the asset.

        Returns:
            The DataFameAsset that has been added to this datasource.
        """
        asset: DataFrameAsset = DataFrameAsset(
            name=name,
            batch_metadata=batch_metadata or {},
        )
        asset.dataframe = dataframe
        return self._add_asset(asset=asset)

great_expectations/datasource/fluent/spark_datasource.py
diff --git a/great_expectations/datasource/fluent/schemas/SparkAzureBlobStorageDatasource.json b/great_expectations/datasource/fluent/schemas/SparkAzureBlobStorageDatasource.json
index 803ca39437a0..4a08c4eba4ab 100644
--- a/great_expectations/datasource/fluent/schemas/SparkAzureBlobStorageDatasource.json
+++ b/great_expectations/datasource/fluent/schemas/SparkAzureBlobStorageDatasource.json
@@ -91,6 +91,11 @@
             "default": true,
             "type": "boolean"
         },
+        "persist": {
+            "title": "Persist",
+            "default": true,
+            "type": "boolean"
+        },
         "azure_options": {
             "title": "Azure Options",
             "default": {},
diff --git a/great_expectations/datasource/fluent/schemas/SparkDBFSDatasource.json b/great_expectations/datasource/fluent/schemas/SparkDBFSDatasource.json
index aeef5affe279..053e8de1106e 100644
--- a/great_expectations/datasource/fluent/schemas/SparkDBFSDatasource.json
+++ b/great_expectations/datasource/fluent/schemas/SparkDBFSDatasource.json
@@ -91,6 +91,11 @@
             "default": true,
             "type": "boolean"
         },
+        "persist": {
+            "title": "Persist",
+            "default": true,
+            "type": "boolean"
+        },
         "base_directory": {
             "title": "Base Directory",
             "type": "string",
diff --git a/great_expectations/datasource/fluent/schemas/SparkDatasource.json b/great_expectations/datasource/fluent/schemas/SparkDatasource.json
index 8197a10dea0a..49c8b437121a 100644
--- a/great_expectations/datasource/fluent/schemas/SparkDatasource.json
+++ b/great_expectations/datasource/fluent/schemas/SparkDatasource.json
@@ -53,6 +53,11 @@
             "title": "Force Reuse Spark Context",
             "default": true,
             "type": "boolean"
+        },
+        "persist": {
+            "title": "Persist",
+            "default": true,
+            "type": "boolean"
         }
     },
     "required": [
diff --git a/great_expectations/datasource/fluent/schemas/SparkFilesystemDatasource.json b/great_expectations/datasource/fluent/schemas/SparkFilesystemDatasource.json
index 2a6fb524e683..f8076b41fda6 100644
--- a/great_expectations/datasource/fluent/schemas/SparkFilesystemDatasource.json
+++ b/great_expectations/datasource/fluent/schemas/SparkFilesystemDatasource.json
@@ -91,6 +91,11 @@
             "default": true,
             "type": "boolean"
         },
+        "persist": {
+            "title": "Persist",
+            "default": true,
+            "type": "boolean"
+        },
         "base_directory": {
             "title": "Base Directory",
             "type": "string",
diff --git a/great_expectations/datasource/fluent/schemas/SparkGoogleCloudStorageDatasource.json b/great_expectations/datasource/fluent/schemas/SparkGoogleCloudStorageDatasource.json
index 1cdfe76965a8..e959279b9547 100644
--- a/great_expectations/datasource/fluent/schemas/SparkGoogleCloudStorageDatasource.json
+++ b/great_expectations/datasource/fluent/schemas/SparkGoogleCloudStorageDatasource.json
@@ -91,6 +91,11 @@
             "default": true,
             "type": "boolean"
         },
+        "persist": {
+            "title": "Persist",
+            "default": true,
+            "type": "boolean"
+        },
         "bucket_or_name": {
             "title": "Bucket Or Name",
             "type": "string"
diff --git a/great_expectations/datasource/fluent/schemas/SparkS3Datasource.json b/great_expectations/datasource/fluent/schemas/SparkS3Datasource.json
index b10e5b15fe46..e6c1d4fee9f4 100644
--- a/great_expectations/datasource/fluent/schemas/SparkS3Datasource.json
+++ b/great_expectations/datasource/fluent/schemas/SparkS3Datasource.json
@@ -91,6 +91,11 @@
             "default": true,
             "type": "boolean"
         },
+        "persist": {
+            "title": "Persist",
+            "default": true,
+            "type": "boolean"
+        },
         "bucket": {
             "title": "Bucket",
             "type": "string"
diff --git a/great_expectations/datasource/fluent/sources.pyi b/great_expectations/datasource/fluent/sources.pyi
index c4bece62fcea..ee7885903382 100644
--- a/great_expectations/datasource/fluent/sources.pyi
+++ b/great_expectations/datasource/fluent/sources.pyi
@@ -320,6 +320,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
     ) -> SparkDatasource: ...
     def update_spark(
         self,
@@ -327,6 +328,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
     ) -> SparkDatasource: ...
     def add_or_update_spark(
         self,
@@ -334,6 +336,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
     ) -> SparkDatasource: ...
     def delete_spark(
         self,
@@ -347,6 +350,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         base_directory: pathlib.Path = ...,
         data_context_root_directory: Union[pathlib.Path, None] = ...,
     ) -> SparkFilesystemDatasource: ...
@@ -358,6 +362,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         base_directory: pathlib.Path = ...,
         data_context_root_directory: Union[pathlib.Path, None] = ...,
     ) -> SparkFilesystemDatasource: ...
@@ -369,6 +374,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         base_directory: pathlib.Path = ...,
         data_context_root_directory: Union[pathlib.Path, None] = ...,
     ) -> SparkFilesystemDatasource: ...
@@ -384,6 +390,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         base_directory: pathlib.Path = ...,
         data_context_root_directory: Union[pathlib.Path, None] = ...,
     ) -> SparkDBFSDatasource: ...
@@ -395,6 +402,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         base_directory: pathlib.Path = ...,
         data_context_root_directory: Union[pathlib.Path, None] = ...,
     ) -> SparkDBFSDatasource: ...
@@ -406,6 +414,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         base_directory: pathlib.Path = ...,
         data_context_root_directory: Union[pathlib.Path, None] = ...,
     ) -> SparkDBFSDatasource: ...
@@ -421,6 +430,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         bucket: str = ...,
         boto3_options: dict[str, Union[ConfigStr, Any]] = ...,
     ) -> SparkS3Datasource: ...
@@ -432,6 +442,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         bucket: str = ...,
         boto3_options: dict[str, Union[ConfigStr, Any]] = ...,
     ) -> SparkS3Datasource: ...
@@ -443,6 +454,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         bucket: str = ...,
         boto3_options: dict[str, Union[ConfigStr, Any]] = ...,
     ) -> SparkS3Datasource: ...
@@ -458,6 +470,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         bucket_or_name: str = ...,
         gcs_options: dict[str, Union[ConfigStr, Any]] = ...,
     ) -> SparkGoogleCloudStorageDatasource: ...
@@ -469,6 +482,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         bucket_or_name: str = ...,
         gcs_options: dict[str, Union[ConfigStr, Any]] = ...,
     ) -> SparkGoogleCloudStorageDatasource: ...
@@ -480,6 +494,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         bucket_or_name: str = ...,
         gcs_options: dict[str, Union[ConfigStr, Any]] = ...,
     ) -> SparkGoogleCloudStorageDatasource: ...
@@ -495,6 +510,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         azure_options: dict[str, Any] = ...,
     ) -> SparkAzureBlobStorageDatasource: ...
     def update_spark_abs(  # noqa: PLR0913
@@ -505,6 +521,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         azure_options: dict[str, Any] = ...,
     ) -> SparkAzureBlobStorageDatasource: ...
     def add_or_update_spark_abs(  # noqa: PLR0913
@@ -515,6 +532,7 @@ class _SourceFactories:
         *,
         spark_config: SparkConfig | None = None,
         force_reuse_spark_context: bool = True,
+        persist: bool = True,
         azure_options: dict[str, Any] = ...,
     ) -> SparkAzureBlobStorageDatasource: ...
     def delete_spark_abs(
diff --git a/great_expectations/datasource/fluent/spark_datasource.py b/great_expectations/datasource/fluent/spark_datasource.py
index 46ae5f4eb8c6..6f84685a8991 100644
--- a/great_expectations/datasource/fluent/spark_datasource.py
+++ b/great_expectations/datasource/fluent/spark_datasource.py
@@ -63,6 +63,7 @@ class _SparkDatasource(Datasource):
     # 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:
diff --git a/tests/datasource/fluent/test_spark_datasource.py b/tests/datasource/fluent/test_spark_datasource.py
index a1029998af8a..71a7e7baae45 100644
--- a/tests/datasource/fluent/test_spark_datasource.py
+++ b/tests/datasource/fluent/test_spark_datasource.py
@@ -112,8 +112,10 @@ def test_spark_data_asset_batch_metadata(
 
 
 @pytest.mark.spark
+@pytest.mark.parametrize("persist", [True, False])
 def test_spark_config_passed_to_execution_engine(
     empty_data_context: AbstractDataContext,
+    persist,
     spark_session,
 ):
     spark_config: SparkConfig | None = {
@@ -128,6 +130,7 @@ def test_spark_config_passed_to_execution_engine(
         name="my_spark_datasource",
         spark_config=spark_config,
         force_reuse_spark_context=False,
+        persist=persist,
     )
     spark_config = json.loads(json.dumps(spark_config), parse_int=str, parse_float=str)
     assert is_candidate_subset_of_target(

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
