teradataml 17.20.0.6__py3-none-any.whl → 20.0.0.0__py3-none-any.whl

teradataml/catalog/__init__.py

@@ -1,4 +1,2 @@
-from teradataml.catalog.model_cataloging import save_model, list_models, delete_model,\
-    describe_model, retrieve_model, publish_model
 from teradataml.catalog.byom import save_byom, delete_byom, list_byom, retrieve_byom,\
     set_byom_catalog, set_license, get_license

teradataml/catalog/byom.py

@@ -26,9 +26,11 @@ from teradataml.options.display import display
 from teradataml.common.constants import ModelCatalogingConstants as mac
 from teradataml.options.configure import configure
 from teradataml.utils.utils import execute_sql
+from teradatasqlalchemy.telemetry.queryband import collect_queryband
 
 validator = _Validators()
 
+
 def __check_if_model_exists(model_id,
                             table_name,
                             schema_name=None,
@@ -101,6 +103,7 @@ def __check_if_model_exists(model_id,
 
     return True if num_rows == 1 else False
 
+
 def __set_validate_catalog_parameters(table_name=None, schema_name=None):
     """
      DESCRIPTION:
@@ -133,7 +136,7 @@ def __set_validate_catalog_parameters(table_name=None, schema_name=None):
             Types: str
 
     RETURNS:
-        bool.
+        List of "table_name" and "schema_name".
 
     RAISES:
         ValueError
@@ -159,6 +162,7 @@ def __set_validate_catalog_parameters(table_name=None, schema_name=None):
     return [table_name, schema_name]
 
 
+@collect_queryband(queryband="stByomCtlg")
 def set_byom_catalog(table_name,
                      schema_name=None):
     """
@@ -222,6 +226,8 @@ def set_byom_catalog(table_name,
     print("The model cataloging parameters are set to table_name='{}' and "
           "schema_name='{}'".format(table_name, schema_name))
 
+
+@collect_queryband(queryband="svByom")
 def save_byom(model_id,
               model_file,
               table_name=None,
@@ -318,7 +324,7 @@ def save_byom(model_id,
                     * datetime.date
                     * datetime.time
                  2. "additional_columns" does not accept keys model_id and model.
-            Types: str
+            Types: dict
 
         additional_columns_types:
             Optional Argument.
@@ -541,7 +547,7 @@ def save_byom(model_id,
             __check_if_model_exists(
                 model_id, table_name, schema_name, raise_error_if_model_found=True)
             if len(additional_columns_types) != 0:
-                warnings.warn("Argument additional_columns_types is ignored since table already exists.")
+                warnings.warn("Argument additional_columns_types is ignored since table already exists.", stacklevel=2)
 
             # Gather column name and type information from existing table
             existing_table_df = DataFrame(in_schema(schema_name, table_name))
@@ -636,7 +642,7 @@ def save_byom(model_id,
             if len(undefined_column_types) > 0:
                 warnings.warn("Specified table does not exist and data types of {0} "\
                         "columns are not provided. Taking default datatypes."\
-                              .format(", ".join(undefined_column_types)))
+                              .format(", ".join(undefined_column_types)), stacklevel=2)
 
             # Create empty vantage table using sqlalchemy object.
             _create_table(
@@ -651,6 +657,8 @@ def save_byom(model_id,
         columns_clause = ", ".join(column_names)
         values_clause = ", ".join(("?" for _ in range(len(column_names))))
         insert_model = f"insert into {table_name} ({columns_clause}) values ({values_clause});"
+        # Empty queryband buffer before SQL call.
+        UtilFuncs._set_queryband()
         execute_sql(insert_model, tuple([insert_parameters]))
         print("Model is saved.")
 
@@ -660,6 +668,8 @@ def save_byom(model_id,
         error_code = MessageCodes.MODEL_CATALOGING_OPERATION_FAILED
         raise TeradataMlException(Messages.get_message(error_code, "save", str(err)), error_code)
 
+
+@collect_queryband(queryband="dltByom")
 def delete_byom(model_id, table_name=None, schema_name=None):
     """
     DESCRIPTION:
@@ -811,6 +821,8 @@ def delete_byom(model_id, table_name=None, schema_name=None):
 
     try:
         delete_model = f"delete from {table_name} where model_id = (?)"
+        # Empty queryband buffer before SQL call.
+        UtilFuncs._set_queryband()
         execute_sql(delete_model, tuple([model_id]))
         print("Model is deleted.")
 
@@ -821,6 +833,8 @@ def delete_byom(model_id, table_name=None, schema_name=None):
         error_msg = Messages.get_message(error_code, "delete", str(err))
         raise TeradataMlException(error_msg, error_code)
 
+
+@collect_queryband(queryband="stLcns")
 def set_license(license,
                 table_name=None,
                 schema_name=None,
@@ -1030,6 +1044,9 @@ def set_license(license,
                                   schema_name=schema_name, check_constraint='id between 1 and 1')
                     query = "insert into {}.{} values (1, '{}')".format(
                         schema_name, license_table, configure._byom_model_catalog_license)
+
+                    # Empty queryband buffer before SQL call.
+                    UtilFuncs._set_queryband()
                     execute_sql(query)
                 except:
                     raise
@@ -1049,6 +1066,7 @@ def set_license(license,
                                           MessageCodes.TABLE_ALREADY_EXISTS)
 
 
+@collect_queryband(queryband="gtLcns")
 def get_license():
     """
         DESCRIPTION:
@@ -1105,6 +1123,7 @@ def get_license():
         print('Set the license information using set_license() function.')
 
 
+@collect_queryband(queryband="rtrvByom")
 def retrieve_byom(model_id,
                   table_name=None,
                   schema_name=None,
@@ -1112,7 +1131,8 @@ def retrieve_byom(model_id,
                   is_license_column=False,
                   license_table_name=None,
                   license_schema_name=None,
-                  require_license=False):
+                  require_license=False,
+                  return_addition_columns=False):
     """
     DESCRIPTION:
         Function to retrieve a saved model. Output of this function can be
@@ -1195,6 +1215,13 @@ def retrieve_byom(model_id,
             Default value: False
             Types: bool
 
+        return_addition_columns:
+            Optional Argument.
+            Specifies whether to return additional columns saved during save_byom() along with
+            model_id and model columns.
+            Default value: False
+            Types: bool
+
         Note:
             The following table describes the system behaviour in different scenarios:
             +----------------+-------------+-----------------------+-------------+-------------------------------------+
@@ -1243,6 +1270,7 @@ def retrieve_byom(model_id,
         Model is saved.
         >>> save_byom('model6', model_file, 'byom_models', schema_name='test')
         Model is saved.
+        >>> # Save the license in an addtional column named "license_data" in the model table.
         >>> save_byom('licensed_model1', model_file, 'byom_licensed_models', additional_columns={"license_data": "A5sUL9KU_kP35Vq"})
         Created the model table 'byom_licensed_models' as it does not exist.
         Model is saved.
@@ -1402,6 +1430,15 @@ def retrieve_byom(model_id,
                                             model
         model_id
         licensed_model1  b'504B03041400080808...'
+
+        # Example 12 - Retrieve a model with id 'licensed_model1' from the table along with all
+        #              additional columns saved during save_byom().
+        >>> df = retrieve_byom('licensed_model1', table_name='byom_licensed_models',
+                               return_addition_columns=True)
+        >>> df
+                                            model     license_data
+        model_id
+        licensed_model1  b'504B03041400080808...'  A5sUL9KU_kP35Vq
     """
 
 
@@ -1469,6 +1506,9 @@ def retrieve_byom(model_id,
     #   product should hold only record in the DataFrame.
 
     if not license:
+        if return_addition_columns:
+            # Return all columns if return_addition_columns is True.
+            return model_details
         return model_details.select(["model_id", "model"])
 
     # Lambda function for attaching the license to model DataFrame.
@@ -1515,13 +1555,16 @@ def retrieve_byom(model_id,
             error_code, "retrieve", "Table which holds license key should have only one row.")
         raise TeradataMlException(error_msg, error_code)
 
-    model_details = model_details.select(["model_id", "model"])
+    if not return_addition_columns:
+        # Return only model_id and model columns if return_addition_columns is False.
+        model_details = model_details.select(["model_id", "model"])
 
     # Make sure license is the column name for license key.
     license_df = license_df.assign(drop_columns=True, license=license_df[license])
     return model_details.join(license_df, how="cross")
 
 
+@collect_queryband(queryband="lstByom")
 def list_byom(table_name=None, schema_name=None, model_id=None):
     """
     DESCRIPTION:

teradataml/catalog/function_argument_mapper.py

@@ -176,8 +176,6 @@ class _ArgumentMapper(_ArgumentMapperSuper):
             function_json_file = None
 
             # The path of the JSON files for the function is expected to be:
-            # teradataml/analytics/mle/<function_name>_mle.json - for MLE functions
-            # OR
             # teradataml/analytics/sqle/<function_name>_sqle.json - for SQLE functions
             if engine == ModelCatalogingConstants.MODEL_ENGINE_ML.value:
                 function_json_file = os.path.join(curr_dir,