Fixing integration tests

Author: harsh97

ads/feature_store/dataset.py

@@ -865,10 +865,11 @@ def _update_from_oci_dataset_model(self, oci_dataset: OCIDataset) -> "Dataset":
                         features_list.append(output_feature)
 
                     value = {self.CONST_ITEMS: features_list}
-                else:
+                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)
-
         return self
 
     def materialise(

ads/feature_store/execution_strategy/spark/spark_execution.py

@@ -1,7 +1,5 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
-import json
-
 # 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/
 
@@ -29,8 +27,6 @@
     raise
 
 from ads.feature_store.common.enums import (
-    FeatureStoreJobType,
-    LifecycleState,
     EntityType,
     ExpectationType,
 )
@@ -47,6 +43,11 @@
 
 from ads.feature_store.feature_statistics.statistics_service import StatisticsService
 from ads.feature_store.common.utils.utility import validate_input_feature_details
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ads.feature_store.feature_group import FeatureGroup
+    from ads.feature_store.dataset import Dataset
 
 logger = logging.getLogger(__name__)
 
@@ -76,7 +77,10 @@ def __init__(self, metastore_id: str = None):
         self._jvm = self._spark_context._jvm
 
     def ingest_feature_definition(
-        self, feature_group, feature_group_job: FeatureGroupJob, dataframe
+        self,
+        feature_group: "FeatureGroup",
+        feature_group_job: FeatureGroupJob,
+        dataframe,
     ):
         try:
             self._save_offline_dataframe(dataframe, feature_group, feature_group_job)
@@ -90,7 +94,7 @@ def ingest_dataset(self, dataset, dataset_job: DatasetJob):
             raise SparkExecutionException(e).with_traceback(e.__traceback__)
 
     def delete_feature_definition(
-        self, feature_group, feature_group_job: FeatureGroupJob
+        self, feature_group: "FeatureGroup", feature_group_job: FeatureGroupJob
     ):
         """
         Deletes a feature definition from the system.
@@ -122,7 +126,7 @@ def delete_feature_definition(
                 output_details=output_details,
             )
 
-    def delete_dataset(self, dataset, dataset_job: DatasetJob):
+    def delete_dataset(self, dataset: "Dataset", dataset_job: DatasetJob):
         """
         Deletes a dataset from the system.
 
@@ -154,7 +158,7 @@ def delete_dataset(self, dataset, dataset_job: DatasetJob):
             )
 
     @staticmethod
