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

wedata/feature_store/training_set_client/training_set_client.py

@@ -0,0 +1,196 @@
+import json
+import logging
+import os
+from collections import defaultdict
+from types import ModuleType
+from typing import Any, Dict, List, Optional, Set, Union
+
+import mlflow
+import yaml
+from mlflow.models import Model, ModelSignature
+from mlflow.utils.file_utils import TempDir, YamlSafeDumper, read_yaml
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import struct
+
+from feature_store.entities.feature_column_info import FeatureColumnInfo
+from feature_store.entities.feature_function import FeatureFunction
+from feature_store.entities.feature_lookup import FeatureLookup
+from feature_store.entities.feature_spec import FeatureSpec
+from feature_store.entities.training_set import TrainingSet
+from feature_store.spark_client.spark_client import SparkClient
+
+from feature_store.constants.constants import (
+    _NO_RESULT_TYPE_PASSED,
+    _PREBUILT_ENV_URI,
+    _USE_SPARK_NATIVE_JOIN,
+    _WARN,
+    MODEL_DATA_PATH_ROOT,
+    PREDICTION_COLUMN_NAME,
+)
+
+from feature_store.utils import common_utils, training_set_utils
+from feature_store.utils.feature_spec_utils import convert_to_yaml_string
+
+_logger = logging.getLogger(__name__)
+
+FEATURE_SPEC_GRAPH_MAX_COLUMN_INFO = 1000
+
+
+class TrainingSetClient:
+    def __init__(
+        self,
+        spark_client: SparkClient
+    ):
+        self._spark_client = spark_client
+
+    def create_training_set(
+        self,
+        feature_spec: FeatureSpec,
+        feature_column_infos: List[FeatureColumnInfo],
+        label_names: List[str],
+        df: DataFrame,
+        ft_metadata: training_set_utils._FeatureTableMetadata,
+        kwargs,
+    ):
+        uc_function_infos = training_set_utils.get_uc_function_infos(
+            self._spark_client,
+            {odci.udf_name for odci in feature_spec.on_demand_column_infos},
+        )
+
+        # TODO(divyagupta-db): Move validation from _validate_join_feature_data in feature_lookup_utils.py
+        #  to a helper function called here and in score_batch.
+
+        # Add consumer of each feature and instrument as final step
+        consumer_feature_table_map = defaultdict(list)
+        for feature in feature_column_infos:
+            consumer_feature_table_map[feature.table_name].append(feature.feature_name)
+        consumed_udf_names = [f.udf_name for f in feature_spec.function_infos]
+
+        # Spark query planning is known to cause spark driver to crash if there are many feature tables to PiT join.
+        # See https://docs.google.com/document/d/1EyA4vvlWikTJMeinsLkxmRAVNlXoF1eqoZElOdqlWyY/edit
+        # So we disable native join by default.
+        training_set_utils.warn_if_non_photon_for_native_spark(
+            kwargs.get(_USE_SPARK_NATIVE_JOIN, False), self._spark_client
+        )
+        return TrainingSet(
+            feature_spec,
+            df,
+            label_names,
+            ft_metadata.feature_table_metadata_map,
+            ft_metadata.feature_table_data_map,
+            uc_function_infos,
+            kwargs.get(_USE_SPARK_NATIVE_JOIN, False),
+        )
+
+    def create_training_set_from_feature_lookups(
+        self,
+        df: DataFrame,
+        feature_lookups: List[Union[FeatureLookup, FeatureFunction]],
+        label: Union[str, List[str], None],
+        exclude_columns: List[str],
+        **kwargs,
+    ) -> TrainingSet:
+
+        # 获取特征查找列表和特征函数列表
+        features = feature_lookups
+        feature_lookups = [f for f in features if isinstance(f, FeatureLookup)]
+        feature_functions = [f for f in features if isinstance(f, FeatureFunction)]
+
+        # 如果未提供标签，则用空列表初始化label_names
+        label_names = common_utils.as_list(label, [])
+        del label
+
+        # 校验数据集和标签
+        training_set_utils.verify_df_and_labels(df, label_names, exclude_columns)
+
+        # 获取特征表元数据
+        ft_metadata = training_set_utils.get_table_metadata(
+            self._spark_client,
+            {fl.table_name for fl in feature_lookups}
+        )
+
+        column_infos = training_set_utils.get_column_infos(
+            feature_lookups,
+            feature_functions,
+            ft_metadata,
+            df_columns=df.columns,
+            label_names=label_names,
+        )
+
+        training_set_utils.validate_column_infos(
+            self._spark_client,
+            ft_metadata,
+            column_infos.source_data_column_infos,
+            column_infos.feature_column_infos,
+            column_infos.on_demand_column_infos,
+            label_names,
+        )
+
+        # Build feature_spec locally for comparison with the feature spec yaml generated by the
+        # FeatureStore backend. This will be removed once the migration is validated.
+        feature_spec = training_set_utils.build_feature_spec(
+            feature_lookups,
+            ft_metadata,
+            column_infos,
+            exclude_columns
+        )
+
+        return self.create_training_set(
+            feature_spec,
+            column_infos.feature_column_infos,
+            label_names,
+            df,
+            ft_metadata,
+            kwargs=kwargs,
+        )
+
+
+
+
+
+    def create_feature_spec(
+        self,
+        name: str,
+        features: List[Union[FeatureLookup, FeatureFunction]],
+        sparkClient: SparkClient,
+        exclude_columns: List[str] = [],
+    ) -> FeatureSpec:
+
+        feature_lookups = [f for f in features if isinstance(f, FeatureLookup)]
+        feature_functions = [f for f in features if isinstance(f, FeatureFunction)]
+
+        # Maximum of 100 FeatureFunctions is supported
+        if len(feature_functions) > training_set_utils.MAX_FEATURE_FUNCTIONS:
+            raise ValueError(
+                f"A maximum of {training_set_utils.MAX_FEATURE_FUNCTIONS} FeatureFunctions are supported."
+            )
+
+        # Get feature table metadata and column infos
+        ft_metadata = training_set_utils.get_table_metadata(
+            self._spark_client,
+            {fl.table_name for fl in feature_lookups}
+        )
+        column_infos = training_set_utils.get_column_infos(
+            feature_lookups,
+            feature_functions,
+            ft_metadata,
+        )
+
+        column_infos = training_set_utils.add_inferred_source_columns(column_infos)
+
+        training_set_utils.validate_column_infos(
+            self._spark_client,
+            ft_metadata,
+            column_infos.source_data_column_infos,
+            column_infos.feature_column_infos,
+            column_infos.on_demand_column_infos,
+        )
+
+        feature_spec = training_set_utils.build_feature_spec(
+            feature_lookups,
+            ft_metadata,
+            column_infos,
+            exclude_columns
+        )
+
+        return feature_spec

