# Copyright (c) 2025 Oracle and/or its affiliates.
|
#
|
# This program is free software; you can redistribute it and/or modify
|
# it under the terms of the GNU General Public License, version 2.0, as
|
# published by the Free Software Foundation.
|
#
|
# This program is designed to work with certain software (including
|
# but not limited to OpenSSL) that is licensed under separate terms,
|
# as designated in a particular file or component or in included license
|
# documentation. The authors of MySQL hereby grant you an
|
# additional permission to link the program and your derivative works
|
# with the separately licensed software that they have either included with
|
# the program or referenced in the documentation.
|
#
|
# Without limiting anything contained in the foregoing, this file,
|
# which is part of MySQL Connector/Python, is also subject to the
|
# Universal FOSS Exception, version 1.0, a copy of which can be found at
|
# http://oss.oracle.com/licenses/universal-foss-exception.
|
#
|
# This program is distributed in the hope that it will be useful, but
|
# WITHOUT ANY WARRANTY; without even the implied warranty of
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
|
# See the GNU General Public License, version 2.0, for more details.
|
#
|
# You should have received a copy of the GNU General Public License
|
# along with this program; if not, write to the Free Software Foundation, Inc.,
|
# 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
|
"""HeatWave ML model utilities for MySQL Connector/Python.
|
|
Provides classes to manage training, prediction, scoring, and explanations
|
via MySQL HeatWave stored procedures.
|
"""
|
import copy
|
import json
|
|
from enum import Enum
|
from typing import Any, Dict, Optional, Union
|
|
import numpy as np
|
import pandas as pd
|
|
from mysql.ai.utils import (
|
VAR_NAME_SPACE,
|
atomic_transaction,
|
convert_to_df,
|
execute_sql,
|
format_value_sql,
|
get_random_name,
|
source_schema,
|
sql_response_to_df,
|
sql_table_from_df,
|
sql_table_to_df,
|
table_exists,
|
temporary_sql_tables,
|
validate_name,
|
)
|
from mysql.connector.abstracts import MySQLConnectionAbstract
|
|
|
class ML_TASK(Enum): # pylint: disable=invalid-name
|
"""Enumeration of supported ML tasks for HeatWave."""
|
|
CLASSIFICATION = "classification"
|
REGRESSION = "regression"
|
FORECASTING = "forecasting"
|
ANOMALY_DETECTION = "anomaly_detection"
|
LOG_ANOMALY_DETECTION = "log_anomaly_detection"
|
RECOMMENDATION = "recommendation"
|
TOPIC_MODELING = "topic_modeling"
|
|
@staticmethod
|
def get_task_string(task: Union[str, "ML_TASK"]) -> str:
|
"""
|
Return the string representation of a machine learning task.
|
|
Args:
|
task (Union[str, ML_TASK]): The task to convert.
|
Accepts either a task enum member (ML_TASK) or a string.
|
|
Returns:
|
str: The string value of the ML task.
|
"""
|
|
if isinstance(task, str):
|
return task
|
|
return task.value
|
|
|
class _MyModelCommon:
|
"""
|
Common utilities and workflow for MySQL HeatWave ML models.
|
|
This class handles model lifecycle steps such as loading, fitting, scoring,
|
making predictions, and explaining models or predictions. Not intended for
|
direct instantiation, but as a superclass for heatwave model wrappers.
|
|
Attributes:
|
db_connection: MySQL connector database connection.
|
task: ML task, e.g., "classification" or "regression".
|
model_name: Identifier of model in MySQL.
|
schema_name: Database schema used for operations and temp tables.
|
"""
|
|
def __init__(
|
self,
|
db_connection: MySQLConnectionAbstract,
|
task: Union[str, ML_TASK] = ML_TASK.CLASSIFICATION,
|
model_name: Optional[str] = None,
|
):
|
"""
|
Instantiate _MyMLModelCommon.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-train.html
|
A full list of supported tasks can be found under "Common ML_TRAIN Options"
|
|
Args:
|
db_connection: MySQL database connection.
|
task: ML task type (default: "classification").
|
model_name: Name to register the model within MySQL (default: None).
|
|
Raises:
|
ValueError: If the schema name is not valid
|
DatabaseError:
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
|
Returns:
|
None
|
"""
|
self.db_connection = db_connection
|
self.task = ML_TASK.get_task_string(task)
|
self.schema_name = source_schema(db_connection)
|
|
with atomic_transaction(self.db_connection) as cursor:
|
execute_sql(cursor, "CALL sys.ML_CREATE_OR_UPGRADE_CATALOG();")
|
|
if model_name is None:
|
model_name = get_random_name(self._is_model_name_available)
|
|
self.model_var = f"{VAR_NAME_SPACE}.{model_name}"
|
self.model_var_score = f"{self.model_var}.score"
|
|
self.model_name = model_name
|
validate_name(model_name)
|
|
with atomic_transaction(self.db_connection) as cursor:
|
execute_sql(cursor, f"SET @{self.model_var} = %s;", params=(model_name,))
|
|
def _delete_model(self) -> bool:
|
"""
|
Deletes the model from the model catalog if present
|
|
Raises:
|
DatabaseError:
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
|
Returns:
|
Whether the model was deleted
|
"""
|
current_user = self._get_user()
|
|
qualified_model_catalog = f"ML_SCHEMA_{current_user}.MODEL_CATALOG"
|
delete_model = (
|
f"DELETE FROM {qualified_model_catalog} "
|
f"WHERE model_handle = @{self.model_var}"
|
)
|
|
with atomic_transaction(self.db_connection) as cursor:
|
execute_sql(cursor, delete_model)
|
return cursor.rowcount > 0
|
|
def _get_model_info(self, model_name: str) -> Optional[dict]:
|
"""
|
Retrieves the model info from the model_catalog
|
|
Args:
|
model_var: The model alias to retrieve
|
|
Returns:
|
The model info from the model_catalog (None if the model is not present in the catalog)
|
|
Raises:
|
DatabaseError:
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
"""
|
|
def process_col(elem: Any) -> Any:
|
if isinstance(elem, str):
|
try:
|
elem = json.loads(elem)
|
except json.JSONDecodeError:
|
pass
|
return elem
|
|
current_user = self._get_user()
|
|
qualified_model_catalog = f"ML_SCHEMA_{current_user}.MODEL_CATALOG"
|
model_exists = (
|
f"SELECT * FROM {qualified_model_catalog} WHERE model_handle = %s"
|
)
|
|
with atomic_transaction(self.db_connection) as cursor:
|
execute_sql(cursor, model_exists, params=(model_name,))
|
model_info_df = sql_response_to_df(cursor)
|
|
if model_info_df.empty:
|
result = None
|
else:
|
unprocessed_result = model_info_df.to_json(orient="records")
|
unprocessed_result_json = json.loads(unprocessed_result)[0]
|
result = {
|
key: process_col(elem)
|
for key, elem in unprocessed_result_json.items()
|
}
|
|
return result
|
|
def get_model_info(self) -> Optional[dict]:
|
"""
|
Checks if the model name is available.
|
Model info is present in the catalog only if the model was previously fitted.
|
|
Returns:
|
True if the model name is not part of the model catalog
|
|
Raises:
|
DatabaseError:
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
"""
|
return self._get_model_info(self.model_name)
|
|
def _is_model_name_available(self, model_name: str) -> bool:
|
"""
|
Checks if the model name is available
|
|
Returns:
|
True if the model name is not part of the model catalog
|
|
Raises:
|
DatabaseError:
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
"""
|
return self._get_model_info(model_name) is None
|
|
def _load_model(self) -> None:
|
"""
|
Loads the model specified by `self.model_name` into MySQL.
|
After loading, the model is ready to handle ML operations.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-model-load.html
|
|
Raises:
|
DatabaseError:
|
If the model is not initialized, i.e., fit or import has not been called
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
|
Returns:
|
None
|
"""
|
with atomic_transaction(self.db_connection) as cursor:
|
load_model_query = f"CALL sys.ML_MODEL_LOAD(@{self.model_var}, NULL);"
|
execute_sql(cursor, load_model_query)
|
|
def _get_user(self) -> str:
|
"""
|
Fetch the current database user (without host).
|
|
Returns:
|
The username string associated with the connection.
|
|
Raises:
|
DatabaseError:
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
ValueError: If the user name includes unsupported characters
|
"""
|
with atomic_transaction(self.db_connection) as cursor:
|
cursor.execute("SELECT CURRENT_USER()")
|
current_user = cursor.fetchone()[0].split("@")[0]
|
|
return validate_name(current_user)
|
|
def explain_model(self) -> dict:
|
"""
|
Get model explanations, such as detailed feature importances.
|
|
Returns:
|
dict: Feature importances and model explainability data.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-model-explanations.html
|
|
Raises:
|
DatabaseError:
|
If the model is not initialized, i.e., fit or import has not been called
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
ValueError:
|
If the model does not exist in the model catalog.
|
Should only occur if model was not fitted or was deleted.
|
"""
|
self._load_model()
|
with atomic_transaction(self.db_connection) as cursor:
|
current_user = self._get_user()
|
|
qualified_model_catalog = f"ML_SCHEMA_{current_user}.MODEL_CATALOG"
|
explain_query = (
|
f"SELECT model_explanation FROM {qualified_model_catalog} "
|
f"WHERE model_handle = @{self.model_var}"
|
)
|
|
execute_sql(cursor, explain_query)
|
df = sql_response_to_df(cursor)
|
|
return df.iloc[0, 0]
|
|
def _fit(
|
self,
|
table_name: str,
|
target_column_name: Optional[str],
|
options: Optional[dict],
|
) -> None:
|
"""
|
Fit an ML model using a referenced SQL table and target column.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-train.html
|
A full list of supported options can be found under "Common ML_TRAIN Options"
|
|
Args:
|
table_name: Name of the training data table.
|
target_column_name: Name of the target/label column.
|
options: Additional fit/config options (may override defaults).
|
|
Raises:
|
DatabaseError:
|
If provided options are invalid or unsupported.
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
ValueError: If the table or target_column name is not valid
|
|
Returns:
|
None
|
"""
|
validate_name(table_name)
|
if target_column_name is not None:
|
validate_name(target_column_name)
|
target_col_string = f"'{target_column_name}'"
|
else:
|
target_col_string = "NULL"
|
|
if options is None:
|
options = {}
|
options = copy.deepcopy(options)
|
options["task"] = self.task
|
|
self._delete_model()
|
|
with atomic_transaction(self.db_connection) as cursor:
|
placeholders, parameters = format_value_sql(options)
|
execute_sql(
|
cursor,
|
(
|
"CALL sys.ML_TRAIN("
|
f"'{self.schema_name}.{table_name}', "
|
f"{target_col_string}, "
|
f"{placeholders}, "
|
f"@{self.model_var}"
|
")"
|
),
|
params=parameters,
|
)
|
|
def _predict(
|
self, table_name: str, output_table_name: str, options: Optional[dict]
|
) -> None:
|
"""
|
Predict on a given data table and write results to an output table.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-predict-table.html
|
A full list of supported options can be found under "ML_PREDICT_TABLE Options"
|
|
Args:
|
table_name: Name of the SQL table with input data.
|
output_table_name: Name for the SQL output table to contain predictions.
|
options: Optional prediction options.
|
|
Returns:
|
None
|
|
Raises:
|
DatabaseError:
|
If provided options are invalid or unsupported,
|
or if the model is not initialized, i.e., fit or import has not
|
been called
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
ValueError: If the table or output_table name is not valid
|
"""
|
validate_name(table_name)
|
validate_name(output_table_name)
|
|
self._load_model()
|
with atomic_transaction(self.db_connection) as cursor:
|
placeholders, parameters = format_value_sql(options)
|
execute_sql(
|
cursor,
|
(
|
"CALL sys.ML_PREDICT_TABLE("
|
f"'{self.schema_name}.{table_name}', "
|
f"@{self.model_var}, "
|
f"'{self.schema_name}.{output_table_name}', "
|
f"{placeholders}"
|
")"
|
),
|
params=parameters,
|
)
|
|
def _score(
|
self,
|
table_name: str,
|
target_column_name: str,
|
metric: str,
|
options: Optional[dict],
|
) -> float:
|
"""
|
Evaluate model performance with a scoring metric.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-score.html
|
A full list of supported options can be found under
|
"Options for Recommendation Models" and
|
"Options for Anomaly Detection Models"
|
|
Args:
|
table_name: Table with features and ground truth.
|
target_column_name: Column of true target labels.
|
metric: String name of the metric to compute.
|
options: Optional dictionary of further scoring options.
|
|
Returns:
|
float: Computed score from the ML system.
|
|
Raises:
|
DatabaseError:
|
If provided options are invalid or unsupported,
|
or if the model is not initialized, i.e., fit or import has not
|
been called
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
ValueError: If the table or target_column name or metric is not valid
|
"""
|
validate_name(table_name)
|
validate_name(target_column_name)
|
validate_name(metric)
|
|
self._load_model()
|
with atomic_transaction(self.db_connection) as cursor:
|
placeholders, parameters = format_value_sql(options)
|
execute_sql(
|
cursor,
|
(
|
"CALL sys.ML_SCORE("
|
f"'{self.schema_name}.{table_name}', "
|
f"'{target_column_name}', "
|
f"@{self.model_var}, "
|
"%s, "
|
f"@{self.model_var_score}, "
|
f"{placeholders}"
|
")"
|
),
|
params=[metric, *parameters],
|
)
|
execute_sql(cursor, f"SELECT @{self.model_var_score}")
|
df = sql_response_to_df(cursor)
|
|
return df.iloc[0, 0]
|
|
def _explain_predictions(
|
self, table_name: str, output_table_name: str, options: Optional[dict]
|
) -> pd.DataFrame:
|
"""
|
Produce explanations for model predictions on provided data.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-explain-table.html
|
A full list of supported options can be found under "ML_EXPLAIN_TABLE Options"
|
|
Args:
|
table_name: Name of the SQL table with input data.
|
output_table_name: Name for the SQL table to store explanations.
|
options: Optional dictionary (default:
|
{"prediction_explainer": "permutation_importance"}).
|
|
Returns:
|
DataFrame: Prediction explanations from the output SQL table.
|
|
Raises:
|
DatabaseError:
|
If provided options are invalid or unsupported,
|
or if the model is not initialized, i.e., fit or import has not
|
been called
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
ValueError: If the table or output_table name is not valid
|
"""
|
validate_name(table_name)
|
validate_name(output_table_name)
|
|
if options is None:
|
options = {"prediction_explainer": "permutation_importance"}
|
|
self._load_model()
|
|
with atomic_transaction(self.db_connection) as cursor:
|
placeholders, parameters = format_value_sql(options)
|
execute_sql(
|
cursor,
|
(
|
"CALL sys.ML_EXPLAIN_TABLE("
|
f"'{self.schema_name}.{table_name}', "
|
f"@{self.model_var}, "
|
f"'{self.schema_name}.{output_table_name}', "
|
f"{placeholders}"
|
")"
|
),
|
params=parameters,
|
)
|
execute_sql(cursor, f"SELECT * FROM {self.schema_name}.{output_table_name}")
|
df = sql_response_to_df(cursor)
|
|
return df
|
|
|
class MyModel(_MyModelCommon):
|
"""
|
Convenience class for managing the ML workflow using pandas DataFrames.
|
|
Methods convert in-memory DataFrames into temp SQL tables before delegating to the
|
_MyMLModelCommon routines, and automatically clean up temp resources.
|
"""
|
|
def fit(
|
self,
|
X: Union[pd.DataFrame, np.ndarray], # pylint: disable=invalid-name
|
y: Optional[Union[pd.DataFrame, np.ndarray]],
|
options: Optional[dict] = None,
|
) -> None:
|
"""
|
Fit a model using DataFrame inputs.
|
|
If an 'id' column is defined in either dataframe, it will be used as the primary key.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-train.html
|
A full list of supported options can be found under "Common ML_TRAIN Options"
|
|
Args:
|
X: Features DataFrame.
|
y: (Optional) Target labels DataFrame or Series. If None, only X is used.
|
options: Additional options to pass to training.
|
|
Returns:
|
None
|
|
Raises:
|
DatabaseError:
|
If provided options are invalid or unsupported.
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
|
Notes:
|
Combines X and y as necessary. Creates a temporary table in the schema for training,
|
and deletes it afterward.
|
"""
|
X, y = convert_to_df(X), convert_to_df(y)
|
|
with (
|
atomic_transaction(self.db_connection) as cursor,
|
temporary_sql_tables(self.db_connection) as temporary_tables,
|
):
|
if y is not None:
|
if isinstance(y, pd.DataFrame):
|
# keep column name if it exists
|
target_column_name = y.columns[0]
|
else:
|
target_column_name = get_random_name(
|
lambda name: name not in X.columns
|
)
|
|
if target_column_name in X.columns:
|
raise ValueError(
|
f"Target column y with name {target_column_name} already present "
|
"in feature dataframe X"
|
)
|
|
df_combined = X.copy()
|
df_combined[target_column_name] = y
|
final_df = df_combined
|
else:
|
target_column_name = None
|
final_df = X
|
|
_, table_name = sql_table_from_df(cursor, self.schema_name, final_df)
|
temporary_tables.append((self.schema_name, table_name))
|
|
self._fit(table_name, target_column_name, options)
|
|
def predict(
|
self,
|
X: Union[pd.DataFrame, np.ndarray], # pylint: disable=invalid-name
|
options: Optional[dict] = None,
|
) -> pd.DataFrame:
|
"""
|
Generate model predictions using DataFrame input.
|
|
If an 'id' column is defined in either dataframe, it will be used as the primary key.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-predict-table.html
|
A full list of supported options can be found under "ML_PREDICT_TABLE Options"
|
|
Args:
|
X: DataFrame containing prediction features (no labels).
|
options: Additional prediction settings.
|
|
Returns:
|
DataFrame with prediction results as returned by HeatWave.
|
|
Raises:
|
DatabaseError:
|
If provided options are invalid or unsupported,
|
or if the model is not initialized, i.e., fit or import has not
|
been called
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
|
Notes:
|
Temporary SQL tables are created and deleted for input/output.
|
"""
|
X = convert_to_df(X)
|
|
with (
|
atomic_transaction(self.db_connection) as cursor,
|
temporary_sql_tables(self.db_connection) as temporary_tables,
|
):
|
_, table_name = sql_table_from_df(cursor, self.schema_name, X)
|
temporary_tables.append((self.schema_name, table_name))
|
|
output_table_name = get_random_name(
|
lambda table_name: not table_exists(
|
cursor, self.schema_name, table_name
|
)
|
)
|
temporary_tables.append((self.schema_name, output_table_name))
|
|
self._predict(table_name, output_table_name, options)
|
predictions = sql_table_to_df(cursor, self.schema_name, output_table_name)
|
|
# ml_results is text but known to always follow JSON format
|
predictions["ml_results"] = predictions["ml_results"].map(json.loads)
|
|
return predictions
|
|
def score(
|
self,
|
X: Union[pd.DataFrame, np.ndarray], # pylint: disable=invalid-name
|
y: Union[pd.DataFrame, np.ndarray],
|
metric: str,
|
options: Optional[dict] = None,
|
) -> float:
|
"""
|
Score the model using X/y data and a selected metric.
|
|
If an 'id' column is defined in either dataframe, it will be used as the primary key.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-score.html
|
A full list of supported options can be found under
|
"Options for Recommendation Models" and
|
"Options for Anomaly Detection Models"
|
|
Args:
|
X: DataFrame of features.
|
y: DataFrame or Series of labels.
|
metric: Metric name (e.g., "balanced_accuracy").
|
options: Optional ml scoring options.
|
|
Raises:
|
DatabaseError:
|
If provided options are invalid or unsupported,
|
or if the model is not initialized, i.e., fit or import has not
|
been called
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
|
Returns:
|
float: Computed score.
|
"""
|
X, y = convert_to_df(X), convert_to_df(y)
|
|
with (
|
atomic_transaction(self.db_connection) as cursor,
|
temporary_sql_tables(self.db_connection) as temporary_tables,
|
):
|
target_column_name = get_random_name(lambda name: name not in X.columns)
|
df_combined = X.copy()
|
df_combined[target_column_name] = y
|
final_df = df_combined
|
|
_, table_name = sql_table_from_df(cursor, self.schema_name, final_df)
|
temporary_tables.append((self.schema_name, table_name))
|
|
score = self._score(table_name, target_column_name, metric, options)
|
|
return score
|
|
def explain_predictions(
|
self,
|
X: Union[pd.DataFrame, np.ndarray], # pylint: disable=invalid-name
|
options: Dict = None,
|
) -> pd.DataFrame:
|
"""
|
Explain model predictions using provided data.
|
|
If an 'id' column is defined in either dataframe, it will be used as the primary key.
|
|
References:
|
https://dev.mysql.com/doc/heatwave/en/mys-hwaml-ml-explain-table.html
|
A full list of supported options can be found under
|
"ML_EXPLAIN_TABLE Options"
|
|
Args:
|
X: DataFrame for which predictions should be explained.
|
options: Optional dictionary of explainability options.
|
|
Returns:
|
DataFrame containing explanation details (feature attributions, etc.)
|
|
Raises:
|
DatabaseError:
|
If provided options are invalid or unsupported, or if the model is not initialized,
|
i.e., fit or import has not been called
|
If a database connection issue occurs.
|
If an operational error occurs during execution.
|
|
Notes:
|
Temporary input/output tables are cleaned up after explanation.
|
"""
|
X = convert_to_df(X)
|
|
with (
|
atomic_transaction(self.db_connection) as cursor,
|
temporary_sql_tables(self.db_connection) as temporary_tables,
|
):
|
|
_, table_name = sql_table_from_df(cursor, self.schema_name, X)
|
temporary_tables.append((self.schema_name, table_name))
|
|
output_table_name = get_random_name(
|
lambda table_name: not table_exists(
|
cursor, self.schema_name, table_name
|
)
|
)
|
temporary_tables.append((self.schema_name, output_table_name))
|
|
explanations = self._explain_predictions(
|
table_name, output_table_name, options
|
)
|
|
return explanations
|
