ODSC-45790: Improve error reporting. Improve get_feature_df output. H… (#281)

Author: KshitizLohia

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 and collection.items:
+            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
@@ -652,6 +700,7 @@ def create(self, **kwargs) -> "Dataset":
         # Create dataset
         logger.info("Saving dataset.")
         self.oci_dataset = self._to_oci_dataset(**kwargs).create()
+        self._update_from_oci_dataset_model(self.oci_dataset)
         self.with_id(self.oci_dataset.id)
         return self
 
@@ -729,11 +778,10 @@ def _update_from_oci_dataset_model(self, oci_dataset: OCIDataset) -> "Dataset":
 
         # Update the main properties
         self.oci_dataset = oci_dataset
-        dataset_details = oci_dataset.to_dict()
 
         for infra_attr, dsc_attr in self.attribute_map.items():
-            if infra_attr in dataset_details:
-                self.set_spec(infra_attr, dataset_details[infra_attr])
+            if dsc_attr in self.oci_dataset.attribute_map:
+                self.set_spec(infra_attr, getattr(self.oci_dataset, dsc_attr))
 
         return self
 
@@ -1013,6 +1061,10 @@ def to_dict(self) -> Dict:
         for key, value in spec.items():
             if hasattr(value, "to_dict"):
                 value = value.to_dict()
+            if hasattr(value, "attribute_map"):
+                value = self.oci_dataset.client.base_client.sanitize_for_serialization(
+                    value
+                )
             spec[key] = value
 
         return {

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/integration/feature_store/test_base.py

@@ -10,6 +10,7 @@
 
 import oci
 import pandas as pd
+from ads.feature_store.entity import Entity
 from great_expectations.core import ExpectationSuite, ExpectationConfiguration
 import ads
 import os
@@ -22,7 +23,7 @@
 
 client_kwargs = dict(
     retry_strategy=oci.retry.NoneRetryStrategy,
-    service_endpoint=os.getenv("service_endpoint"),
+    service_endpoint="http://127.0.0.1:21000/20230101",
 )
 ads.set_auth(client_kwargs=client_kwargs)
 
@@ -36,20 +37,26 @@
 
 
 def transformation_with_kwargs(data_frame, **kwargs):
-    is_area_enabled = kwargs.get('is_area_enabled')
+    is_area_enabled = kwargs.get("is_area_enabled")
 
     if is_area_enabled:
         # Calculate petal area and sepal area
-        data_frame["petal_area"] = data_frame["petal_length"] * data_frame["petal_width"]
-        data_frame["sepal_area"] = data_frame["sepal_length"] * data_frame["sepal_width"]
+        data_frame["petal_area"] = (
+            data_frame["petal_length"] * data_frame["petal_width"]
+        )
+        data_frame["sepal_area"] = (
+            data_frame["sepal_length"] * data_frame["sepal_width"]
+        )
 
     # Return the updated DataFrame
     return data_frame
 
 
 class FeatureStoreTestCase:
     # networks compartment in feature store
-    TIME_NOW = str.format("{}_{}",datetime.utcnow().strftime("%Y_%m_%d_%H_%M_%S"),int(random()*1000))
+    TIME_NOW = str.format(
+        "{}_{}", datetime.utcnow().strftime("%Y_%m_%d_%H_%M_%S"), int(random() * 1000)
+    )
     TENANCY_ID = "ocid1.tenancy.oc1..aaaaaaaa462hfhplpx652b32ix62xrdijppq2c7okwcqjlgrbknhgtj2kofa"
     COMPARTMENT_ID = "ocid1.tenancy.oc1..aaaaaaaa462hfhplpx652b32ix62xrdijppq2c7okwcqjlgrbknhgtj2kofa"
     METASTORE_ID = "ocid1.datacatalogmetastore.oc1.iad.amaaaaaabiudgxyap7tizm4gscwz7amu7dixz7ml3mtesqzzwwg3urvvdgua"
@@ -376,9 +383,11 @@ def create_entity_resource(self, feature_store) -> "Entity":
         return entity
 
     def create_transformation_resource(self, feature_store) -> "Transformation":
-        transformation = feature_store.create_transformation(source_code_func=transformation_with_kwargs,
-                                                             display_name="transformation_with_kwargs",
-                                                             transformation_mode=TransformationMode.PANDAS)
+        transformation = feature_store.create_transformation(
+            source_code_func=transformation_with_kwargs,
+            display_name="transformation_with_kwargs",
+            transformation_mode=TransformationMode.PANDAS,
+        )
         return transformation
 
     def define_feature_group_resource(

tests/integration/feature_store/test_dataset_complex.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*--
+
+# Copyright (c) 2023 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+import pytest
+import unittest
+
+from ads.feature_store.dataset import Dataset
+
+from ads.feature_store.entity import Entity
+
+from ads.feature_store.feature_store import FeatureStore
+
+from ads.feature_store.feature_group import FeatureGroup
+
+from ads.feature_store.feature_option_details import FeatureOptionDetails
+from tests.integration.feature_store.test_base import FeatureStoreTestCase
+
+
+class TestDatasetComplex(FeatureStoreTestCase):
+    """Contains integration tests for Dataset Delta changes."""
+
+    @pytest.fixture()
+    def feature_store(self) -> FeatureStore:
+        feature_store = self.define_feature_store_resource().create()
+        yield feature_store
+        # self.clean_up_feature_store(feature_store)
+
+    @pytest.fixture()
+    def entity(self, feature_store: FeatureStore):
+        entity = self.create_entity_resource(feature_store)
+        yield entity
+        # self.clean_up_entity(entity)
+
+    @pytest.fixture()
+    def feature_group(self, entity, feature_store) -> "FeatureGroup":
+        feature_group = self.define_feature_group_resource(
+            entity.oci_fs_entity.id, feature_store.oci_fs.id
+        ).create()
+        yield feature_group
+        # self.clean_up_feature_group(feature_group)
+
+    def test_manual_dataset(
+        self,
+        feature_store: FeatureStore,
+        entity: Entity,
+        feature_group: FeatureGroup,
+    ):
+        query = """
+             (SELECT
+                name, games, goals
+                FROM tblMadrid WHERE name = 'ronaldo')
+             UNION
+             (SELECT
+                name, games, goals
+                FROM tblBarcelona WHERE name = 'messi')
+            ORDER BY goals"""
+        name = self.get_name("fireside_football_debate")
+        dataset_resource = (
+            Dataset()
+            .with_description("dataset description")
+            .with_compartment_id(self.COMPARTMENT_ID)
+            .with_name(name)
+            .with_entity_id(entity_id=entity.id)
+            .with_feature_store_id(feature_store_id=feature_store.id)
+            .with_query(query)
+            .with_feature_groups([feature_group])
+        ).create()
+        assert len(dataset_resource.feature_groups) == 1
+        assert dataset_resource.feature_groups[0].id == feature_group.id
+        assert dataset_resource.get_spec(
+            Dataset.CONST_FEATURE_GROUP
+        ).is_manual_association
+        return dataset_resource

tests/integration/feature_store/test_dataset_delta.py

@@ -7,6 +7,8 @@
 import pytest
 import unittest
 
+from ads.feature_store.dataset import Dataset
+
 from ads.feature_store.feature_option_details import FeatureOptionDetails
 from tests.integration.feature_store.test_base import FeatureStoreTestCase
 
@@ -125,6 +127,9 @@ def test_dataset_materialise_overwrite(self, feature_group, dataset):
         dataset.materialise(ingestion_mode=IngestionMode.OVERWRITE)
 
         df = dataset.preview(row_count=50)
+        assert (
+            dataset.get_spec(Dataset.CONST_FEATURE_GROUP).is_manual_association == False
+        )
         assert df.count() == 14
         assert len(df.columns) == 6
 

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")