wedata/feature_store/utils/common_utils.py

@@ -0,0 +1,96 @@
+"""
+通用工具函数
+"""
+
+from collections import Counter
+from typing import Any, List
+
+from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository
+from mlflow.store.artifact.runs_artifact_repo import RunsArtifactRepository
+
+
+def is_artifact_uri(uri):
+    """
+    Checks the artifact URI is associated with a MLflow model or run.
+    The actual URI can be a model URI, model URI + subdirectory, or model URI + path to artifact file.
+    """
+    return ModelsArtifactRepository.is_models_uri(
+        uri
+    ) or RunsArtifactRepository.is_runs_uri(uri)
+
+def as_list(obj, default=None):
+    if not obj:
+        return default
+    elif isinstance(obj, list):
+        return obj
+    else:
+        return [obj]
+
+def get_duplicates(elements: List[Any]) -> List[Any]:
+    """
+    Returns duplicate elements in the order they first appear.
+    """
+    element_counts = Counter(elements)
+    duplicates = []
+    for e in element_counts.keys():
+        if element_counts[e] > 1:
+            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.
+    Passes single-quoted, comma delimited duplicates to the error template.
+    """
+    duplicate_strings = get_duplicates(strings)
+    if duplicate_strings:
+        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`".
+    Use this function to sanitize identifiers such as column names in SQL and PySpark.
+    """
+    return f"`{identifier.replace('`', '``')}`"
+
+
+def sanitize_identifiers(identifiers: List[str]):
+    """
+    Sanitize and wrap the identifiers in a list with backquotes.
+    """
+    return [sanitize_identifier(i) for i in identifiers]
+
+
+def sanitize_multi_level_name(multi_level_name: str):
+    """
+    Sanitize a multi-level name (such as an Unity Catalog table name) by sanitizing each segment
+    and joining the results. For example, "ca+t.fo`o.ba$r" becomes "`ca+t`.`fo``o`.`ba$r`".
+    """
+    segments = multi_level_name.split(".")
+    return ".".join(sanitize_identifiers(segments))
+
+
+def unsanitize_identifier(identifier: str):
+    """
+    Unsanitize an identifier. Useful when we get a possibly sanitized identifier from Spark or
+    somewhere else, but we need an unsanitized one.
+    Note: This function does not check the correctness of the identifier passed in. e.g. `foo``
+    is not a valid sanitized identifier. When given such invalid input, this function returns
+    invalid output.
+    """
+    if len(identifier) >= 2 and identifier[0] == "`" and identifier[-1] == "`":
+        return identifier[1:-1].replace("``", "`")
+    else:
+        return identifier
+
+
+# strings containing \ or ' can break sql statements, so escape them.
+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))