Odsc 45790 support complex queries (#291)

Author: KshitizLohia

ads/feature_store/dataset.py

@@ -1,15 +1,20 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
-import json
 import logging
 from copy import deepcopy
 from datetime import datetime
 from typing import Dict, List, Union
 
 import pandas
+import pandas as pd
 from great_expectations.core import ExpectationSuite
 
 from ads import deprecated
+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 (
@@ -29,6 +34,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
@@ -116,6 +122,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",
@@ -133,6 +140,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:
@@ -151,6 +159,7 @@ def __init__(self, spec: Dict = None, **kwargs) -> None:
         super().__init__(spec=spec, **deepcopy(kwargs))
         # Specify oci Dataset instance
         self.dataset_job = None
+        self._is_manual_association: bool = False
         self._spark_engine = None
         self.oci_dataset = self._to_oci_dataset(**kwargs)
         self.lineage = OCILineage(**kwargs)
@@ -183,6 +192,16 @@ def spark_engine(self):
             self._spark_engine = SparkEngine(get_metastore_id(self.feature_store_id))
         return self._spark_engine
 
+    @property
+    def is_manual_association(self):
+        collection: DatasetFeatureGroupCollection = self.get_spec(
+            self.CONST_FEATURE_GROUP
+        )
+        if collection and collection.is_manual_association is not None:
+            return collection.is_manual_association
+        else:
+            return self._is_manual_association
+
     @property
     def kind(self) -> str:
         """The kind of the object as showing in a YAML."""
@@ -530,6 +549,54 @@ 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))
+
+        self._is_manual_association = True
+        return self.set_spec(
+            self.CONST_FEATURE_GROUP,
+            DatasetFeatureGroupCollection(items=collection, is_manual_association=True),
+        )
+
+    def feature_groups_to_df(self):
+        return pd.DataFrame.from_records(
+            [
+                feature_group.oci_feature_group.to_df_record()
+                for feature_group in self.feature_groups
+            ]
+        )
+
     @property
     def partition_keys(self) -> List[str]:
         return self.get_spec(self.CONST_PARTITION_KEYS)
@@ -641,7 +708,7 @@ def show(self, rankdir: str = GraphOrientation.LEFT_RIGHT) -> None:
                 f"Can't get lineage information for Feature group id {self.id}"
             )
 
-    def create(self, **kwargs) -> "Dataset":
+    def create(self, validate_sql=False, **kwargs) -> "Dataset":
         """Creates dataset  resource.
 
         !!! note "Lazy"
@@ -654,6 +721,8 @@ def create(self, **kwargs) -> "Dataset":
         kwargs
             Additional kwargs arguments.
             Can be any attribute that `oci.feature_store.models.Dataset` accepts.
+        validate_sql:
+            Boolean value indicating whether to validate sql before creating dataset
 
         Returns
         -------
@@ -674,13 +743,17 @@ def create(self, **kwargs) -> "Dataset":
         if self.statistics_config is None:
             self.statistics_config = StatisticsConfig()
 
+        if validate_sql is True:
+            self.spark_engine.sql(self.get_spec(self.CONST_QUERY))
+
         payload = deepcopy(self._spec)
         payload.pop("id", None)
         logger.debug(f"Creating a dataset resource with payload {payload}")
 
         # 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
 
@@ -793,8 +866,7 @@ def _update_from_oci_dataset_model(self, oci_dataset: OCIDataset) -> "Dataset":
 
                     value = {self.CONST_ITEMS: features_list}
                 else:
-                    value = dataset_details[infra_attr]
-
+                    value = getattr(self.oci_dataset, dsc_attr)
                 self.set_spec(infra_attr, value)
 
         return self
@@ -1134,6 +1206,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

@@ -180,7 +180,7 @@ def __init__(self, spec: Dict = None, **kwargs) -> None:
         # Specify oci FeatureGroup instance
         self.feature_group_job = None
         self._spark_engine = None
-        self.oci_feature_group = self._to_oci_feature_group(**kwargs)
+        self.oci_feature_group: OCIFeatureGroup = self._to_oci_feature_group(**kwargs)
         self.dsc_job = OCIFeatureGroupJob()
         self.lineage = OCILineage(**kwargs)
 
@@ -750,8 +750,12 @@ def get_features_df(self) -> "pd.DataFrame":
         """
         records = []
         for feature in self.features:
