Adding sql validation/pretty printing

Author: harsh97

ads/feature_store/dataset.py

@@ -1,12 +1,12 @@
 #!/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
@@ -159,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)
@@ -191,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."""
@@ -572,8 +583,18 @@ def with_feature_groups(self, feature_groups: List["FeatureGroup"]) -> "Dataset"
         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)
+            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
@@ -687,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=True, **kwargs) -> "Dataset":
         """Creates dataset  resource.
 
         !!! note "Lazy"
@@ -700,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
         -------
@@ -720,6 +743,9 @@ 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}")

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)
 
@@ -940,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.
@@ -1165,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.