code changes for the partitiong support

Author: yogesh266

ads/feature_store/entity.py

@@ -281,16 +281,17 @@ def create(self, **kwargs) -> "Entity":
         return self
 
     def _build_feature_group(
-        self,
-        primary_keys,
-        input_feature_details,
-        expectation_suite: ExpectationSuite = None,
-        expectation_type: ExpectationType = ExpectationType.NO_EXPECTATION,
-        statistics_config: Union[StatisticsConfig, bool] = True,
-        transformation_id: str = None,
-        name: str = None,
-        description: str = None,
-        compartment_id: str = None,
+            self,
+            primary_keys,
+            partition_keys,
+            input_feature_details,
+            expectation_suite: ExpectationSuite = None,
+            expectation_type: ExpectationType = ExpectationType.NO_EXPECTATION,
+            statistics_config: Union[StatisticsConfig, bool] = True,
+            transformation_id: str = None,
+            name: str = None,
+            description: str = None,
+            compartment_id: str = None,
     ):
         feature_group_resource = (
             FeatureGroup()
@@ -302,6 +303,7 @@ def _build_feature_group(
             )
             .with_entity_id(self.id)
             .with_transformation_id(transformation_id)
+            .with_partition_keys(partition_keys)
             .with_primary_keys(primary_keys)
             .with_input_feature_details(input_feature_details)
             .with_statistics_config(statistics_config)
@@ -314,24 +316,27 @@ def _build_feature_group(
         return feature_group_resource
 
     def create_feature_group(
-        self,
-        primary_keys: List[str],
-        input_feature_details: List[FeatureDetail] = None,
-        schema_details_dataframe: Union[DataFrame, pd.DataFrame] = None,
-        expectation_suite: ExpectationSuite = None,
-        expectation_type: ExpectationType = ExpectationType.NO_EXPECTATION,
-        statistics_config: Union[StatisticsConfig, bool] = True,
-        transformation_id: str = None,
-        name: str = None,
-        description: str = None,
-        compartment_id: str = None,
+            self,
+            primary_keys: List[str],
+            partition_keys: List[str] = None,
+            input_feature_details: List[FeatureDetail] = None,
+            schema_details_dataframe: Union[DataFrame, pd.DataFrame] = None,
+            expectation_suite: ExpectationSuite = None,
+            expectation_type: ExpectationType = ExpectationType.NO_EXPECTATION,
+            statistics_config: Union[StatisticsConfig, bool] = True,
+            transformation_id: str = None,
+            name: str = None,
+            description: str = None,
+            compartment_id: str = None,
     ) -> "FeatureGroup":
         """Creates FeatureGroup  resource.
 
         Parameters
         ----------
         primary_keys: List[str]
             List of primary keys.
+        partition_keys: List[str]
+            List of partition_keys to partition the materialized data.
         input_feature_details: List[FeatureDetail]
             Raw feature schema for the input features.
         schema_details_dataframe: Union[DataFrame, pd.DataFrame]
@@ -377,6 +382,7 @@ def create_feature_group(
 
         self.oci_feature_group = self._build_feature_group(
             primary_keys,
+            partition_keys,
             raw_feature_details,
             expectation_suite,
             expectation_type,
@@ -405,7 +411,7 @@ def delete_feature_group(self):
 
     @classmethod
     def list_feature_group(
-        cls, compartment_id: str = None, **kwargs
+            cls, compartment_id: str = None, **kwargs
     ) -> List["FeatureGroup"]:
         """Lists FeatureGroup resources in a given compartment.
 
@@ -426,7 +432,7 @@ def list_feature_group(
 
     @classmethod
     def list_feature_group_df(
-        cls, compartment_id: str = None, **kwargs
+            cls, compartment_id: str = None, **kwargs
     ) -> "pandas.DataFrame":
         """Lists FeatureGroup resources in a given compartment as pandas dataframe.
 
@@ -446,14 +452,14 @@ def list_feature_group_df(
         return FeatureGroup.list_df(compartment_id, **kwargs)
 
     def _build_dataset(
-        self,
-        query: str,
-        name: str = None,
-        description: str = None,
-        compartment_id: str = None,
-        expectation_suite: ExpectationSuite = None,
-        expectation_type: ExpectationType = ExpectationType.NO_EXPECTATION,
-        statistics_config: Union[StatisticsConfig, bool] = True,
+            self,
+            query: str,
+            name: str = None,
+            description: str = None,
+            compartment_id: str = None,
+            expectation_suite: ExpectationSuite = None,
+            expectation_type: ExpectationType = ExpectationType.NO_EXPECTATION,
+            statistics_config: Union[StatisticsConfig, bool] = True,
     ):
         dataset_resource = (
             Dataset()
@@ -477,14 +483,14 @@ def _build_dataset(
         return dataset_resource
 
     def create_dataset(
-        self,
-        query: str,
-        name: str = None,
-        description: str = None,
-        compartment_id: str = None,
-        expectation_suite: ExpectationSuite = None,
-        expectation_type: ExpectationType = ExpectationType.NO_EXPECTATION,
-        statistics_config: Union[StatisticsConfig, bool] = True,
+            self,
+            query: str,
+            name: str = None,
+            description: str = None,
+            compartment_id: str = None,
+            expectation_suite: ExpectationSuite = None,
+            expectation_type: ExpectationType = ExpectationType.NO_EXPECTATION,
+            statistics_config: Union[StatisticsConfig, bool] = True,
     ) -> "Dataset":
         """Creates Dataset resource.
 
@@ -560,7 +566,7 @@ def list_dataset(cls, compartment_id: str = None, **kwargs) -> List["Dataset"]:
 
     @classmethod
     def list_dataset_df(
-        cls, compartment_id: str = None, **kwargs
+            cls, compartment_id: str = None, **kwargs
     ) -> "pandas.DataFrame":
         """Lists Dataset resources in a given compartment as pandas dataframe.
 

ads/feature_store/execution_strategy/delta_lake/delta_lake_service.py

@@ -37,6 +37,7 @@ def write_dataframe_to_delta_lake(
         dataflow_output,
         target_table_name,
         delta_table_primary_key,
+        partition_keys,
         ingestion_mode,
         raw_schema,
         feature_options=None,
@@ -47,6 +48,7 @@ def write_dataframe_to_delta_lake(
             dataflow_output (DataFrame): The data frame that needs to be written to the Delta table.
             target_table_name (str): The name of the target Delta table.
             delta_table_primary_key (List[dict]): The list of primary keys for the target Delta table.
+            partition_keys(List[dict]): The List of partition Keys.
             ingestion_mode (str): The ingestion mode for the data load.
             raw_schema (StructType): The schema of the raw data being ingested.
             feature_options (Dict[str, Union[str, int, float, bool]]): Optional. The dictionary containing feature options.
@@ -55,6 +57,8 @@ def write_dataframe_to_delta_lake(
             None.
         """
 
+        print(delta_table_primary_key, partition_keys, "Aaaya Kya")
+
         logger.info(f"target table name {target_table_name}")
         if (
             self.spark_engine.is_delta_table_exists(target_table_name)
@@ -109,7 +113,7 @@ def write_dataframe_to_delta_lake(
             logger.info(f"Upsert ops for target table {target_table_name} ended")
         else:
             self.save_delta_dataframe(
-                dataflow_output, ingestion_mode, target_table_name, feature_options
+                dataflow_output, ingestion_mode, target_table_name, feature_options, partition_keys
             )
 
     def __execute_delta_merge_insert_update(
@@ -193,7 +197,7 @@ def __execute_delta_merge_insert_update_all(
         )
 
     def save_delta_dataframe(
-        self, dataframe, dataframe_ingestion_mode, table_name, feature_options=None
+        self, dataframe, dataframe_ingestion_mode, table_name, feature_options=None, partition_keys=None
     ):
         """
         Saves a DataFrame to a Delta table with the specified options.
@@ -205,6 +209,12 @@ def save_delta_dataframe(
             feature_options (dict): Optional feature options to use when saving the DataFrame.
 
         """
+        delta_partition_keys = []
+
+        partition_keys_items = partition_keys["items"]
+        if partition_keys_items:
+            delta_partition_keys = [partition_key.get("name") for partition_key in partition_keys_items]
+
         if feature_options and feature_options.get("featureOptionWriteConfigDetails"):
             feature_delta_write_option_config = feature_options.get(
                 "featureOptionWriteConfigDetails"
@@ -216,9 +226,9 @@ def save_delta_dataframe(
 
             dataframe.write.format("delta").options(
                 **self.get_delta_write_config(feature_delta_write_option_config)
-            ).mode(dataframe_ingestion_mode).saveAsTable(table_name)
+            ).mode(dataframe_ingestion_mode).partitionBy(delta_partition_keys).saveAsTable(table_name)
         else:
-            dataframe.write.format("delta").mode(dataframe_ingestion_mode).saveAsTable(
+            dataframe.write.format("delta").mode(dataframe_ingestion_mode).partitionBy(delta_partition_keys).saveAsTable(
                 table_name
             )
 

ads/feature_store/execution_strategy/spark/spark_execution.py

@@ -260,6 +260,7 @@ def _save_offline_dataframe(
                 featured_data,
                 target_table,
                 feature_group.primary_keys,
+                feature_group.partition_keys,
                 feature_group_job.ingestion_mode,
                 featured_data.schema,
                 feature_group_job.feature_option_details,

ads/feature_store/feature_group.py

@@ -127,6 +127,7 @@ class FeatureGroup(Builder):
     CONST_ITEMS = "items"
     CONST_PRIMARY_KEY_NAME = "name"
     CONST_PRIMARY_KEYS = "primaryKeys"
+    CONST_PARTITION_KEYS = "partitionKeys"
     CONST_EXPECTATION_DETAILS = "expectationDetails"
     CONST_INPUT_FEATURE_DETAILS = "inputFeatureDetails"
     CONST_OUTPUT_FEATURE_DETAILS = "outputFeatureDetails"
@@ -156,6 +157,7 @@ class FeatureGroup(Builder):
         CONST_OUTPUT_FEATURE_DETAILS: "output_feature_details",
         CONST_STATISTICS_CONFIG: "statistics_config",
         CONST_INFER_SCHEMA: "is_infer_schema",
+        CONST_PARTITION_KEYS: "partition_keys"
     }
 
     def __init__(self, spec: Dict = None, **kwargs) -> None:
@@ -325,6 +327,37 @@ def with_primary_keys(self, primary_keys: List[str]) -> "FeatureGroup":
             },
         )
 
+    @property
+    def partition_keys(self) -> List[str]:
+        return self.get_spec(self.CONST_PARTITION_KEYS)
+
+    @partition_keys.setter
+    def partition_keys(self, value: List[str]):
+        self.with_partition_keys(value)
+
+    def with_partition_keys(self, partition_keys: List[str]) -> "FeatureGroup":
+        """Sets the partition keys of the feature group.
+
+        Parameters
+        ----------
+        partition_keys: List[str]
+            The List of partition keys for the feature group.
+
+        Returns
+        -------
+        FeatureGroup
+            The FeatureGroup instance (self)
+        """
+        return self.set_spec(
+            self.CONST_PARTITION_KEYS,
+            {
+                self.CONST_ITEMS: [
+                    {self.CONST_PRIMARY_KEY_NAME: partition_key}
+                    for partition_key in partition_keys or []
+                ]
+            },
+        )
+
     @property
     def feature_store_id(self) -> str:
         return self.get_spec(self.CONST_FEATURE_STORE_ID)