-    def _validate_expectation(expectation_type, validation_output):
+    def _validate_expectation(expectation_type, validation_output: dict):
         """
         Validates the expectation based on the given expectation type and the validation output.
 

ads/feature_store/statistics.py

@@ -1,17 +1,94 @@
-import pandas as pd
-from typing import Dict
-from copy import deepcopy
+import matplotlib.figure
+from matplotlib.gridspec import GridSpec
+from matplotlib.figure import Figure
 
 from ads.feature_store.response.response_builder import ResponseBuilder
-from ads.jobs.builders.base import Builder
-from ads.common import utils
+
+import matplotlib.pyplot as plt
+import matplotlib.font_manager as fm
+import json
+
+
+def add_plots_for_stat(fig: Figure, feature: str, stat: dict):
+    freq_dist = stat.get(Statistics.CONST_FREQUENCY_DISTRIBUTION)
+    top_k = stat.get(Statistics.CONST_TOP_K_FREQUENT)
+    probability = stat.get(Statistics.CONST_PROBABILITY_DISTRIBUTION)
+
+    def subplot_generator():
+        plot_count = 0
+        if stat.get(Statistics.CONST_FREQUENCY_DISTRIBUTION) is not None:
+            plot_count += 1
+        # if stat.get(Statistics.CONST_TOP_K_FREQUENT) is not None:
+        #     plot_count += 1
+        if stat.get(Statistics.CONST_PROBABILITY_DISTRIBUTION) is not None:
+            plot_count += 1
+
+        for i in range(0, plot_count):
+            yield fig.add_subplot(1, plot_count, i + 1)
+
+    subplots = subplot_generator()
+    if freq_dist is not None:
+        axs = next(subplots)
+        fig.suptitle(feature, fontproperties=fm.FontProperties(weight="bold"))
+        axs.hist(
+            x=freq_dist.get("bins"),
+            weights=freq_dist.get("frequency"),
+            cumulative=False,
+            color="teal",
+            mouseover=True,
+            animated=True,
+        )
+
+        axs.set_xlabel(
+            "Bins", fontdict={"fontproperties": fm.FontProperties(size="xx-small")}
+        )
+        axs.set_ylabel(
+            "Frequency", fontdict={"fontproperties": fm.FontProperties(size="xx-small")}
+        )
+        axs.set_title(
+            "Frequency Distribution",
+            fontdict={"fontproperties": fm.FontProperties(size="small")},
+        )
+        axs.set_xticks(freq_dist.get("bins"))
+    if probability is not None:
+        axs = next(subplots)
+        fig.suptitle(feature, fontproperties=fm.FontProperties(weight="bold"))
+        axs.bar(
+            probability.get("bins"),
+            probability.get("density"),
+            color="teal",
+            mouseover=True,
+            animated=True,
+        )
+        axs.set_xlabel(
+            "Bins", fontdict={"fontproperties": fm.FontProperties(size="xx-small")}
+        )
+        axs.set_ylabel(
+            "Density", fontdict={"fontproperties": fm.FontProperties(size="xx-small")}
+        )
+        axs.set_title(
+            "Probability Distribution",
+            fontdict={"fontproperties": fm.FontProperties(size="small")},
+        )
+        axs.set_xticks(probability.get("bins"))
+
+
+def subfigure_generator(count: int, fig: Figure):
+    rows = count
+    subfigs = fig.subfigures(rows, 1)
+    for i in range(0, rows):
+        yield subfigs[i]
 
 
 class Statistics(ResponseBuilder):
     """
     Represents statistical information.
     """
 
+    CONST_FREQUENCY_DISTRIBUTION = "FrequencyDistribution"
+    CONST_PROBABILITY_DISTRIBUTION = "ProbabilityDistribution"
+    CONST_TOP_K_FREQUENT = "TopKFrequentElements"
+
     @property
     def kind(self) -> str:
         """
@@ -23,3 +100,26 @@ def kind(self) -> str:
             The kind of the statistics object, which is always "statistics".
         """
         return "statistics"
+
+    def to_viz(self):
+        if self.content is not None:
+            stats: dict = json.loads(self.content)
+            fig: Figure = plt.figure(figsize=(20, 20), dpi=150)
+            plt.subplots_adjust(hspace=3)
+
+            stats = {
+                feature: stat
+                for feature, stat in stats.items()
+                if Statistics.__graph_exists__(stat)
+            }
+            subfigures = subfigure_generator(len(stats), fig)
+            for feature, stat in stats.items():
+                sub_figure = next(subfigures)
+                add_plots_for_stat(sub_figure, feature, stat)
+
+    @staticmethod
+    def __graph_exists__(stat: dict):
+        return (
+            stat.get(Statistics.CONST_FREQUENCY_DISTRIBUTION) != None
+            or stat.get(Statistics.CONST_PROBABILITY_DISTRIBUTION) != None
+        )

tests/integration/feature_store/test_base.py

@@ -22,8 +22,8 @@
 
 
 client_kwargs = dict(
-    retry_strategy=oci.retry.NoneRetryStrategy,
-    service_endpoint=os.getenv("service_endpoint"),
+    retry_strategy=oci.retry.NoneRetryStrategy(),
+    fs_service_endpoint=os.getenv("service_endpoint"),
 )
 ads.set_auth(client_kwargs=client_kwargs)
 

tests/integration/feature_store/test_dataset_complex.py

@@ -70,8 +70,6 @@ def test_manual_dataset(
         ).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
+        assert dataset_resource.is_manual_association
         dataset_resource.delete()
         return dataset_resource