ODSC-45790: Improve error reporting. Improve get_feature_df output. Handle manual featuregroup assosciation with dataset to support complex queries

Author: harsh97

ads/feature_store/dataset.py

@@ -8,6 +8,11 @@
 
 import pandas
 from great_expectations.core import ExpectationSuite
+from oci.feature_store.models import (
+    DatasetFeatureGroupCollection,
+    DatasetFeatureGroupSummary,
+)
+
 from ads.common import utils
 from ads.common.oci_mixin import OCIModelMixin
 from ads.feature_store.common.enums import (
@@ -26,6 +31,7 @@
     OciExecutionStrategyProvider,
 )
 from ads.feature_store.feature import DatasetFeature
+from ads.feature_store.feature_group import FeatureGroup
 from ads.feature_store.feature_group_expectation import Expectation
 from ads.feature_store.feature_option_details import FeatureOptionDetails
 from ads.feature_store.service.oci_dataset import OCIDataset
@@ -113,6 +119,7 @@ class Dataset(Builder):
     CONST_ITEMS = "items"
     CONST_LAST_JOB_ID = "jobId"
     CONST_MODEL_DETAILS = "modelDetails"
+    CONST_FEATURE_GROUP = "datasetFeatureGroups"
 
     attribute_map = {
         CONST_ID: "id",
@@ -130,6 +137,7 @@ class Dataset(Builder):
         CONST_LIFECYCLE_STATE: "lifecycle_state",
         CONST_MODEL_DETAILS: "model_details",
         CONST_PARTITION_KEYS: "partition_keys",
+        CONST_FEATURE_GROUP: "dataset_feature_groups",
     }
 
     def __init__(self, spec: Dict = None, **kwargs) -> None:
@@ -503,6 +511,44 @@ def with_model_details(self, model_details: ModelDetails) -> "Dataset":
 
         return self.set_spec(self.CONST_MODEL_DETAILS, model_details.to_dict())
 
+    @property
+    def feature_groups(self) -> List["FeatureGroup"]:
+        collection: "DatasetFeatureGroupCollection" = self.get_spec(
+            self.CONST_FEATURE_GROUP
+        )
+        feature_groups: List["FeatureGroup"] = []
+        if collection:
+            for datasetFGSummary in collection.items:
+                feature_groups.append(
+                    FeatureGroup.from_id(datasetFGSummary.feature_group_id)
+                )
+
+        return feature_groups
+
+    @feature_groups.setter
+    def feature_groups(self, feature_groups: List["FeatureGroup"]):
+        self.with_feature_groups(feature_groups)
+
+    def with_feature_groups(self, feature_groups: List["FeatureGroup"]) -> "Dataset":
+        """Sets the model details for the dataset.
+
+        Parameters
+        ----------
+        feature_groups: List of feature groups
+        Returns
+        -------
+        Dataset
+            The Dataset instance (self).
+
+        """
+        collection: List["DatasetFeatureGroupSummary"] = []
+        for group in feature_groups:
+            collection.append(DatasetFeatureGroupSummary(feature_group_id=group.id))
+
+        return self.set_spec(
+            self.CONST_FEATURE_GROUP, DatasetFeatureGroupCollection(items=collection)
+        )
+
     @property
     def partition_keys(self) -> List[str]:
         return self.get_spec(self.CONST_PARTITION_KEYS)
@@ -560,7 +606,9 @@ def add_models(self, model_details: ModelDetails) -> "Dataset":
                 f"Dataset update Failed with : {type(ex)} with error message: {ex}"
             )
             if existing_model_details:
-                self.with_model_details(ModelDetails().with_items(existing_model_details["items"]))
+                self.with_model_details(
+                    ModelDetails().with_items(existing_model_details["items"])
+                )
             else:
                 self.with_model_details(ModelDetails().with_items([]))
                 return self

ads/feature_store/execution_strategy/spark/spark_execution.py

@@ -16,6 +16,7 @@
     show_validation_summary,
 )
 from ads.feature_store.execution_strategy.engine.spark_engine import SparkEngine
+import traceback
 
 try:
     from pyspark.sql import DataFrame
@@ -291,8 +292,9 @@ def _save_offline_dataframe(
 
         except Exception as ex:
             error_details = str(ex)
+            tb = traceback.format_exc()
             logger.error(
-                f"FeatureGroup Materialization Failed with : {type(ex)} with error message: {ex}"
+                f"FeatureGroup Materialization Failed with : {type(ex)} with error message: {ex} and stacktrace {tb}",
             )
 
         show_ingestion_summary(
@@ -427,8 +429,9 @@ def _save_dataset_input(self, dataset, dataset_job: DatasetJob):
 
         except Exception as ex:
             error_details = str(ex)
+            tb = traceback.format_exc()
             logger.error(
-                f"Dataset Materialization Failed with : {type(ex)} with error message: {ex}"
+                f"Dataset Materialization Failed with : {type(ex)} with error message: {ex} and stacktrace {tb}"
             )
 
         show_ingestion_summary(

ads/feature_store/feature_group.py

@@ -748,7 +748,6 @@ def get_features_df(self) -> "pd.DataFrame":
                 {
                     "name": feature.feature_name,
                     "type": feature.feature_type,
-                    "feature_group_id": feature.feature_group_id,
                 }
             )
         return pd.DataFrame.from_records(records)

tests/unitary/with_extras/feature_store/test_dataset.py

@@ -259,12 +259,12 @@ def test__to_oci_fs_entity(self, mock_load_key_file, mock_config_from_file):
     @patch.object(SparkSessionSingleton, "__init__", return_value=None)
     @patch.object(SparkSessionSingleton, "get_spark_session")
     def test_materialise(self, spark, get_spark_session, mock_update):
-            with patch.object(DatasetJob, "create") as mock_dataset_job:
-                with patch.object(FeatureStore, "from_id"):
-                    with patch.object(DatasetJob, "_mark_job_complete"):
-                        mock_dataset_job.return_value = self.mock_dsc_dataset_job
-                        self.mock_dsc_dataset.with_id(DATASET_OCID)
-                        self.mock_dsc_dataset.materialise()
+        with patch.object(DatasetJob, "create") as mock_dataset_job:
+            with patch.object(FeatureStore, "from_id"):
+                with patch.object(DatasetJob, "_mark_job_complete"):
+                    mock_dataset_job.return_value = self.mock_dsc_dataset_job
+                    self.mock_dsc_dataset.with_id(DATASET_OCID)
+                    self.mock_dsc_dataset.materialise()
 
     @patch.object(SparkSessionSingleton, "__init__", return_value=None)
     @patch.object(SparkSessionSingleton, "get_spark_session")