[FEATURE] DataAsset uses partitioner from BatchConfig

docs/docusaurus/docs/snippets/how_to_connect_to_a_sql_table.py

@@ -8,6 +8,7 @@
 import warnings
 
 import great_expectations as gx
+from great_expectations.core.partitioners import PartitionerColumnValue
 from great_expectations.datasource.fluent import GxDatasourceWarning
 
 sqlite_database_path = str(
@@ -76,11 +77,14 @@
 }
 
 # <snippet name="docs/docusaurus/docs/snippets/how_to_connect_to_a_sql_table.py add_vendor_id_splitter">
-table_asset.add_partitioner_column_value("vendor_id")
+partitioner = PartitionerColumnValue(column_name="vendor_id")
 # </snippet>
 
 # <snippet name="docs/docusaurus/docs/snippets/how_to_connect_to_a_sql_table.py build_vendor_id_batch_request">
-my_batch_request = my_asset.build_batch_request({"vendor_id": 1})
+my_batch_request = my_asset.build_batch_request(
+    options={"vendor_id": 1},
+    partitioner=partitioner,
+)
 # </snippet>
 
 batches = my_asset.get_batch_list_from_batch_request(my_batch_request)

docs/docusaurus/docs/snippets/organize_batches_in_sqlite_datasource.py

@@ -2,6 +2,7 @@
 import tempfile
 
 import great_expectations as gx
+from great_expectations.core.partitioners import PartitionerYearAndMonth
 
 temp_dir = tempfile.TemporaryDirectory()
 full_path_to_project_directory = pathlib.Path(temp_dir.name).resolve()
@@ -38,10 +39,10 @@
 
 # Python
 # <snippet name="docs/docusaurus/docs/snippets/organize_batches_in_sqlite_datasource.py add_splitter_year_and_month">
-my_table_asset.add_partitioner_year_and_month(column_name="pickup_datetime")
+partitioner = PartitionerYearAndMonth(column_name="pickup_datetime")
 # </snippet>
 
-my_batch_request = my_table_asset.build_batch_request()
+my_batch_request = my_table_asset.build_batch_request(partitioner=partitioner)
 batches = my_table_asset.get_batch_list_from_batch_request(my_batch_request)
 
 assert len(batches) == 12
@@ -51,11 +52,14 @@
 my_asset = my_table_asset.add_sorters(["+year", "-month"])
 # </snippet>
 
-assert my_asset.batch_request_options == ("year", "month")
+assert my_asset.get_batch_request_options_keys(partitioner=partitioner) == (
+    "year",
+    "month",
+)
 
 # Python
 # <snippet name="docs/docusaurus/docs/snippets/organize_batches_in_sqlite_datasource.py my_batch_list">
-my_batch_request = my_table_asset.build_batch_request()
+my_batch_request = my_table_asset.build_batch_request(partitioner=partitioner)
 batches = my_table_asset.get_batch_list_from_batch_request(my_batch_request)
 # </snippet>
 

docs/sphinx_api_docs_source/public_api_excludes.py

@@ -828,4 +828,30 @@
         name="get_or_create_spark_session",
         filepath=pathlib.Path("great_expectations/core/util.py"),
     ),
+    IncludeExcludeDefinition(
+        reason="This method does not need to be accessed by users, and will eventually be removed from docs.",
+        name="get_batch_request_options_keys",
+        filepath=pathlib.Path(
+            "great_expectations/datasource/fluent/file_path_data_asset.py"
+        ),
+    ),
+    IncludeExcludeDefinition(
+        reason="This method does not need to be accessed by users, and will eventually be removed from docs.",
+        name="get_batch_request_options_keys",
+        filepath=pathlib.Path(
+            "great_expectations/datasource/fluent/pandas_datasource.py"
+        ),
+    ),
+    IncludeExcludeDefinition(
+        reason="This method does not need to be accessed by users, and will eventually be removed from docs.",
+        name="get_batch_request_options_keys",
+        filepath=pathlib.Path(
+            "great_expectations/datasource/fluent/spark_datasource.py"
+        ),
+    ),
+    IncludeExcludeDefinition(
+        reason="This method does not need to be accessed by users, and will eventually be removed from docs.",
+        name="get_batch_request_options_keys",
+        filepath=pathlib.Path("great_expectations/datasource/fluent/sql_datasource.py"),
+    ),
 ]

great_expectations/core/partitioners.py

@@ -1,31 +1,37 @@
 from typing import List, Literal, Union
 
+from great_expectations._docs_decorators import public_api
 from great_expectations.compatibility import pydantic
 
 
+@public_api
 class PartitionerYear(pydantic.BaseModel):
     column_name: str
     method_name: Literal["partition_on_year"] = "partition_on_year"
 
 
+@public_api
 class PartitionerYearAndMonth(pydantic.BaseModel):
     column_name: str
     method_name: Literal["partition_on_year_and_month"] = "partition_on_year_and_month"
 
 