-            records.append({"name": feature.feature_name, "type": feature.feature_type})
-
+            records.append(
+                {
+                    "name": feature.feature_name,
+                    "type": feature.feature_type,
+                }
+            )
         return pd.DataFrame.from_records(records)
 
     def get_input_features_df(self) -> "pd.DataFrame":
@@ -936,7 +940,7 @@ def get_last_job(self) -> "FeatureGroupJob":
             return fg_job[0]
         return FeatureGroupJob.from_id(self.job_id)
 
-    def select(self, features: Optional[List[str]] = []) -> Query:
+    def select(self, features: Optional[List[str]] = ()) -> Query:
         """
         Selects a subset of features from the feature group and returns a Query object that can be used to view the
         resulting dataframe.
@@ -1161,28 +1165,9 @@ def list_df(cls, compartment_id: str = None, **kwargs) -> "pd.DataFrame":
         for oci_feature_group in OCIFeatureGroup.list_resource(
             compartment_id, **kwargs
         ):
-            records.append(
-                {
-                    "id": oci_feature_group.id,
-                    "name": oci_feature_group.name,
-                    "description": oci_feature_group.description,
-                    "time_created": oci_feature_group.time_created.strftime(
-                        utils.date_format
-                    ),
-                    "time_updated": oci_feature_group.time_updated.strftime(
-                        utils.date_format
-                    ),
-                    "lifecycle_state": oci_feature_group.lifecycle_state,
-                    "created_by": f"...{oci_feature_group.created_by[-6:]}",
-                    "compartment_id": f"...{oci_feature_group.compartment_id[-6:]}",
-                    "primary_keys": oci_feature_group.primary_keys,
-                    "feature_store_id": oci_feature_group.feature_store_id,
-                    "entity_id": oci_feature_group.entity_id,
-                    "input_feature_details": oci_feature_group.input_feature_details,
-                    "expectation_details": oci_feature_group.expectation_details,
-                    "statistics_config": oci_feature_group.statistics_config,
-                }
-            )
+            oci_feature_group: OCIFeatureGroup = oci_feature_group
+            records.append(oci_feature_group.to_df_record())
+
         return pd.DataFrame.from_records(records)
 
     @classmethod

ads/feature_store/service/oci_feature_group.py

@@ -5,6 +5,8 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import datetime
+import pandas as pd
+from ads.common import utils
 
 import oci
 from oci.feature_store.models import (
@@ -64,9 +66,9 @@ class OCIFeatureGroup(OCIFeatureStoreMixin, oci.feature_store.models.FeatureGrou
         Gets feature group by OCID.
     Examples
     --------
-    >>> oci_feature_group = OCIFeatureGroup.from_id("<feature_group_id>")
-    >>> oci_feature_group.description = "A brand new description"
-    >>> oci_feature_group.delete()
+    >>> self = OCIFeatureGroup.from_id("<feature_group_id>")
+    >>> self.description = "A brand new description"
+    >>> self.delete()
     """
 
     def create(self) -> "OCIFeatureGroup":
@@ -122,6 +124,27 @@ def delete(self):
         """
         self.client.delete_feature_group(self.id)
 
+    def to_df(self):
+        return pd.DataFrame.from_records([self.to_df_record()])
+
+    def to_df_record(self):
+        return {
+            "id": self.id,
+            "name": self.name,
+            "description": self.description,
+            "time_created": self.time_created.strftime(utils.date_format),
+            "time_updated": self.time_updated.strftime(utils.date_format),
+            "lifecycle_state": self.lifecycle_state,
+            "created_by": f"...{self.created_by[-6:]}",
+            "compartment_id": f"...{self.compartment_id[-6:]}",
+            "primary_keys": self.primary_keys,
+            "feature_store_id": self.feature_store_id,
+            "entity_id": self.entity_id,
+            "input_feature_details": self.input_feature_details,
+            "expectation_details": self.expectation_details,
+            "statistics_config": self.statistics_config,
+        }
+
     @classmethod
     def from_id(cls, id: str) -> "OCIFeatureGroup":
         """Gets feature group resource  by id.

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
@@ -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(