tencent-wedata-feature-engineering-dev 0.1.0__py3-none-any.whl

wedata/feature_store/utils/feature_lookup_utils.py

@@ -0,0 +1,564 @@
+import copy
+import datetime
+import logging
+import re
+from collections import defaultdict
+from functools import reduce
+from typing import Dict, List, Optional, Tuple
+
+from pyspark.sql import DataFrame
+from pyspark.sql import functions as F
+import pyspark.sql.functions as psf
+
+
+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 wedata.feature_store.utils import uc_utils
+
+_logger = logging.getLogger(__name__)
+
+
+def _spark_asof_join_features(
+    df: DataFrame,
+    df_lookup_keys: List[str],
+    df_timestamp_lookup_key: str,
+    feature_table_data: DataFrame,
+    feature_table_keys: List[str],
+    feature_table_timestamp_key: str,
+    feature_to_output_name: Dict[str, str],
+    lookback_window_seconds: Optional[float] = None,
+    use_spark_native_join: Optional[bool] = False,
+) -> DataFrame:
+    # Alias feature table's keys to DataFrame lookup keys
+    ft_key_aliases = [
+        feature_table_data[ft_key].alias(df_key)
+        for (ft_key, df_key) in zip(feature_table_keys, df_lookup_keys)
+    ]
+    # Alias features to corresponding output names
+    ft_features = [
+        (feature_name, output_name)
+        for feature_name, output_name in feature_to_output_name.items()
+        # Skip join if feature it is already in DataFrame and therefore overridden
+        if output_name not in df.columns
+    ]
+    ft_feature_aliases = [
+        feature_table_data[feature_name].alias(output_name)
+        for feature_name, output_name in ft_features
+    ]
+    # Alias feature table's timestamp key to DataFrame timestamp lookup keys
+    ft_timestamp_key_aliases = [
+        feature_table_data[feature_table_timestamp_key].alias(df_timestamp_lookup_key)
+    ]
+    # Select key, timestamp key, and feature columns from feature table
+    feature_and_keys = feature_table_data.select(
+        ft_key_aliases + ft_timestamp_key_aliases + ft_feature_aliases
+    )
+
+    _logger.debug(
+        "Using native spark for point in time join"
+        if use_spark_native_join
+        else "Using tempo for point in time join"
+    )
+
+    if use_spark_native_join:
+        joined_df = _spark_asof_join_features_native(
+            labels_df=df,
+            features_df=feature_and_keys,
+            primary_keys=df_lookup_keys,
+            timestamp_key=df_timestamp_lookup_key,
+            lookback_window_seconds=lookback_window_seconds,
+        )
+    else:
+        joined_df = _spark_asof_join_features_tempo(
+            df=df,
+            feature_df=feature_and_keys,
+            lookup_keys=df_lookup_keys,
+            timestamp_key=df_timestamp_lookup_key,
+            lookback_window=lookback_window_seconds,
+            ft_features=ft_features,
+        )
+    return joined_df
+
+
+def _spark_asof_join_features_tempo(
+        df: DataFrame,
+        feature_df: DataFrame,
+        lookup_keys: List[str],
+        timestamp_key: str,
+        ft_features: List[Tuple[str, str]],
+        lookback_window: Optional[float] = None
+) -> DataFrame:
+    """
+    自定义实现as-of连接
+    :param df: 主表DataFrame
+    :param feature_df: 特征表DataFrame
+    :param lookup_keys: 连接键列表
+    :param timestamp_key: 时间戳列名
+    :param lookback_window: 最大回溯时间(秒)
+    :return: 连接后的DataFrame
+    """
+    from wedata.tempo.tsdf import TSDF
+    # 1. 只保留键列和时间戳列
+    df_tsdf = TSDF(df, ts_col=timestamp_key, partition_cols=lookup_keys)
+    ft_tsdf = TSDF(feature_df, ts_col=timestamp_key, partition_cols=lookup_keys)
+    # 进行as-of连接
+    joined_df = df_tsdf.asofJoin(
+        ft_tsdf,
+        left_prefix="left",
+        right_prefix="right",
+        skipNulls=False,
+        tolerance=lookback_window
+        if lookback_window is not None
+        else None,
+    ).df
+
+    # 去掉前缀，恢复列名
+    left_aliases = [
+        joined_df[f"left_{column_name}"].alias(column_name)
+        for column_name in df.columns
+        if column_name not in lookup_keys
+    ]
+    right_aliases = [
+        joined_df[f"right_{output_name}"].alias(output_name)
+        for (_, output_name) in ft_features
+    ]
+    return joined_df.select(lookup_keys + left_aliases + right_aliases)
+
+
+def _spark_asof_join_features_native(
+    labels_df: DataFrame,
+    features_df: DataFrame,
+    primary_keys: List[str],
+    timestamp_key: str,
+    lookback_window_seconds: Optional[float] = None,
+):
+    """
+    Performs an as-of join operation between two dataframes using native Spark operations.
+    Uses broadcast join for label dataset when within a size threshold to improve
+    efficiency of join operation with the assumption that size(labels_df) << size(features_df).
+    TODO(ML-40580): automatically switch labels_df and features_df based on size
+    The join operation is performed as follows:
+        1. Drop non-join key (primary and timestamp keys) columns from labels and features DataFrames
+        2. Broadcast join labels onto features DataFrame if within broadcast threshold.
+        3. Select maximum timestamp for each primary key
+        4. Rejoin non-primary key columns from features DataFrame to get features data
+        5. Rejoin non-primary key columns from labels DataFrame to get joint data
+
+    Parameters:
+    labels_df (DataFrame): The labels dataframe to join.
+    features_df (DataFrame): The features dataframe to join.
+    primary_keys (List[str]): The primary keys used for joining.
+    timestamp_key (str): The timestamp key used for joining.
+    lookback_window_seconds (Optional[float]): The lookback window in seconds.
+        If provided, the join operation will only consider records within this window.
+
+    Returns:
+    DataFrame: The result of the as-of join operation.
+    """
+    labels_df_keys_only = labels_df.select(
+        [F.col(key) for key in primary_keys] + [F.col(timestamp_key)]
+    )
+
+    # Broadcast labels DataFrame if within the broadcast threshold
+    if _df_in_size_threshold(labels_df_keys_only, BROADCAST_JOIN_THRESHOLD.get()):
+        labels_df_keys_only = F.broadcast(labels_df_keys_only)
+
+    # Drop non-primary key columns from features DataFrame
+    features_df_keys_only = features_df.select(
+        [F.col(key).alias(f"__features_pk_{key}") for key in primary_keys]
+        + [F.col(timestamp_key).alias("__features_tk")]
+    )
+
+    # Create join conditions
+    join_conditions = [
+        labels_df_keys_only[key] == features_df_keys_only[f"__features_pk_{key}"]
+        for key in primary_keys
+    ]
+    join_conditions = reduce(lambda x, y: x & y, join_conditions)
+    join_conditions &= (
+        labels_df_keys_only[timestamp_key] >= features_df_keys_only["__features_tk"]
+    )
+    if lookback_window_seconds is not None:
+        join_conditions &= (
+            psf.unix_timestamp(labels_df_keys_only[timestamp_key])
+            - psf.unix_timestamp(features_df_keys_only["__features_tk"])
+        ) <= lookback_window_seconds
+
+    # Join labels and features DataFrames
+    labels_df_keys_with_features_keys = labels_df_keys_only.join(
+        features_df_keys_only, on=join_conditions, how="left"
+    )
+
+    # Find the features max timestamps for each primary keys and timestamp key in labels
+    labels_df_keys_with_features_keys = labels_df_keys_with_features_keys.groupBy(
+        [labels_df_keys_only[key] for key in primary_keys] + [F.col(timestamp_key)]
+    ).agg(F.max("__features_tk").alias("__max_ts"))
+
+    if _df_in_size_threshold(
+        labels_df_keys_with_features_keys, BROADCAST_JOIN_THRESHOLD.get()
+    ):
+        labels_df_keys_with_features_keys = F.broadcast(
+            labels_df_keys_with_features_keys
+        )
+
+    # Rejoin features DataFrame to get the features data
+    join_conditions = [
+        features_df[key] == labels_df_keys_with_features_keys[key]
+        for key in primary_keys
+    ]
+    join_conditions = reduce(lambda x, y: x & y, join_conditions)
+    join_conditions &= (
+        features_df[timestamp_key] == labels_df_keys_with_features_keys["__max_ts"]
+    )
+
+    features = features_df.join(
+        labels_df_keys_with_features_keys,
+        on=join_conditions,
+        how="inner",
+    )
+
+    pk_columns_to_drop = [
+        labels_df_keys_with_features_keys[key] for key in primary_keys
+    ]
+    features = features.drop(*pk_columns_to_drop).drop(
+        features_df[timestamp_key], labels_df_keys_with_features_keys["__max_ts"]
+    )
+    features = features.dropDuplicates(primary_keys + [timestamp_key])
+    # Rejoin labels DataFrame if columns were dropped
+    joint_df = labels_df.join(features, on=primary_keys + [timestamp_key], how="left")
+    return joint_df
+
+
+def _df_in_size_threshold(df, threshold) -> float:
+    # Default to within threshold if can not find
+    try:
+        num_bytes = _get_df_size_from_spark_plan(df)
+    except Exception as e:
+        num_bytes = 0
+    return num_bytes <= threshold
+
+
+def _get_df_size_from_spark_plan(df: DataFrame) -> float:
+    """
+    获取DataFrame的估算大小(字节数)
+    替代方案：直接从DataFrame的SparkSession获取执行计划信息
+
+    参数:
+        df: 要计算大小的Spark DataFrame
+
+    返回:
+        float: 估算的字节数
+
+    异常:
+        ValueError: 如果无法从执行计划中解析出大小信息
+    """
+    # 直接从DataFrame获取SparkSession
+    spark = df.sql_ctx.sparkSession
+
+    # 创建临时视图
+    df.createOrReplaceTempView("temp_view_for_size")
+
+    # 获取执行计划
+    plan = spark.sql("explain cost select * from temp_view_for_size").collect()[0][0]
+
+    # 解析大小信息
+    search_result = re.search(r"sizeInBytes=.*(['\)])", plan, re.MULTILINE)
+    if search_result is None:
+        raise ValueError("无法从Spark执行计划中获取sizeInBytes信息")
+
+    # 提取大小和单位
+    result = search_result.group(0).replace(")", "")
+    size, units = result.split("=")[1].split()
+
+    # 单位转换映射
+    units_map = {
+        "TiB": 1024**4,  # 太字节
+        "GiB": 1024**3,  # 吉字节
+        "MiB": 1024**2,  # 兆字节
+        "KiB": 1024,     # 千字节
+        "B": 1           # 字节(处理没有单位的情况)
+    }
+
+    # 清理单位字符串并转换
+    clean_units = units.rstrip(",")
+    return float(size) * units_map.get(clean_units, 1)  # 默认返回原始值
+
+
+def _spark_join_features(
+    df: DataFrame,
+    df_keys: List[str],
+    feature_table_data: DataFrame,
+    feature_table_keys: List[str],
+    feature_to_output_name: Dict[str, str],
+) -> DataFrame:
+    """
+    Helper to join `feature_name` from `feature_table_data` into `df`.
+
+    This join uses a temporary table that contains only the keys and feature
+    from the feature table. The temporary table aliases the keys to match
+    the lookup keys and the feature to match the output_name.
+
+    Aliasing the keys allows us to join on name instead of by column,
+    which prevents duplicate column names after the join.
+    (see: https://kb.databricks.com/data/join-two-dataframes-duplicated-columns.html)
+
+    The joined-in feature is guaranteed to be unique because FeatureSpec
+    columns must be unique and the join is skipped if the feature
+    already exists in the DataFrame.
+    """
+
+    # Alias feature table's keys to DataFrame lookup keys
+    ft_key_aliases = [
+        feature_table_data[ft_key].alias(df_key)
+        for (ft_key, df_key) in zip(feature_table_keys, df_keys)
+    ]
+    # Alias features to corresponding output names
+    ft_feature_aliases = [
+        feature_table_data[feature_name].alias(output_name)
+        for feature_name, output_name in feature_to_output_name.items()
+        # Skip join if feature it is already in DataFrame and therefore overridden
+        if output_name not in df.columns
+    ]
+    # Select key and feature columns from feature table
+    feature_and_keys = feature_table_data.select(ft_key_aliases + ft_feature_aliases)
+    # Join feature to feature table
+    return df.join(feature_and_keys, df_keys, how="left")
+
+
+def _validate_join_keys(
+    feature_column_info: FeatureColumnInfo,
+    df: DataFrame,
+    feature_table_metadata: FeatureTable,
+    feature_table_data: DataFrame,
+    is_timestamp_key: bool = False,
+):
+    join_error_phrase = (
+        f"Unable to join feature table '{feature_column_info.table_name}'"
+    )
+    feature_column_info_keys = (
+        feature_column_info.timestamp_lookup_key
+        if is_timestamp_key
+        else feature_column_info.lookup_key
+    )
+    feature_table_keys = (
+        feature_table_metadata.timestamp_keys
+        if is_timestamp_key
+        else feature_table_metadata.primary_keys
+    )
+
+    lookup_key_kind = "timestamp lookup key" if is_timestamp_key else "lookup key"
+    feature_table_key_kind = "timestamp key" if is_timestamp_key else "primary key"
+
+    # Validate df has necessary keys
+    missing_df_keys = list(
+        filter(lambda df_key: df_key not in df.columns, feature_column_info_keys)
+    )
+    if missing_df_keys:
+        missing_keys = ", ".join([f"'{key}'" for key in missing_df_keys])
+        raise ValueError(
+            f"{join_error_phrase} because {lookup_key_kind} {missing_keys} not found in DataFrame."
+        )
+    # Validate feature table has necessary keys
+    missing_ft_keys = list(
+        filter(
+            lambda ft_key: ft_key not in feature_table_data.columns, feature_table_keys
+        )
+    )
+    if missing_ft_keys:
+        missing_keys = ", ".join([f"'{key}'" for key in missing_ft_keys])
+        raise ValueError(
+            f"{join_error_phrase} because {feature_table_key_kind} {missing_keys} not found in feature table."
+        )
+
+    # Validate number of feature table keys matches number of df lookup keys
+    if len(feature_column_info_keys) != len(feature_table_keys):
+        raise ValueError(
+            f"{join_error_phrase} because "
+            f"number of {feature_table_key_kind}s ({feature_table_keys}) "
+            f"does not match "
+            f"number of {lookup_key_kind}s ({feature_column_info_keys})."
+        )
+
+    # Validate feature table keys match types of df keys. The number of keys is expected to be the same.
+    # for (df_key, ft_key) in zip(feature_column_info_keys, feature_table_keys):
+    #     df_key_type = DataType.from_spark_type(df.schema[df_key].dataType)
+    #     ft_key_type = DataType.from_spark_type(
+    #         feature_table_data.schema[ft_key].dataType
+    #     )
+    #     if df_key_type != ft_key_type:
+    #         raise ValueError(
+    #             f"{join_error_phrase} because {feature_table_key_kind} '{ft_key}' has type '{DataType.to_string(ft_key_type)}' "
+    #             f"but corresponding {lookup_key_kind} '{df_key}' has type '{DataType.to_string(df_key_type)}' in DataFrame."
+    #         )
+
+
+def _validate_join_feature_data(
+    df: DataFrame,
+    features_to_join: List[FeatureColumnInfo],
+    feature_table_metadata_map: Dict[str, FeatureTable],
+    feature_table_data_map: Dict[str, DataFrame],
+):
+    for feature_info in features_to_join:
+        feature_table_metadata = feature_table_metadata_map[feature_info.table_name]
+        feature_table_data = feature_table_data_map[feature_info.table_name]
+        # Validate feature table primary keys match length/type of df lookup keys
+        _validate_join_keys(
+            feature_info,
+            df,
+            feature_table_metadata,
+            feature_table_data,
+            is_timestamp_key=False,
+        )
+        # Validate feature table timestamp keys match length/type of df timestamp lookup keys
+        _validate_join_keys(
+            feature_info,
+            df,
+            feature_table_metadata,
+            feature_table_data,
+            is_timestamp_key=True,
+        )
+
+
+def join_feature_data_if_not_overridden(
+    feature_spec: FeatureSpec,
+    df: DataFrame,
+    features_to_join: List[FeatureColumnInfo],
+    feature_table_metadata_map: Dict[str, FeatureTable],
+    feature_table_data_map: Dict[str, DataFrame],
+    use_spark_native_join: Optional[bool] = False,
+) -> DataFrame:
+    """
+    Joins `df` with features specified by `feature_spec.feature_column_infos` if they do not already exist.
+
+    Return column order is df.columns + newly joined features. The newly joined feature order is not guaranteed to
+    match `feature_spec.feature_column_infos` as feature lookups are first grouped by table for efficiency.
+
+    Before joining, it checks that:
+    1. Feature table keys match length and types of `df` lookup keys specified by FeatureSpec
+    2. `df` contains lookup keys specified by FeatureSpec
+    3. Feature table timestamp lookup keys match length and types of `df` timestamp lookup keys if specified by FeatureSpec
+    4. `df` contains timestamp lookup keys if specified by FeatureSpec
+    """
+    _validate_join_feature_data(
+        df=df,
+        features_to_join=features_to_join,
+        feature_table_metadata_map=feature_table_metadata_map,
+        feature_table_data_map=feature_table_data_map,
+    )
+
+    # Helper class to group all unique combinations of feature table names and lookup keys.
+    # All features in each of these groups will be JOINed with the training df using a single JOIN.
+    class JoinDataKey:
+        def __init__(
+            self,
+            feature_table: str,
+            lookup_key: List[str],
+            timestamp_lookup_key: List[str],
+            lookback_window: Optional[datetime.timedelta] = None,
+        ):
+            self.feature_table = feature_table
+            self.lookup_key = lookup_key
+            self.timestamp_lookup_key = timestamp_lookup_key
+            self.lookback_window = lookback_window
+
+        def __hash__(self):
+            return (
+                hash(self.feature_table)
+                + hash(tuple(self.lookup_key))
+                + hash(tuple(self.timestamp_lookup_key))
+                + hash(self.lookback_window)
+            )
+
+        def __eq__(self, other):
+            return (
+                self.feature_table == other.feature_table
+                and self.lookup_key == other.lookup_key
+                and self.timestamp_lookup_key == other.timestamp_lookup_key
+                and self.lookback_window == other.lookback_window
+            )
+
+    # Iterate through the list of FeatureColumnInfo and group features by name of the
+    # feature table and lookup key(s) and timestamp lookup key(s)
+    table_join_data = defaultdict(dict)
+    lookback_windows = {
+        t.table_name: t.lookback_window for t in feature_spec.table_infos
+    }
+    for feature_info in features_to_join:
+        join_data_key = JoinDataKey(
+            feature_info.table_name,
+            feature_info.lookup_key,
+            feature_info.timestamp_lookup_key,
+            lookback_windows[feature_info.table_name],
+        )
+        table_join_data[join_data_key][
+            feature_info.feature_name
+        ] = feature_info.output_name
+
+    for join_data_key, feature_to_output_name in table_join_data.items():
+
+        feature_table_metadata = feature_table_metadata_map[join_data_key.feature_table]
+        feature_table_data = feature_table_data_map[join_data_key.feature_table]
+
+        if join_data_key.timestamp_lookup_key:
+            # If lookback window is set to 0, then perform exact join instead of asof join to get perf benefits.
+            if (
+                join_data_key.lookback_window is not None
+                and join_data_key.lookback_window == 0
+            ):
+                df = _spark_join_features(
+                    df=df,
+                    df_keys=join_data_key.lookup_key
+                    + join_data_key.timestamp_lookup_key,
+                    feature_table_data=feature_table_data,
+                    feature_table_keys=feature_table_metadata.primary_keys
+                    + feature_table_metadata.timestamp_keys,
+                    feature_to_output_name=feature_to_output_name,
+                )
+            else:
+                df = _spark_asof_join_features(
+                    df=df,
+                    df_lookup_keys=join_data_key.lookup_key,
+                    df_timestamp_lookup_key=join_data_key.timestamp_lookup_key[0],
+                    feature_table_data=feature_table_data,
+                    feature_table_keys=feature_table_metadata.primary_keys,
+                    feature_table_timestamp_key=feature_table_metadata.timestamp_keys[
+                        0
+                    ],
+                    feature_to_output_name=feature_to_output_name,
+                    lookback_window_seconds=join_data_key.lookback_window,
+                    use_spark_native_join=use_spark_native_join,
+                )
+        else:
+            df = _spark_join_features(
+                df=df,
+                df_keys=join_data_key.lookup_key,
+                feature_table_data=feature_table_data,
+                feature_table_keys=feature_table_metadata.primary_keys,
+                feature_to_output_name=feature_to_output_name,
+            )
+    return df
+
+
+def get_feature_lookups_with_full_table_names(
+        feature_lookups: List[FeatureLookup], current_catalog: str, current_schema: str
+) -> List[FeatureLookup]:
+    """
+    Takes in a list of FeatureLookups, and returns copies with reformatted table names.
+    """
+    table_names = {fl.table_name for fl in feature_lookups}
+    uc_utils._check_qualified_table_names(table_names)
+    uc_utils._verify_all_tables_are_either_in_uc_or_in_hms(
+        table_names, current_catalog, current_schema
+    )
+    standardized_feature_lookups = []
+    for fl in feature_lookups:
+        fl_copy = copy.deepcopy(fl)
+        fl_copy._table_name = uc_utils.get_full_table_name(
+            fl_copy.table_name, current_catalog, current_schema
+        )
+        standardized_feature_lookups.append(fl_copy)
+    return standardized_feature_lookups