+@public_api
 class PartitionerYearAndMonthAndDay(pydantic.BaseModel):
     column_name: str
     method_name: Literal[
         "partition_on_year_and_month_and_day"
     ] = "partition_on_year_and_month_and_day"
 
 
+@public_api
 class PartitionerDatetimePart(pydantic.BaseModel):
     datetime_parts: List[str]
     column_name: str
     method_name: Literal["partition_on_date_parts"] = "partition_on_date_parts"
 
 
+@public_api
 class PartitionerDividedInteger(pydantic.BaseModel):
     divisor: int
     column_name: str
@@ -34,24 +40,36 @@ class PartitionerDividedInteger(pydantic.BaseModel):
     ] = "partition_on_divided_integer"
 
 
+@public_api
 class PartitionerModInteger(pydantic.BaseModel):
     mod: int
     column_name: str
     method_name: Literal["partition_on_mod_integer"] = "partition_on_mod_integer"
 
 
+@public_api
 class PartitionerColumnValue(pydantic.BaseModel):
     column_name: str
     method_name: Literal["partition_on_column_value"] = "partition_on_column_value"
 
 
+@public_api
 class PartitionerMultiColumnValue(pydantic.BaseModel):
     column_names: List[str]
     method_name: Literal[
         "partition_on_multi_column_values"
     ] = "partition_on_multi_column_values"
 
 
+@public_api
+class PartitionerConvertedDatetime(pydantic.BaseModel):
+    column_name: str
+    method_name: Literal[
+        "partition_on_converted_datetime"
+    ] = "partition_on_converted_datetime"
+    date_format_string: str
+
+
 Partitioner = Union[
     PartitionerColumnValue,
     PartitionerMultiColumnValue,
@@ -61,4 +79,5 @@ class PartitionerMultiColumnValue(pydantic.BaseModel):
     PartitionerYearAndMonth,
     PartitionerYearAndMonthAndDay,
     PartitionerDatetimePart,
+    PartitionerConvertedDatetime,
 ]

great_expectations/datasource/fluent/file_path_data_asset.py

@@ -22,6 +22,16 @@
 from great_expectations._docs_decorators import public_api
 from great_expectations.compatibility import pydantic
 from great_expectations.compatibility.typing_extensions import override
