ODSC-45790: Addressing review comments

Author: harsh97

ads/feature_store/dataset.py

@@ -517,7 +517,7 @@ def feature_groups(self) -> List["FeatureGroup"]:
             self.CONST_FEATURE_GROUP
         )
         feature_groups: List["FeatureGroup"] = []
-        if collection:
+        if collection and collection.items:
             for datasetFGSummary in collection.items:
                 feature_groups.append(
                     FeatureGroup.from_id(datasetFGSummary.feature_group_id)
@@ -700,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
 
@@ -777,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
 
@@ -1061,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 {

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