Addressing review comments

Author: harsh97

ads/feature_store/dataset.py

@@ -164,7 +164,7 @@ def __init__(self, spec: Dict = None, **kwargs) -> None:
         self.oci_dataset = self._to_oci_dataset(**kwargs)
         self.lineage = OCILineage(**kwargs)
 
-    def _to_oci_dataset(self, **kwargs):
+    def _to_oci_dataset(self, **kwargs) -> OCIDataset:
         """Creates an `OCIDataset` instance from the  `Dataset`.
 
         kwargs
@@ -235,8 +235,8 @@ def name(self) -> str:
         return self.get_spec(self.CONST_NAME)
 
     @name.setter
-    def name(self, name: str) -> "Dataset":
-        return self.with_name(name)
+    def name(self, name: str):
+        self.with_name(name)
 
     def with_name(self, name: str) -> "Dataset":
         """Sets the name.
@@ -865,8 +865,6 @@ def _update_from_oci_dataset_model(self, oci_dataset: OCIDataset) -> "Dataset":
                         features_list.append(output_feature)
 
                     value = {self.CONST_ITEMS: features_list}
-                elif infra_attr == self.CONST_FEATURE_GROUP:
-                    value = getattr(self.oci_dataset, dsc_attr)
                 else:
                     value = dataset_details[infra_attr]
                 self.set_spec(infra_attr, value)
@@ -1207,8 +1205,14 @@ def to_dict(self) -> Dict:
         for key, value in spec.items():
             if hasattr(value, "to_dict"):
                 value = value.to_dict()
-            spec[key] = value
-
+            if key == self.CONST_FEATURE_GROUP:
+                spec[
+                    key
+                ] = self.oci_dataset.client.base_client.sanitize_for_serialization(
+                    value
+                )
+            else:
+                spec[key] = value
         return {
             "kind": self.kind,
             "type": self.type,

ads/feature_store/feature_stat.py

@@ -1,12 +1,19 @@
-import json
 from abc import abstractmethod
 
-from ads.feature_store.common.utils.utility import none_type_safe_json_loads
-from plotly.graph_objs import Figure
+from ads.common.decorator.runtime_dependency import OptionalDependency
+
 from typing import List
-import plotly
-import plotly.graph_objects as go
-from plotly.subplots import make_subplots
+
+try:
+    import plotly
+    from plotly.graph_objs import Figure
+    import plotly.graph_objects as go
+    from plotly.subplots import make_subplots
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        f"The `plotly` module was not found. Please run `pip install "
+        f"{OptionalDependency.FEATURE_STORE}`."
+    )
 
 
 class FeatureStat:
@@ -229,7 +236,7 @@ def to_viz(self):
             fig.layout.title = self.CONST_TITLE_FORMAT.format(self.feature_name)
             fig.update_layout(title_font_size=20)
             fig.update_layout(title_x=0.5)
-            fig.update_layout(showlegend=False)
+            # fig.update_layout(showlegend=False)
             plotly.offline.iplot(
                 fig,
                 filename=self.CONST_PLOT_FORMAT.format(self.feature_name),