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

@@ -10,6 +10,11 @@
 from great_expectations.core import ExpectationSuite
 
 from ads import deprecated
+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 (
@@ -29,6 +34,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
@@ -116,6 +122,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",
@@ -133,6 +140,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:
@@ -530,6 +538,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)

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

@@ -750,8 +750,12 @@ def get_features_df(self) -> "pd.DataFrame":
         """
         records = []
         for feature in self.features:
-            records.append({"name": feature.feature_name, "type": feature.feature_type})
-
+            records.append(
+                {
+                    "name": feature.feature_name,
+                    "type": feature.feature_type,
+                }
+            )
         return pd.DataFrame.from_records(records)
 
     def get_input_features_df(self) -> "pd.DataFrame":