handling both old and new response for to_viz

Author: najiyacl

ads/feature_store/execution_strategy/spark/spark_execution.py

@@ -77,10 +77,10 @@ def __init__(self, metastore_id: str = None):
         self._jvm = self._spark_context._jvm
 
     def ingest_feature_definition(
-        self,
-        feature_group: "FeatureGroup",
-        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)
@@ -94,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: "FeatureGroup", feature_group_job: FeatureGroupJob
+            self, feature_group: "FeatureGroup", feature_group_job: FeatureGroupJob
     ):
         """
         Deletes a feature definition from the system.
@@ -188,7 +188,7 @@ def _validate_expectation(expectation_type, validation_output: dict):
             raise Exception(error_message)
 
     def _save_offline_dataframe(
-        self, data_frame, feature_group, feature_group_job: FeatureGroupJob
+            self, data_frame, feature_group, feature_group_job: FeatureGroupJob
     ):
         """Ingest dataframe to the feature store system. as now this handles both spark dataframe and pandas
         dataframe. in case of pandas after transformation we convert it to spark and write to the delta.
@@ -211,6 +211,7 @@ def _save_offline_dataframe(
         feature_statistics = None
         validation_output = None
         output_features = []
+        version = 2  # after MLM upgrade
 
         try:
             # Create database in hive metastore if not exist
@@ -312,6 +313,7 @@ def _save_offline_dataframe(
             "validation_output": str(validation_output) if validation_output else None,
             "commit_id": "commit_id",
             "feature_statistics": feature_statistics,
+            "version": version
         }
 
         self._update_job_and_parent_details(
@@ -460,7 +462,7 @@ def _save_dataset_input(self, dataset, dataset_job: DatasetJob):
 
     @staticmethod
     def _update_job_and_parent_details(
-        parent_entity, job_entity, output_features=None, output_details=None
+            parent_entity, job_entity, output_features=None, output_details=None
     ):
         """
         Updates the parent and job entities with relevant details.

ads/feature_store/feature_group.py

@@ -1348,7 +1348,12 @@ def get_statistics(self, job_id: str = None) -> "Statistics":
         feature_statistics = (
             output_details.get("featureStatistics") if output_details else None
         )
-        return Statistics(feature_statistics)
+        stat_version = (
+            output_details.get("version") if output_details else None
+        )
+        version = stat_version if stat_version is not None else 1
+
+        return Statistics(feature_statistics, version)
 
     def get_validation_output(self, job_id: str = None) -> "ValidationOutput":
         """Retrieve validation report for the job with job_id

ads/feature_store/response/response_builder.py

@@ -12,7 +12,7 @@ class ResponseBuilder(Builder):
 
     CONST_CONTENT = "content"
 
-    def __init__(self, content: str):
+    def __init__(self, content: str, version: int = 1):
         """
         Initializes a new instance of the validation output class.
 
@@ -21,6 +21,7 @@ def __init__(self, content: str):
         content : str
             The validation output information as a JSON string.
         """
+        self.version = version
         super().__init__()
         self.set_spec(self.CONST_CONTENT, content)
 

ads/feature_store/statistics/abs_feature_value.py

@@ -28,15 +28,20 @@ def __validate__(self):
 
     @classmethod
     @abstractmethod
-    def __from_json__(cls, json_dict: dict):
+    def __from_json__(cls, json_dict: dict, version: int):
+        pass
+
+    @classmethod
+    @abstractmethod
+    def __from_json_v2__(cls, json_dict: dict):
         pass
 
     @classmethod
     def from_json(
-        cls, json_dict: dict, ignore_errors: bool = False
+        cls, json_dict: dict, version: int, ignore_errors: bool = False
     ) -> Union["AbsFeatureValue", None]:
         try:
-            return cls.__from_json__(json_dict=json_dict)
+            return cls.__from_json__(json_dict=json_dict, version=version)
         except Exception as e:
             if ignore_errors:
                 return None

ads/feature_store/statistics/charts/abstract_feature_plot.py

@@ -24,7 +24,7 @@ def add_to_figure(self, fig: Figure, xaxis: int, yaxis: int):
 
     @classmethod
     @abstractmethod
-    def __from_json__(cls, json_dict: dict):
+    def __from_json__(cls, json_dict: dict, version: int):
         pass
 
     @staticmethod

ads/feature_store/statistics/charts/box_plot.py

@@ -44,7 +44,17 @@ def __init__(self, q1: float, q2: float, q3: float):
             super().__init__()
 
         @classmethod
-        def __from_json__(cls, json_dict: dict) -> "BoxPlot.Quartiles":
+        def __from_json__(cls, json_dict: dict, version: int = 1) -> "BoxPlot.Quartiles":
+            if version == 2:
+                return cls.__from_json_v2__(json_dict)
+            return cls(
+                json_dict.get(cls.CONST_Q1),
+                json_dict.get(cls.CONST_Q2),
+                json_dict.get(cls.CONST_Q3),
+            )
+
+        @classmethod
+        def __from_json_v2__(cls, json_dict: dict) -> "BoxPlot.Quartiles":
             metric_data = json_dict.get(AbsFeatureValue.CONST_METRIC_DATA)
             return cls(
                 metric_data[0],
@@ -126,12 +136,12 @@ def get_box_points_from_frequency_distribution(
             return []
 
     @classmethod
-    def __from_json__(cls, json_dict: dict) -> "BoxPlot":
-        quartiles = cls.Quartiles.from_json(json_dict.get(cls.CONST_QUARTILES))
+    def __from_json__(cls, json_dict: dict, version: int = 1) -> "BoxPlot":
+        quartiles = cls.Quartiles.from_json(json_dict.get(cls.CONST_QUARTILES), version)
         return cls(
-            mean=GenericFeatureValue.from_json(json_dict.get(cls.CONST_MEAN)).val,
+            mean=GenericFeatureValue.from_json(json_dict.get(cls.CONST_MEAN),version).val,
             median=quartiles.q2,
-            sd=GenericFeatureValue.from_json(json_dict.get(cls.CONST_SD)).val,
+            sd=GenericFeatureValue.from_json(json_dict.get(cls.CONST_SD),version).val,
             q1=quartiles.q1,
             q3=quartiles.q3,
             box_points=json_dict.get(cls.CONST_BOX_POINTS),

ads/feature_store/statistics/charts/frequency_distribution.py

@@ -30,10 +30,20 @@ def __init__(self, frequency: List, bins: List):
     def __validate__(self):
         assert type(self.frequency) == list
         assert type(self.bins) == list
-        # assert 0 < len(self.frequency) == len(self.bins) > 0
+        assert 0 < len(self.frequency)
+        assert len(self.bins) > 0
 
     @classmethod
-    def __from_json__(cls, json_dict: dict) -> "FrequencyDistribution":
+    def __from_json__(cls, json_dict: dict, version: int = 1) -> "FrequencyDistribution":
+        if version == 2:
+            return cls.__from_json_v2__(json_dict)
+        return FrequencyDistribution(
+            frequency=json_dict.get(cls.CONST_FREQUENCY),
+            bins=json_dict.get(cls.CONST_BINS),
+        )
+
+    @classmethod
+    def __from_json_v2__(cls, json_dict: dict) -> "FrequencyDistribution":
         metric_data = json_dict.get(AbsFeatureValue.CONST_METRIC_DATA)
         return FrequencyDistribution(
             bins=metric_data[0],

ads/feature_store/statistics/charts/probability_distribution.py

@@ -33,7 +33,16 @@ def __validate__(self):
         # assert 0 < len(self.density) == len(self.bins) > 0
 
     @classmethod
-    def __from_json__(cls, json_dict: dict) -> "ProbabilityDistribution":
+    def __from_json__(cls, json_dict: dict, version: int = 1) -> "ProbabilityDistribution":
+        if version == 2:
+            return cls.__from_json_v2__(json_dict)
+        return cls(
+            density=json_dict.get(ProbabilityDistribution.CONST_DENSITY),
+            bins=json_dict.get(ProbabilityDistribution.CONST_BINS),
+        )
+
+    @classmethod
+    def __from_json_v2__(cls, json_dict: dict) -> "ProbabilityDistribution":
         metric_data = json_dict.get(AbsFeatureValue.CONST_METRIC_DATA)
         return cls(
             bins=metric_data[0],

ads/feature_store/statistics/charts/top_k_frequent_elements.py

@@ -44,7 +44,7 @@ def __validate__(self):
 
         @classmethod
         def __from_json__(
-            cls, json_dict: dict
+            cls, json_dict: dict, version: int = 1
         ) -> "TopKFrequentElements.TopKFrequentElement":
             return cls(
                 value=json_dict.get(cls.CONST_VALUE),
@@ -64,7 +64,17 @@ def __validate__(self):
         assert 0 < len(self.values) == len(self.estimates) > 0
 
     @classmethod
-    def __from_json__(cls, json_dict: dict) -> "TopKFrequentElements":
+    def __from_json__(cls, json_dict: dict, version: int = 1) -> "TopKFrequentElements":
+        if version == 2:
+            return cls.__from_json_v2__(json_dict)
+        elements = json_dict.get(cls.CONST_VALUE)
+        top_k_frequent_elements = [cls.TopKFrequentElement.__from_json__(element) for element in elements]
+        values = [element.value for element in top_k_frequent_elements]
+        estimates = [element.estimate for element in top_k_frequent_elements]
+        return cls(values, estimates)
+
+    @classmethod
+    def __from_json_v2__(cls, json_dict: dict) -> "TopKFrequentElements":
         metric_data = json_dict.get(AbsFeatureValue.CONST_METRIC_DATA)
         return cls(
             values=metric_data[0],

ads/feature_store/statistics/feature_stat.py

@@ -51,10 +51,10 @@ def __init__(
         self.box_plot = box_plot
 
     @classmethod
-    def from_json(cls, feature_name: str, json_dict: dict) -> "FeatureStatistics":
+    def from_json(cls, feature_name: str, json_dict: dict, version: int) -> "FeatureStatistics":
         if json_dict is not None:
             frequency_distribution = FrequencyDistribution.from_json(
-                json_dict.get(cls.CONST_FREQUENCY_DISTRIBUTION), ignore_errors=True
+                json_dict.get(cls.CONST_FREQUENCY_DISTRIBUTION), version, ignore_errors=True,
             )
 
             # inject box points for boxplot creation
@@ -66,14 +66,14 @@ def from_json(cls, feature_name: str, json_dict: dict) -> "FeatureStatistics":
             return cls(
                 feature_name,
                 TopKFrequentElements.from_json(
-                    json_dict.get(cls.CONST_TOP_K_FREQUENT), ignore_errors=True
+                    json_dict.get(cls.CONST_TOP_K_FREQUENT), version, ignore_errors=True
                 ),
                 frequency_distribution,
                 ProbabilityDistribution.from_json(
-                    json_dict.get(cls.CONST_PROBABILITY_DISTRIBUTION),
+                    json_dict.get(cls.CONST_PROBABILITY_DISTRIBUTION), version,
                     ignore_errors=True,
                 ),
-                BoxPlot.from_json(json_dict, ignore_errors=True),
+                BoxPlot.from_json(json_dict, version, ignore_errors=True),
             )
         else:
             return cls(feature_name)