added unit test and log level off

Author: yogesh266

ads/feature_store/common/spark_session_singleton.py

@@ -87,6 +87,7 @@ def __init__(self, metastore_id: str = None):
             self.spark_session = spark_builder.getOrCreate()
 
         self.spark_session.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
+        self.spark_session.sparkContext.setLogLevel("OFF")
 
     def get_spark_session(self):
         """Access method to get the spark session."""

ads/feature_store/common/utils/feature_schema_mapper.py

@@ -87,8 +87,8 @@ def map_pandas_type_to_feature_type(feature_name, values):
                 inferred_dtype = current_dtype
             else:
                 if (
-                        current_dtype != inferred_dtype
-                        and current_dtype is not FeatureType.UNKNOWN
+                    current_dtype != inferred_dtype
+                    and current_dtype is not FeatureType.UNKNOWN
                 ):
                     raise TypeError(
                         f"Input feature '{feature_name}' has mixed types, {current_dtype} and {inferred_dtype}. "
@@ -233,7 +233,7 @@ def map_feature_type_to_pandas(feature_type):
 
 
 def convert_pandas_datatype_with_schema(
-        raw_feature_details: List[dict], input_df: pd.DataFrame
+    raw_feature_details: List[dict], input_df: pd.DataFrame
 ):
     feature_detail_map = {}
     for feature_details in raw_feature_details:
@@ -245,13 +245,13 @@ def convert_pandas_datatype_with_schema(
             pandas_type = map_feature_type_to_pandas(feature_type)
             input_df[column] = (
                 input_df[column]
-                    .astype(pandas_type)
-                    .where(pd.notnull(input_df[column]), None)
+                .astype(pandas_type)
+                .where(pd.notnull(input_df[column]), None)
             )
 
 
 def map_spark_type_to_stats_data_type(spark_type):
-    """ Maps the spark data types to MLM library data types
+    """Maps the spark data types to MLM library data types
     args:
         param spark_type: Spark data type input from the feature dataframe on which we need stats
     :return:
@@ -272,7 +272,7 @@ def map_spark_type_to_stats_data_type(spark_type):
 
 
 def map_spark_type_to_stats_variable_type(spark_type):
-    """ Maps the spark data types to MLM library variable types
+    """Maps the spark data types to MLM library variable types
     args:
         param spark_type: Spark data type input from the feature dataframe on which we need stats
     :return:

ads/feature_store/common/utils/transformation_query_validator.py

@@ -13,14 +13,12 @@
     function_name: Transformation function name 
     input : The input placeholder for the FROM clause
 """
-USER_TRANSFORMATION_FUNCTION = \
-    """def {function_name}(input):
+USER_TRANSFORMATION_FUNCTION = """def {function_name}(input):
     sql_query = f\"\"\"{query}\"\"\"
     return sql_query"""
 
 
 class TransformationQueryValidator:
-
     @staticmethod
     def __verify_sql_query_plan(parser_plan, input_symbol: str):
         """
@@ -41,21 +39,23 @@ def __verify_sql_query_plan(parser_plan, input_symbol: str):
         cte = re.findall(r"CTE \[(.*?)\]", plan_string)
         table_names = []
         for plan_item in plan_items:
-            if plan_item['class'] == 'org.apache.spark.sql.catalyst.analysis.UnresolvedRelation':
-                table = plan_item['multipartIdentifier']
-                res = table.strip('][').split(', ')
+            if (
+                plan_item["class"]
+                == "org.apache.spark.sql.catalyst.analysis.UnresolvedRelation"
+            ):
+                table = plan_item["multipartIdentifier"]
+                res = table.strip("][").split(", ")
                 if len(res) >= 2:
-                    raise ValueError(
-                        "FROM Clause has invalid input {0}".format(table))
+                    raise ValueError("FROM Clause has invalid input {0}".format(table))
                 else:
                     if res[0].lower() != input_symbol.lower():
                         raise ValueError(
-                            f"Incorrect table template name, It should be {input_symbol}")
+                            f"Incorrect table template name, It should be {input_symbol}"
+                        )
                     if table not in cte:
                         table_names.append(f"{table}")
                         if len(table_names) > 1:
-                            raise ValueError(
-                                "Multiple tables are not supported ")
+                            raise ValueError("Multiple tables are not supported ")
 
     @staticmethod
     def verify_sql_input(query_input: str, input_symbol: str):
@@ -71,13 +71,17 @@ def verify_sql_input(query_input: str, input_symbol: str):
         try:
             parser_plan = parser.parsePlan(query_input)
         except ParseException as pe:
-            raise ParseException(f"Unable to parse the sql expression, exception occurred:  {pe}")
+            raise ParseException(
+                f"Unable to parse the sql expression, exception occurred:  {pe}"
+            )
 
         # verify if the parser plan has only FROM DATA_SOURCE_INPUT template
         TransformationQueryValidator.__verify_sql_query_plan(parser_plan, input_symbol)
 
     @staticmethod
-    def create_transformation_template(query: str, input_symbol: str, function_name: str):
+    def create_transformation_template(
+        query: str, input_symbol: str, function_name: str
+    ):
         """
         Creates the query transformation function to ensure backend integrity
         Args:
@@ -86,5 +90,7 @@ def create_transformation_template(query: str, input_symbol: str, function_name:
             function_name : The name of the transformation function
         """
         transformation_query = query.replace(input_symbol, "{input}")
-        output = USER_TRANSFORMATION_FUNCTION.format(query=transformation_query, function_name=function_name)
+        output = USER_TRANSFORMATION_FUNCTION.format(
+            query=transformation_query, function_name=function_name
+        )
         return output

ads/feature_store/feature_statistics/statistics_service.py

@@ -36,49 +36,64 @@ class StatisticsService:
 
     @staticmethod
     def compute_stats_with_mlm(
-            statistics_config: StatisticsConfig, input_df: DataFrame
+        statistics_config: StatisticsConfig, input_df: DataFrame
     ):
         feature_metrics = None
-        if bool(input_df.head(1)) and statistics_config and statistics_config.is_enabled:
+        if (
+            bool(input_df.head(1))
+            and statistics_config
+            and statistics_config.is_enabled
+        ):
             feature_schema = {}
             if input_df.schema:
-                StatisticsService.__get_mlm_supported_schema(feature_schema, input_df, statistics_config)
-                feature_metrics = StatisticsService.__get_feature_metric(feature_schema, input_df)
+                StatisticsService.__get_mlm_supported_schema(
+                    feature_schema, input_df, statistics_config
+                )
+                feature_metrics = StatisticsService.__get_feature_metric(
+                    feature_schema, input_df
+                )
             else:
                 raise ValueError("Dataframe schema is missing")
         return feature_metrics
 
     @staticmethod
     def __get_feature_metric(feature_schema: dict, data_frame: DataFrame):
         feature_metrics = None
-        runner = InsightsBuilder(). \
-            with_input_schema(input_schema=feature_schema). \
-            with_data_frame(data_frame=data_frame). \
-            with_engine(engine=EngineDetail(engine_name="spark")). \
-            build()
+        runner = (
+            InsightsBuilder()
+            .with_input_schema(input_schema=feature_schema)
+            .with_data_frame(data_frame=data_frame)
+            .with_engine(engine=EngineDetail(engine_name="spark"))
+            .build()
+        )
         result = runner.run()
         if result and result.profile:
             profile = result.profile
             feature_metrics = json.dumps(profile.to_json()[CONST_FEATURE_METRICS])
         else:
-            logger.warning(f"stats computation failed with MLM for schema {feature_schema}")
+            logger.warning(
+                f"stats computation failed with MLM for schema {feature_schema}"
+            )
         return feature_metrics
 
     @staticmethod
-    def __get_mlm_supported_schema(feature_schema: dict, input_df: DataFrame, statistics_config: StatisticsConfig):
+    def __get_mlm_supported_schema(
+        feature_schema: dict, input_df: DataFrame, statistics_config: StatisticsConfig
+    ):
         relevant_columns = statistics_config.columns
         for field in input_df.schema.fields:
             data_type = map_spark_type_to_stats_data_type(field.dataType)
             if not data_type:
-                logger.warning(f"Unable to map spark data type fields to MLM fields, "
-                               f"Actual data type {field.dataType}")
+                logger.warning(
+                    f"Unable to map spark data type fields to MLM fields, "
+                    f"Actual data type {field.dataType}"
+                )
             elif relevant_columns:
                 if field.name in relevant_columns:
-                    feature_schema[field.name] = FeatureType(data_type,
-                                                             map_spark_type_to_stats_variable_type(
-                                                                 field.dataType))
+                    feature_schema[field.name] = FeatureType(
+                        data_type, map_spark_type_to_stats_variable_type(field.dataType)
+                    )
             else:
-                feature_schema[field.name] = FeatureType(data_type,
-                                                         map_spark_type_to_stats_variable_type(
-                                                             field.dataType))
-
+                feature_schema[field.name] = FeatureType(
+                    data_type, map_spark_type_to_stats_variable_type(field.dataType)
+                )

ads/feature_store/feature_store.py

@@ -431,7 +431,7 @@ def _build_transformation(
         display_name: str = None,
         description: str = None,
         compartment_id: str = None,
-        sql_query: str = None
+        sql_query: str = None,
     ):
         transformation = (
             Transformation()
@@ -455,7 +455,7 @@ def create_transformation(
         display_name: str = None,
         description: str = None,
         compartment_id: str = None,
-        sql_query: str = None
+        sql_query: str = None,
     ) -> "Transformation":
         """Creates transformation resource from feature store.
 
@@ -491,7 +491,7 @@ def create_transformation(
             display_name,
             description,
             compartment_id,
-            sql_query
+            sql_query,
         )
 
         return self.oci_transformation.create()

ads/feature_store/feature_store_registrar.py

@@ -252,7 +252,9 @@ def _create_transformations(self) -> List[Transformation]:
             )
             if transformation.source_code_function:
                 # to encode to base64
-                transformation.source_code_function = transformation.source_code_function
+                transformation.source_code_function = (
+                    transformation.source_code_function
+                )
         return self._transformations.create_models(self._progress)
 
     def _create_feature_groups(self) -> List[FeatureGroup]:

ads/feature_store/transformation.py

@@ -19,7 +19,9 @@
 from ads.common.oci_mixin import OCIModelMixin
 from ads.feature_store.service.oci_transformation import OCITransformation
 from ads.jobs.builders.base import Builder
-from ads.feature_store.common.utils.transformation_query_validator import TransformationQueryValidator
+from ads.feature_store.common.utils.transformation_query_validator import (
+    TransformationQueryValidator,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -278,16 +280,18 @@ def with_source_code_function(self, source_code_func) -> "Transformation":
                 Base64EncoderDecoder.encode(source_code),
             )
         else:
-            return self.set_spec(
-                self.CONST_SOURCE_CODE, None)
+            return self.set_spec(self.CONST_SOURCE_CODE, None)
 
     @property
     def transformation_query_input(self) -> str:
         return self.get_spec(self.CONST_TRANSFORMATION_QUERY_INPUT)
 
     @transformation_query_input.setter
     def transformation_query_input(self, transformation_query):
-        self.set_spec(self.CONST_TRANSFORMATION_QUERY_INPUT, Base64EncoderDecoder.encode(transformation_query))
+        self.set_spec(
+            self.CONST_TRANSFORMATION_QUERY_INPUT,
+            Base64EncoderDecoder.encode(transformation_query),
+        )
 
     def with_transformation_query_input(self, transformation_query) -> "Transformation":
         """Sets the source code function for the transformation.
@@ -303,7 +307,8 @@ def with_transformation_query_input(self, transformation_query) -> "Transformati
             The Transformation instance (self)
         """
         return self.set_spec(
-            self.CONST_TRANSFORMATION_QUERY_INPUT, transformation_query)
+            self.CONST_TRANSFORMATION_QUERY_INPUT, transformation_query
+        )
 
     def with_id(self, id: str) -> "Transformation":
         return self.set_spec(self.CONST_ID, id)
@@ -380,7 +385,9 @@ def create(self, **kwargs) -> "Transformation":
 
         if not self.transformation_query_input:
             if not self.source_code_function:
-                raise ValueError("Transformation source code function must be provided.")
+                raise ValueError(
+                    "Transformation source code function must be provided."
+                )
 
             if not self.display_name:
                 self.display_name = self._transformation_function_name
@@ -390,13 +397,16 @@ def create(self, **kwargs) -> "Transformation":
                     "Transformation display name and function name must be same."
                 )
         elif self.transformation_mode.lower() == TransformationMode.SQL.value:
-            TransformationQueryValidator.verify_sql_input(self.transformation_query_input,
-                                                          self.CONST_DATA_SOURCE_TRANSFORMATION_INPUT)
+            TransformationQueryValidator.verify_sql_input(
+                self.transformation_query_input,
+                self.CONST_DATA_SOURCE_TRANSFORMATION_INPUT,
+            )
             # convert it to transformation function to ensure the integrity with backend
-            output = TransformationQueryValidator \
-                .create_transformation_template(self.transformation_query_input,
-                                                self.CONST_DATA_SOURCE_TRANSFORMATION_INPUT,
-                                                self.CONST_TRANSFORMATION_FUNCTION_NAME)
+            output = TransformationQueryValidator.create_transformation_template(
+                self.transformation_query_input,
+                self.CONST_DATA_SOURCE_TRANSFORMATION_INPUT,
+                self.CONST_TRANSFORMATION_FUNCTION_NAME,
+            )
             self.with_display_name(self.CONST_TRANSFORMATION_FUNCTION_NAME)
             self.with_source_code_function(output)
         else: