wedata-feature-engineering 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

wedata/feature_store/utils/common_utils.py

@@ -1,12 +1,217 @@
 """
 通用工具函数
 """
-
+import os
 from collections import Counter
-from typing import Any, List
+from datetime import datetime, timezone
+from functools import wraps
+from typing import Any, Dict, List, Optional
+from urllib.parse import urlparse
 
+import mlflow
+from mlflow.exceptions import RestException
+from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
 from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository
 from mlflow.store.artifact.runs_artifact_repo import RunsArtifactRepository
+from mlflow.utils import databricks_utils
+
+from wedata.feature_store.constants.constants import MODEL_DATA_PATH_ROOT
+
+
+def validate_table_name(name: str):
+    """
+    验证特征表名规范，仅支持单表名，不能包含点（如<catalog>.<schema>.<table>）
+
+    参数:
+        name: 要验证的表名
+
+    异常:
+        ValueError: 如果表名包含点或不符合规范
+    """
+    if not name or not isinstance(name, str):
+        raise ValueError("Table name must be a non-empty string")
+    if name.count('.') > 0:
+        raise ValueError("Feature table name only supports single table name, cannot contain dots (e.g. <catalog>.<schema>.<table>)")
+    if not name[0].isalpha():
+        raise ValueError("Table name must start with a letter")
+    if not all(c.isalnum() or c == '_' for c in name):
+        raise ValueError("Table name can only contain letters, numbers and underscores")
+
+
+def build_full_table_name(table_name: str) -> str:
+    """
+    构建完整的表名，格式化为`<database>.<table>`形式。
+
+    Args:
+        table_name: 输入的表名（可以是简化的表名或完整表名）。
+
+    Returns:
+        完整表名（`<database>.<table>`）。
+    """
+
+    # 从环境变量中获取当前主账户名称
+    owner_uin = os.environ.get("WEDATA_OWNER_UIN", "default")
+
+    # 如果owner_uin为空，则报错
+    if not owner_uin:
+        raise ValueError("WEDATA_OWNER_UIN environment variable is not set")
+
+    feature_store_database = f"{owner_uin}.{table_name}"
+
+    return feature_store_database
+
+
+def enable_if(condition):
+    """
+    A decorator that conditionally enables a function based on a condition.
+    If the condition is not truthy, calling the function raises a NotImplementedError.
+
+    :param condition: A callable that returns a truthy or falsy value.
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            if not condition():
+                raise NotImplementedError
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def is_empty(target: str):
+    return target is None or len(target.strip()) == 0
+
+
+class _NoDbutilsError(Exception):
+    pass
+
+
+def _get_dbutils():
+    try:
+        import IPython
+
+        ip_shell = IPython.get_ipython()
+        if ip_shell is None:
+            raise _NoDbutilsError
+        return ip_shell.ns_table["user_global"]["dbutils"]
+    except ImportError:
+        raise _NoDbutilsError
+    except KeyError:
+        raise _NoDbutilsError
+
+
+def utc_timestamp_ms_from_iso_datetime_string(date_string: str) -> int:
+    dt = datetime.fromisoformat(date_string)
+    utc_dt = dt.replace(tzinfo=timezone.utc)
+    return 1000 * utc_dt.timestamp()
+
+
+def pip_depependency_pinned_major_version(pip_package_name, major_version):
+    """
+    Generate a pip dependency string that is pinned to a major version, for example: "databricks-feature-lookup==0.*"
+    """
+    return f"{pip_package_name}=={major_version}.*"
+
+
+def add_mlflow_pip_depependency(conda_env, pip_package_name):
+    """
+    Add a new pip dependency to the conda environment taken from the raw MLflow model.
+    """
+    if pip_package_name is None or len(pip_package_name) == 0:
+        raise ValueError(
+            "Unexpected input: missing or empty pip_package_name parameter"
+        )
+
+    found_pip_dependency = False
+    if conda_env is not None:
+        for dep in conda_env["dependencies"]:
+            if isinstance(dep, dict) and "pip" in dep:
+                found_pip_dependency = True
+                pip_deps = dep["pip"]
+                if pip_package_name not in pip_deps:
+                    pip_deps.append(pip_package_name)
+        if "dependencies" in conda_env and not found_pip_dependency:
+            raise ValueError(
+                "Unexpected input: mlflow conda_env did not contain pip as a dependency"
+            )
+
+
+def download_model_artifacts(model_uri, dir):
+    """
+    Downloads model artifacts from model_uri to dir.
+    """
+    if not is_artifact_uri(model_uri):
+        raise ValueError(
+            f"Invalid model URI '{model_uri}'."
+            f"Use ``models:/model_name>/<version_number>`` or "
+            f"``runs:/<mlflow_run_id>/run-relative/path/to/model``."
+        )
+
+    try:
+        repo = get_artifact_repository(model_uri)
+    except RestException as e:
+        raise ValueError(f"The model at '{model_uri}' does not exist.", e)
+
+    artifact_path = os.path.join(mlflow.pyfunc.DATA, MODEL_DATA_PATH_ROOT)
+    if len(repo.list_artifacts(artifact_path)) == 0:
+        raise ValueError(
+            f"No suitable model found at '{model_uri}'. Either no model exists in this "
+            f"artifact location or an existing model was not packaged with Feature Store metadata. "
+            f"Only models logged by FeatureStoreClient.log_model can be used in inference."
+        )
+
+    return repo.download_artifacts(artifact_path="", dst_path=dir)
+
+
+def validate_params_non_empty(params: Dict[str, Any], expected_params: List[str]):
+    """
+    Validate that none of the expected parameters are empty.
+    """
+    for expected_param in expected_params:
+        if expected_param not in params:
+            raise ValueError(
+                f'Internal error: expected parameter "{expected_param}" not found in params dictionary'
+            )
+        param_value = params[expected_param]
+        if not param_value:
+            raise ValueError(f'Parameter "{expected_param}" cannot be empty')
+
+
+def is_in_databricks_job():
+    """
+    Overrides the behavior of the mlflow databricks_utils.is_in_databricks_job().
+    """
+    try:
+        return databricks_utils.get_job_id() is not None
+    except Exception:
+        return False
+
+
+def get_workspace_url() -> Optional[str]:
+    """
+    Overrides the behavior of the mlflow.utils.databricks_utils.get_workspace_url().
+    """
+    workspace_url = databricks_utils.get_workspace_url()
+    if workspace_url and not urlparse(workspace_url).scheme:
+        workspace_url = "https://" + workspace_url
+    return workspace_url
+
+
+def is_in_databricks_env():
+    """
+    Determine if we are running in a Databricks environment.
+    """
+    try:
+        return (
+                is_in_databricks_job()
+                or databricks_utils.is_in_databricks_notebook()
+                or databricks_utils.is_in_databricks_runtime()
+        )
+    except Exception:
+        return False
 
 
 def is_artifact_uri(uri):
@@ -18,6 +223,7 @@ def is_artifact_uri(uri):
         uri
     ) or RunsArtifactRepository.is_runs_uri(uri)
 
+
 def as_list(obj, default=None):
     if not obj:
         return default
@@ -26,6 +232,7 @@ def as_list(obj, default=None):
     else:
         return [obj]
 
+
 def get_duplicates(elements: List[Any]) -> List[Any]:
     """
     Returns duplicate elements in the order they first appear.
@@ -37,6 +244,7 @@ def get_duplicates(elements: List[Any]) -> List[Any]:
             duplicates.append(e)
     return duplicates
 
+
 def validate_strings_unique(strings: List[str], error_template: str):
     """
     Validates all strings are unique, otherwise raise ValueError with the error template and duplicates.
@@ -47,6 +255,7 @@ def validate_strings_unique(strings: List[str], error_template: str):
         duplicates_formatted = ", ".join([f"'{s}'" for s in duplicate_strings])
         raise ValueError(error_template.format(duplicates_formatted))
 
+
 def sanitize_identifier(identifier: str):
     """
     Sanitize and wrap an identifier with backquotes. For example, "a`b" becomes "`a``b`".
@@ -89,8 +298,9 @@ def unsanitize_identifier(identifier: str):
 def escape_sql_string(input_str: str) -> str:
     return input_str.replace("\\", "\\\\").replace("'", "\\'")
 
+
 def get_unique_list_order(elements: List[Any]) -> List[Any]:
     """
     Returns unique elements in the order they first appear.
     """
-    return list(dict.fromkeys(elements))
+    return list(dict.fromkeys(elements))

wedata/feature_store/utils/feature_lookup_utils.py

@@ -10,13 +10,13 @@ from pyspark.sql import DataFrame, Window
 from pyspark.sql import functions as F
 from pyspark.sql.functions import sum, unix_timestamp
 
-from feature_store.entities.environment_variables import BROADCAST_JOIN_THRESHOLD
-from feature_store.entities.feature_column_info import FeatureColumnInfo
-from feature_store.entities.feature_lookup import FeatureLookup
-from feature_store.entities.feature_spec import FeatureSpec
-from feature_store.entities.feature_table import FeatureTable
+from wedata.feature_store.entities.environment_variables import BROADCAST_JOIN_THRESHOLD
+from wedata.feature_store.entities.feature_column_info import FeatureColumnInfo
+from wedata.feature_store.entities.feature_lookup import FeatureLookup
+from wedata.feature_store.entities.feature_spec import FeatureSpec
+from wedata.feature_store.entities.feature_table import FeatureTable
 
-from feature_store.utils import common_utils, validation_utils, uc_utils
+from wedata.feature_store.utils import common_utils, validation_utils, uc_utils
 
 _logger = logging.getLogger(__name__)
 

wedata/feature_store/utils/feature_spec_utils.py

@@ -6,12 +6,12 @@ from typing import Dict, List, Set, Tuple, Type, Union
 import yaml
 from mlflow.utils.file_utils import YamlSafeDumper
 
-from feature_store.entities.column_info import ColumnInfo
-from feature_store.entities.feature_column_info import FeatureColumnInfo
-from feature_store.entities.feature_spec import FeatureSpec
-from feature_store.entities.on_demand_column_info import OnDemandColumnInfo
-from feature_store.entities.source_data_column_info import SourceDataColumnInfo
-from feature_store.utils.topological_sort import topological_sort
+from wedata.feature_store.entities.column_info import ColumnInfo
+from wedata.feature_store.entities.feature_column_info import FeatureColumnInfo
+from wedata.feature_store.entities.feature_spec import FeatureSpec
+from wedata.feature_store.entities.on_demand_column_info import OnDemandColumnInfo
+from wedata.feature_store.entities.source_data_column_info import SourceDataColumnInfo
+from wedata.feature_store.utils.topological_sort import topological_sort
 
 DEFAULT_GRAPH_DEPTH_LIMIT = 5
 

wedata/feature_store/utils/feature_utils.py

@@ -1,11 +1,11 @@
 import copy
 from typing import List, Union
 
-from feature_store.entities.feature_function import FeatureFunction
-from feature_store.entities.feature_lookup import FeatureLookup
-from feature_store.spark_client.spark_client import SparkClient
-from feature_store.utils import uc_utils
-from feature_store.utils.feature_lookup_utils import get_feature_lookups_with_full_table_names
+from wedata.feature_store.entities.feature_function import FeatureFunction
+from wedata.feature_store.entities.feature_lookup import FeatureLookup
+from wedata.feature_store.spark_client.spark_client import SparkClient
+from wedata.feature_store.utils import uc_utils
+from wedata.feature_store.utils.feature_lookup_utils import get_feature_lookups_with_full_table_names
 
 
 def format_feature_lookups_and_functions(

wedata/feature_store/utils/on_demand_utils.py

@@ -0,0 +1,107 @@
+import copy
+from typing import Dict, List
+
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import expr
+
+from wedata.feature_store.entities.feature_function import FeatureFunction
+from wedata.feature_store.entities.function_info import FunctionInfo
+from wedata.feature_store.entities.on_demand_column_info import OnDemandColumnInfo
+from wedata.feature_store.utils import common_utils, uc_utils
+
+
+def _udf_expr(udf_name: str, arguments: List[str]) -> expr:
+    """
+    Generate a Spark SQL expression, e.g. expr("udf_name(col1, col2)")
+    """
+    arguments_str = ", ".join(common_utils.sanitize_identifiers(arguments))
+    return expr(f"{udf_name}({arguments_str})")
+
+
+def _validate_apply_functions_df(
+    df: DataFrame,
+    functions_to_apply: List[OnDemandColumnInfo],
+    uc_function_infos: Dict[str, FunctionInfo],
+):
+    """
+    Validate the following:
+    1. On-demand input columns specified by functions_to_apply exist in the DataFrame.
+    2. On-demand input columns have data types that match those of UDF parameters.
+    """
+    for odci in functions_to_apply:
+        function_info = uc_function_infos[odci.udf_name]
+        types_dict = dict(df.dtypes)
+
+        for p in function_info.input_params:
+            arg_column = odci.input_bindings[p.name]
+            if arg_column not in df.columns:
+                raise ValueError(
+                    f"FeatureFunction argument column '{arg_column}' for UDF '{odci.udf_name}' parameter '{p.name}' "
+                    f"does not exist in provided DataFrame with schema '{df.schema}'."
+                )
+            if types_dict[arg_column] != p.type_text:
+                raise ValueError(
+                    f"FeatureFunction argument column '{arg_column}' for UDF '{odci.udf_name}' parameter '{p.name}' "
+                    f"does not have the expected type. Argument column '{arg_column}' has type "
+                    f"'{types_dict[arg_column]}' and parameter '{p.name}' has type '{p.type_text}'."
+                )
+
+
+def apply_functions_if_not_overridden(
+    df: DataFrame,
+    functions_to_apply: List[OnDemandColumnInfo],
+    uc_function_infos: Dict[str, FunctionInfo],
+) -> DataFrame:
+    """
+    For all on-demand features, in the order defined by the FeatureSpec:
+    If the feature does not already exist, append the evaluated UDF expression.
+    Existing column values or column positions are not modified.
+
+    `_validate_apply_functions_df` validates UDFs can be applied on `df` schema.
+
+    The caller should validate:
+    1. FeatureFunction bound argument columns for UDF parameters exist in FeatureSpec defined features.
+    2. FeatureFunction output feature names are unique.
+    """
+    _validate_apply_functions_df(
+        df=df,
+        functions_to_apply=functions_to_apply,
+        uc_function_infos=uc_function_infos,
+    )
+
+    columns = {}
+    for odci in functions_to_apply:
+        if odci.output_name not in df.columns:
+            function_info = uc_function_infos[odci.udf_name]
+            # Resolve the bound arguments in the UDF parameter order
+            udf_arguments = [
+                odci.input_bindings[p.name] for p in function_info.input_params
+            ]
+            columns[odci.output_name] = _udf_expr(odci.udf_name, udf_arguments)
+    return df.withColumns(columns)
+
+
+def get_feature_functions_with_full_udf_names(
+    feature_functions: List[FeatureFunction], current_catalog: str, current_schema: str
+):
+    """
+    Takes in a list of FeatureFunctions, and returns copies with:
+    1. Fully qualified UDF names.
+    2. If output_name is empty, fully qualified UDF names as output_name.
+    """
+    udf_names = {ff.udf_name for ff in feature_functions}
+    uc_utils._check_qualified_udf_names(udf_names)
+    uc_utils._verify_all_udfs_in_uc(udf_names, current_catalog, current_schema)
+
+    standardized_feature_functions = []
+    for ff in feature_functions:
+        ff_copy = copy.deepcopy(ff)
+        del ff
+
+        ff_copy._udf_name = uc_utils.get_full_udf_name(
+            ff_copy.udf_name, current_catalog, current_schema
+        )
+        if not ff_copy.output_name:
+            ff_copy._output_name = ff_copy.udf_name
+        standardized_feature_functions.append(ff_copy)
+    return standardized_feature_functions

wedata/feature_store/utils/schema_utils.py

@@ -1,6 +1,6 @@
 import logging
 
-from feature_store.constants.constants import _ERROR, _WARN
+from wedata.feature_store.constants.constants import _ERROR, _WARN
 
 _logger = logging.getLogger(__name__)
 

wedata/feature_store/utils/signature_utils.py

@@ -0,0 +1,205 @@
+import logging
+from typing import Any, Dict, Optional
+
+import mlflow
+from mlflow.models import ModelSignature
+from mlflow.types import ColSpec
+from mlflow.types import DataType as MlflowDataType
+from mlflow.types import ParamSchema, Schema
+
+from wedata.feature_store.entities.feature_column_info import FeatureColumnInfo
+from wedata.feature_store.entities.feature_spec import FeatureSpec
+from wedata.feature_store.entities.on_demand_column_info import OnDemandColumnInfo
+from wedata.feature_store.entities.source_data_column_info import SourceDataColumnInfo
+
+_logger = logging.getLogger(__name__)
+
+# Some types (array, map, decimal, timestamp_ntz) are unsupported due to MLflow signatures
+# lacking any equivalent types. We thus cannot construct a ColSpec for any column
+# that uses these types.
+SUPPORTED_TYPE_MAP = {
+    "smallint": MlflowDataType.integer,  # Upcast to integer
+    "int": MlflowDataType.integer,
+    "bigint": MlflowDataType.long,
+    "float": MlflowDataType.float,
+    "double": MlflowDataType.double,
+    "boolean": MlflowDataType.boolean,
+    "date": MlflowDataType.datetime,
+    "timestamp": MlflowDataType.datetime,
+    "string": MlflowDataType.string,
+    "binary": MlflowDataType.binary,
+}
+
+
+def is_unsupported_type(type_str: str):
+    return type_str not in SUPPORTED_TYPE_MAP
+
+
+def convert_spark_data_type_to_mlflow_signature_type(spark_type):
+    """
+    Maps Databricks SQL types to MLflow signature types.
+    docs.databricks.com/sql/language-manual/sql-ref-datatypes.html#language-mappings
+    """
+    return SUPPORTED_TYPE_MAP.get(spark_type)
+
+
+def get_input_schema_from_feature_spec(feature_spec: FeatureSpec) -> Schema:
+    """
+    Produces an MLflow signature schema from a feature spec.
+    Source data columns are marked as required inputs and feature columns
+    (both lookups and on-demand features) are marked as optional inputs.
+
+    :param feature_spec: FeatureSpec object with datatypes for each column.
+    """
+    # If we're missing any data types for any column, we are likely dealing with a
+    # malformed feature spec and should halt signature construction.
+    if any([ci.data_type is None for ci in feature_spec.column_infos]):
+        raise Exception("Training set does not contain column data types.")
+
+    source_data_cols = [
+        ci
+        for ci in feature_spec.column_infos
+        if isinstance(ci.info, SourceDataColumnInfo)
+    ]
+    # Don't create signature if any source data columns (required) are of complex types.
+    if any(
+        [
+            ci.data_type is None or is_unsupported_type(ci.data_type)
+            for ci in source_data_cols
+        ]
+    ):
+        raise Exception(
+            "Input DataFrame contains column data types not supported by "
+            "MLflow model signatures."
+        )
+    required_input_colspecs = [
+        ColSpec(
+            convert_spark_data_type_to_mlflow_signature_type(ci.data_type),
+            ci.info.output_name,
+            required=True,
+        )
+        for ci in source_data_cols
+    ]
+    feature_cols = [
+        ci
+        for ci in feature_spec.column_infos
+        if isinstance(ci.info, (FeatureColumnInfo, OnDemandColumnInfo))
+    ]
+    unsupported_feature_cols = [
+        ci for ci in feature_cols if is_unsupported_type(ci.data_type)
+    ]
+    optional_input_colspecs = [
+        ColSpec(
+            convert_spark_data_type_to_mlflow_signature_type(ci.data_type),
+            ci.output_name,
+            required=False,
+        )
+        for ci in feature_cols
+        if not is_unsupported_type(ci.data_type)
+    ]
+    if unsupported_feature_cols:
+        feat_string = ", ".join(
+            [f"{ci.output_name} ({ci.data_type})" for ci in unsupported_feature_cols]
+        )
+        _logger.warning(
+            f"The following features will not be included in the input schema because their"
+            f" data types are not supported by MLflow model signatures: {feat_string}. "
+            f"These features cannot be overridden during model serving."
+        )
+    return Schema(required_input_colspecs + optional_input_colspecs)
+
+
+def get_output_schema_from_labels(label_type_map: Optional[Dict[str, str]]) -> Schema:
+    """
+    Produces an MLflow signature schema from the provided label type map.
+    :param label_type_map: Map label column name -> data type
+    """
+    if not label_type_map:
+        raise Exception("Training set does not contain a label.")
+    if any([is_unsupported_type(dtype) for dtype in label_type_map.values()]):
+        raise Exception(
+            "Labels are of data types not supported by MLflow model signatures."
+        )
+    else:
+        output_colspecs = [
+            ColSpec(
+                convert_spark_data_type_to_mlflow_signature_type(spark_type),
+                col_name,
+                required=True,
+            )
+            for col_name, spark_type in label_type_map.items()
+        ]
+        return Schema(output_colspecs)
+
+
+def get_mlflow_signature_from_feature_spec(
+    feature_spec: FeatureSpec,
+    label_type_map: Optional[Dict[str, str]],
+    override_output_schema: Optional[Schema],
+    params: Optional[Dict[str, Any]] = None,
+) -> Optional[ModelSignature]:
+    """
+    Produce an MLflow signature from a feature spec and label type map.
+    Source data columns are marked as required inputs and feature columns
+    (both lookups and on-demand features) are marked as optional inputs.
+
+    Reads output types from the cached label -> datatype map in the training set.
+    If override_output_schema is provided, it will always be used as the output schema.
+
+    :param feature_spec: FeatureSpec object with datatypes for each column.
+    :param label_type_map: Map of label column name -> datatype
+    :param override_output_schema: User-provided output schema to use if provided.
+    """
+    kwargs = {}
+    kwargs["inputs"] = get_input_schema_from_feature_spec(feature_spec)
+    try:
+        output_schema = override_output_schema or get_output_schema_from_labels(
+            label_type_map
+        )
+        kwargs["outputs"] = output_schema
+    except Exception as e:
+        _logger.warning(f"Could not infer an output schema: {e}")
+
+    if params:
+        try:
+            from mlflow.types.utils import _infer_param_schema
+
+            kwargs["params"] = _infer_param_schema(params)
+        except Exception as e:
+            _logger.warning(f"Could not infer params schema: {e}")
+
+    return mlflow.models.ModelSignature(**kwargs)
+
+
+def drop_signature_inputs_and_invalid_params(signature):
+    """
+    Drop ModelSignature inputs field and invalid params from params field.
+    This is useful for feature store model's raw_model.
+    Feature store model's input schema does not apply to raw_model's input,
+    so we drop the inputs field of raw_model's signature.
+    Feature store model's result_type param enables setting and overriding
+    a default result_type for predictions, but this interferes with params
+    passed to MLflow's predict function, so we drop result_type from
+    the params field of raw_model's signature.
+
+    :param signature: ModelSignature object.
+    """
+    if signature:
+        outputs_schema = signature.outputs
+        params_schema = signature.params if hasattr(signature, "params") else None
+        try:
+            # Only for mlflow>=2.6.0 ModelSignature contains params attribute
+            if params_schema:
+                updated_params_schema = ParamSchema(
+                    [param for param in params_schema if param.name != "result_type"]
+                )
+                return ModelSignature(
+                    outputs=outputs_schema, params=updated_params_schema
+                )
+            if outputs_schema:
+                return ModelSignature(outputs=outputs_schema)
+        except TypeError:
+            _logger.warning(
+                "ModelSignature without inputs is not supported, please upgrade "
+                "mlflow >= 2.7.0 to use the feature."
+            )