2828from ads .feature_store .execution_strategy .execution_strategy_provider import (
2929 OciExecutionStrategyProvider ,
3030)
31- from ads .feature_store .feature import Feature
3231from ads .feature_store .feature import DatasetFeature
3332from ads .feature_store .feature_group_expectation import Expectation
3433from ads .feature_store .feature_option_details import FeatureOptionDetails
@@ -247,7 +246,7 @@ def id(self) -> str:
247246 return self .get_spec (self .CONST_ID )
248247
249248 @property
250- def features (self ) -> List [Feature ]:
249+ def features (self ) -> List [DatasetFeature ]:
251250 return [
252251 DatasetFeature (** feature_dict )
253252 for feature_dict in self .get_spec (self .CONST_OUTPUT_FEATURE_DETAILS )[
@@ -259,7 +258,6 @@ def features(self) -> List[Feature]:
259258 def with_id (self , id : str ) -> "Dataset" :
260259 return self .set_spec (self .CONST_ID , id )
261260
262-
263261 def with_job_id (self , dataset_job_id : str ) -> "Dataset" :
264262 """Sets the job_id for the last running job.
265263
@@ -721,12 +719,12 @@ def delete(self):
721719
722720 dataset_execution_strategy .delete_dataset (self , dataset_job )
723721
724- def get_features (self ) -> List [Feature ]:
722+ def get_features (self ) -> List [DatasetFeature ]:
725723 """
726724 Returns all the features in the dataset.
727725
728726 Returns:
729- List[Feature ]
727+ List[DatasetFeature ]
730728 """
731729
732730 return self .features
@@ -740,13 +738,9 @@ def get_features_df(self) -> "pandas.DataFrame":
740738 """
741739 records = []
742740 for feature in self .features :
743- records .append (
744- {
745- "name" : feature .feature_name ,
746- "type" : feature .feature_type
747- }
748- )
741+ records .append ({"name" : feature .feature_name , "type" : feature .feature_type })
749742 return pandas .DataFrame .from_records (records )
743+
750744 def update (self , ** kwargs ) -> "Dataset" :
751745 """Updates Dataset in the feature store.
752746
@@ -794,9 +788,7 @@ def _update_from_oci_dataset_model(self, oci_dataset: OCIDataset) -> "Dataset":
794788 # May not need if we fix the backend and add feature_group_id to the output_feature
795789 features_list = []
796790 for output_feature in dataset_details [infra_attr ]["items" ]:
797- output_feature ["datasetId" ] = dataset_details [
798- self .CONST_ID
799- ]
791+ output_feature ["datasetId" ] = dataset_details [self .CONST_ID ]
800792 features_list .append (output_feature )
801793
802794 value = {self .CONST_ITEMS : features_list }
0 commit comments