+from great_expectations.core.partitioners import (
+    PartitionerColumnValue,
+    PartitionerDatetimePart,
+    PartitionerDividedInteger,
+    PartitionerModInteger,
+    PartitionerMultiColumnValue,
+    PartitionerYear,
+    PartitionerYearAndMonth,
+    PartitionerYearAndMonthAndDay,
+)
 from great_expectations.datasource.fluent.batch_request import (
     BatchRequest,
     BatchRequestOptions,
@@ -113,6 +123,20 @@ class _FilePathDataAsset(DataAsset):
     _test_connection_error_message: str = pydantic.PrivateAttr(
         "Could not connect to your asset"
     )
+    _partitioner_implementation_map: dict[
+        type[Partitioner], type[SparkPartitioner]
+    ] = pydantic.PrivateAttr(
+        default={
+            PartitionerYear: SparkPartitionerYear,
+            PartitionerYearAndMonth: SparkPartitionerYearAndMonth,
+            PartitionerYearAndMonthAndDay: SparkPartitionerYearAndMonthAndDay,
+            PartitionerColumnValue: SparkPartitionerColumnValue,
+            PartitionerDatetimePart: SparkPartitionerDatetimePart,
+            PartitionerDividedInteger: SparkPartitionerDividedInteger,
+            PartitionerModInteger: SparkPartitionerModInteger,
+            PartitionerMultiColumnValue: SparkPartitionerMultiColumnValue,
+        }
+    )
 
     class Config:
         """
@@ -140,6 +164,18 @@ def __init__(self, **data):
         )
         self._all_group_names = self._regex_parser.get_all_group_names()
 
+    def get_partitioner_implementation(
+        self, abstract_partitioner: Partitioner
+    ) -> SparkPartitioner:
+        PartitionerClass = self._partitioner_implementation_map.get(
+            type(abstract_partitioner)
+        )
+        if PartitionerClass is None:
+            raise ValueError(
+                f"Requested Partitioner `{abstract_partitioner.method_name}` is not implemented for this DataAsset. "
+            )
+        return PartitionerClass(**abstract_partitioner.dict())
+
     @property
     @override
     def batch_request_options(
@@ -167,6 +203,18 @@ def batch_request_options(
             + partitioner_options
         )
 
+    @override
+    def get_batch_request_options_keys(
+        self, partitioner: Optional[Partitioner]
+    ) -> tuple[str, ...]:
+        option_keys: tuple[str, ...] = tuple(self._all_group_names) + (
+            FILE_PATH_BATCH_SPEC_KEY,
+        )
+        if partitioner:
+            spark_partitioner = self.get_partitioner_implementation(partitioner)
+            option_keys += tuple(spark_partitioner.param_names)
+        return option_keys
+
     @public_api
     @override
     def build_batch_request(
@@ -206,8 +254,13 @@ def build_batch_request(
                         f"not a string: {value}"
                     )
 
-        if options is not None and not self._valid_batch_request_options(options):
-            allowed_keys = set(self.batch_request_options)
+        if options is not None and not self._batch_request_options_are_valid(
+            options=options,
+            partitioner=partitioner,
+        ):
+            allowed_keys = set(
+                self.get_batch_request_options_keys(partitioner=partitioner)
+            )
             actual_keys = set(options.keys())
             raise gx_exceptions.InvalidBatchRequestError(
                 "Batch request options should only contain keys from the following set:\n"
@@ -233,9 +286,14 @@ def _validate_batch_request(self, batch_request: BatchRequest) -> None:
         if not (
             batch_request.datasource_name == self.datasource.name
             and batch_request.data_asset_name == self.name
-            and self._valid_batch_request_options(batch_request.options)
+            and self._batch_request_options_are_valid(
+                options=batch_request.options, partitioner=batch_request.partitioner
+            )
         ):
-            options = {option: None for option in self.batch_request_options}
+            valid_options = self.get_batch_request_options_keys(
+                partitioner=batch_request.partitioner
+            )
+            options = {option: None for option in valid_options}
             expect_batch_request_form = BatchRequest(
                 datasource_name=self.datasource.name,
                 data_asset_name=self.name,
@@ -317,11 +375,17 @@ def _get_batch_definition_list(
         Returns:
             List of batch definitions.
         """
-        if self.partitioner:
+        if batch_request.partitioner:
+            spark_partitioner = self.get_partitioner_implementation(
+                batch_request.partitioner
+            )
             # Remove the partitioner kwargs from the batch_request to retrieve the batch and add them back later to the batch_spec.options
-            batch_request_options_counts = Counter(self.batch_request_options)
+            valid_options = self.get_batch_request_options_keys(
+                partitioner=batch_request.partitioner
+            )
+            batch_request_options_counts = Counter(valid_options)
             batch_request_copy_without_partitioner_kwargs = copy.deepcopy(batch_request)
-            for param_name in self.partitioner.param_names:
+            for param_name in spark_partitioner.param_names:
                 # If the option appears twice (e.g. from asset regex and from partitioner) then don't remove.
                 if batch_request_options_counts[param_name] == 1:
                     batch_request_copy_without_partitioner_kwargs.options.pop(
@@ -366,12 +430,15 @@ def _batch_spec_options_from_batch_request(
             ),
         }
 
-        if self.partitioner:
-            batch_spec_options["partitioner_method"] = self.partitioner.method_name
-            partitioner_kwargs = self.partitioner.partitioner_method_kwargs()
+        if batch_request.partitioner:
+            spark_partitioner = self.get_partitioner_implementation(
+                batch_request.partitioner
+            )
+            batch_spec_options["partitioner_method"] = spark_partitioner.method_name
+            partitioner_kwargs = spark_partitioner.partitioner_method_kwargs()
             partitioner_kwargs[
                 "batch_identifiers"
-            ] = self.partitioner.batch_request_options_to_batch_spec_kwarg_identifiers(
+            ] = spark_partitioner.batch_request_options_to_batch_spec_kwarg_identifiers(
                 batch_request.options
             )
             batch_spec_options["partitioner_kwargs"] = partitioner_kwargs

great_expectations/datasource/fluent/interfaces.py

@@ -228,6 +228,13 @@ def batch_request_options(self) -> tuple[str, ...]:
             """One needs to implement "batch_request_options" on a DataAsset subclass."""
         )
 
+    def get_batch_request_options_keys(
+        self, partitioner: Optional[Partitioner]
+    ) -> tuple[str, ...]:
+        raise NotImplementedError(
+            """One needs to implement "get_batch_request_options_keys" on a DataAsset subclass."""
+        )
+
     def build_batch_request(
         self,
         options: Optional[BatchRequestOptions] = None,
@@ -353,8 +360,11 @@ def get_batch_config(self, batch_config_name: str) -> BatchConfig:
             raise KeyError(f"Multiple keys for {batch_config_name} found")
         return batch_configs[0]
 
-    def _valid_batch_request_options(self, options: BatchRequestOptions) -> bool:
-        return set(options.keys()).issubset(set(self.batch_request_options))
+    def _batch_request_options_are_valid(
+        self, options: BatchRequestOptions, partitioner: Optional[Partitioner]
+    ) -> bool:
+        valid_options = self.get_batch_request_options_keys(partitioner=partitioner)
+        return set(options.keys()).issubset(set(valid_options))
 
     def _get_batch_metadata_from_batch_request(
         self, batch_request: BatchRequest

great_expectations/datasource/fluent/pandas_datasource.py

@@ -119,6 +119,10 @@ def test_connection(self) -> None:
     def batch_request_options(self) -> tuple[str, ...]:
         return tuple()
 
+    @override
+    def get_batch_request_options_keys(self, partitioner):
+        return tuple()
+
     @override
     def get_batch_list_from_batch_request(
         self, batch_request: BatchRequest