tencent-wedata-feature-engineering-dev 0.1.50__py3-none-any.whl → 0.2.0__py3-none-any.whl

wedata/feature_store/utils/env_utils.py

@@ -1,108 +0,0 @@
-import os
-
-
-class EnvironmentError(Exception):
-    pass
-
-
-def get_project_id() -> str:
-    """
-    获取当前项目名称
-    
-    Returns:
-        str: 项目ID
-        
-    Raises:
-        ValueError: 当环境变量 WEDATA_PROJECT_ID 未设置时
-    """
-    project_id = os.environ.get("WEDATA_PROJECT_ID")
-    if project_id:
-        return project_id
-    raise EnvironmentError("environment variable WEDATA_PROJECT_ID is not set, please check environment configuration")
-
-
-def get_cloud_secret() -> (str, str):
-    """
-    获取云上密钥
-
-    Returns:
-        tuple: 包含云上密钥的元组
-    """
-    secret_id = os.environ.get("WEDATA_CLOUD_TEMP_SECRET_ID")
-    secret_key = os.environ.get("WEDATA_CLOUD_TEMP_SECRET_KEY")
-    return secret_id, secret_key
-
-
-def get_region() -> str:
-    """
-    获取当前地域
-    """
-    region_dlc = os.environ.get("DLC_REGION")
-    region_emr = os.environ.get("KERNEL_REGION")
-    region = region_dlc if region_dlc else region_emr
-    if not region:
-        raise EnvironmentError("environment variable DLC_REGION or KERNEL_REGION is not set, "
-                               "please check environment configuration")
-    return region
-
-
-def get_database_name(database_name: str) -> str:
-    """
-    获取数据库名称
-
-    Args:
-        database_name: 数据库名称
-
-    Returns:
-        str: 数据库名称
-
-    Raises:
-        EnvironmentError: 当环境变量 WEDATA_DEFAULT_FEATURE_STORE_DATABASE 未设置时
-    """
-    feature_store_database_name = os.environ.get("WEDATA_DEFAULT_FEATURE_STORE_DATABASE")
-    if database_name:
-        return database_name
-    elif feature_store_database_name:
-        return feature_store_database_name
-    raise EnvironmentError("environment variable WEDATA_DEFAULT_FEATURE_STORE_DATABASE is not set, "
-                           "please check environment configuration")
-
-
-def get_engine_name() -> str:
-    """
-    获取引擎名称
-    """
-    # 因为DLC有特殊，所以先判断DLC，如果没有再判断EMR
-    if get_engine_type() == "DLC":
-        return _get_variable("KERNEL_ENGINE")
-    return _get_variable("KERNEL_ENGINE_NAME")
-
-
-def get_engine_id() -> str:
-    """
-    获取引擎ID
-    """
-    return _get_variable("KERNEL_ENGINE")
-
-
-def get_engine_type() -> str:
-    """
-    判断引擎类型
-    """
-    return "DLC" if os.environ.get("DLC_REGION") else "EMR"
-
-
-def get_feast_remote_url() -> str:
-    """
-    获取Feast远程URL
-    """
-    return _get_variable("KERNEL_FEAST_REMOTE_ADDRESS")
-
-
-def _get_variable(variable_key: str, is_raise: bool = True, default_value: str = None) -> str:
-    val = os.environ.get(variable_key, default_value)
-    if not val:
-        if is_raise:
-            raise EnvironmentError(f"environment variable {variable_key} is not set, "
-                                   f"please check environment configuration")
-    return val

wedata/tempo/interpol.py

@@ -1,448 +0,0 @@
-from __future__ import annotations
-
-from typing import Callable, List, Optional, Union
-
-from pyspark.sql.dataframe import DataFrame
-import pyspark.sql.functions as sfn
-from pyspark.sql.types import NumericType
-from pyspark.sql.window import Window
-
-import wedata.tempo.resample as t_resample
-import wedata.tempo.tsdf as t_tsdf
-import wedata.tempo.utils as t_utils
-
-# Interpolation fill options
-method_options = ["zero", "null", "bfill", "ffill", "linear"]
-
-
-class Interpolation:
-    def __init__(self, is_resampled: bool):
-        self.is_resampled = is_resampled
-
-    def __validate_fill(self, method: str) -> None:
-        """
-        Validate if the fill provided is within the allowed list of values.
-
-        :param fill: Fill type e.g. "zero", "null", "bfill", "ffill", "linear"
-        """
-        if method not in method_options:
-            raise ValueError(
-                f"Please select from one of the following fill options: {method_options}"
-            )
-
-    def __validate_col(
-        self,
-        df: DataFrame,
-        partition_cols: Optional[List[str]],
-        target_cols: List[str],
-        ts_col: str,
-        ts_col_dtype: Optional[str] = None,  # NB: added for testing purposes only
-    ) -> None:
-        """
-        Validate if target column exists and is of numeric type, and validates if partition column exists.
-
-        :param df: DataFrame to be validated
-        :param partition_cols: Partition columns to be validated
-        :param target_col: Target column to be validated
-        :param ts_col: Timestamp column to be validated
-        """
-
-        if partition_cols is not None:
-            for column in partition_cols:
-                if column not in str(df.columns):
-                    raise ValueError(
-                        f"Partition Column: '{column}' does not exist in DataFrame."
-                    )
-        for column in target_cols:
-            if column not in str(df.columns):
-                raise ValueError(
-                    f"Target Column: '{column}' does not exist in DataFrame."
-                )
-
-        if ts_col not in str(df.columns):
-            raise ValueError(
-                f"Timestamp Column: '{ts_col}' does not exist in DataFrame."
-            )
-
-        if ts_col_dtype is None:
-            ts_col_dtype = df.select(ts_col).dtypes[0][1]
-        if ts_col_dtype != "timestamp":
-            raise ValueError("Timestamp Column needs to be of timestamp type.")
-
-    def __calc_linear_spark(
-        self, df: DataFrame, ts_col: str, target_col: str
-    ) -> DataFrame:
-        """
-        Native Spark function for calculating linear interpolation on a DataFrame.
-
-        :param df: prepared dataframe to be interpolated
-        :param ts_col: timeseries column name
-        :param target_col: column to be interpolated
-        """
-        interpolation_expr = f"""
-        case when is_interpolated_{target_col} = false then {target_col}
-            when {target_col} is null then
-            (next_null_{target_col} - previous_{target_col})
-            /(unix_timestamp(next_timestamp_{target_col})-unix_timestamp(previous_timestamp_{target_col}))
-            *(unix_timestamp({ts_col}) - unix_timestamp(previous_timestamp_{target_col}))
-            + previous_{target_col}
-        else
-            (next_{target_col}-{target_col})
-            /(unix_timestamp(next_timestamp)-unix_timestamp(previous_timestamp))
-            *(unix_timestamp({ts_col}) - unix_timestamp(previous_timestamp))
-            + {target_col}
-        end as {target_col}
-        """
-
-        # remove target column to avoid duplication during interpolation expression
-        cols: List[str] = df.columns
-        cols.remove(target_col)
-        interpolated: DataFrame = df.selectExpr(*cols, interpolation_expr)
-        # Preserve column order
-        return interpolated.select(*df.columns)
-
-    def _is_valid_method_for_column(
-        self, series: DataFrame, method: str, col_name: str
-    ) -> bool:
-        """
-        zero and linear interpolation are only valid for numeric columns
-        """
-        if method in ["linear", "zero"]:
-            return isinstance(series.schema[col_name].dataType, NumericType)
-        else:
-            return True
-
-    def __interpolate_column(
-        self,
-        series: DataFrame,
-        ts_col: str,
-        target_col: str,
-        method: str,
-    ) -> DataFrame:
-        """
-        Apply interpolation to column.
-
-        :param series: input DataFrame
-        :param ts_col: timestamp column name
-        :param target_col: column to interpolate
-        :param method: interpolation function to fill missing values
-        """
-
-        if not self._is_valid_method_for_column(series, method, target_col):
-            raise ValueError(
-                f"Interpolation method '{method}' is not supported for column "
-                f"'{target_col}' of type '{series.schema[target_col].dataType}'. "
-                f"Only NumericType columns are supported."
-            )
-
-        output_df: DataFrame = series
-
-        # create new column for if target column is interpolated
-        flag_expr = f"""
-        CASE WHEN {target_col} is null and is_ts_interpolated = false THEN true
-             WHEN is_ts_interpolated = true THEN true
-        ELSE false
-        END AS is_interpolated_{target_col}
-        """
-        output_df = output_df.withColumn(
-            f"is_interpolated_{target_col}", sfn.expr(flag_expr)
-        )
-
-        # Handle zero fill
-        if method == "zero":
-            output_df = output_df.withColumn(
-                target_col,
-                sfn.when(
-                    sfn.col(f"is_interpolated_{target_col}") == False,  # noqa: E712
-                    sfn.col(target_col),
-                ).otherwise(sfn.lit(0)),
-            )
-
-        # Handle null fill
-        if method == "null":
-            output_df = output_df.withColumn(
-                target_col,
-                sfn.when(
-                    sfn.col(f"is_interpolated_{target_col}") == False,  # noqa: E712
-                    sfn.col(target_col),
-                ).otherwise(None),
-            )
-
-        # Handle forward fill
-        if method == "ffill":
-            output_df = output_df.withColumn(
-                target_col,
-                sfn.when(
-                    sfn.col(f"is_interpolated_{target_col}") == True,  # noqa: E712
-                    sfn.col(f"previous_{target_col}"),
-                ).otherwise(sfn.col(target_col)),
-            )
-        # Handle backwards fill
-        if method == "bfill":
-            output_df = output_df.withColumn(
-                target_col,
-                # Handle case when subsequent value is null
-                sfn.when(
-                    (sfn.col(f"is_interpolated_{target_col}") == True)  # noqa: E712
-                    & (
-                        sfn.col(f"next_{target_col}").isNull()
-                        & (sfn.col(f"{ts_col}_{target_col}").isNull())
-                    ),
-                    sfn.col(f"next_null_{target_col}"),
-                ).otherwise(
-                    # Handle standard backwards fill
-                    sfn.when(
-                        sfn.col(f"is_interpolated_{target_col}") == True,  # noqa: E712
-                        sfn.col(f"next_{target_col}"),
-                    ).otherwise(sfn.col(f"{target_col}"))
-                ),
-            )
-
-        # Handle linear fill
-        if method == "linear":
-            output_df = self.__calc_linear_spark(
-                output_df,
-                ts_col,
-                target_col,
-            )
-
-        return output_df
-
-    def __generate_time_series_fill(
-        self, df: DataFrame, partition_cols: Optional[List[str]], ts_col: str
-    ) -> DataFrame:
-        """
-        Create additional timeseries columns for previous and next timestamps
-
-        :param df: input DataFrame
-        :param partition_cols: partition column names
-        :param ts_col: timestamp column name
-        """
-        return df.withColumn(
-            "previous_timestamp",
-            sfn.col(ts_col),
-        ).withColumn(
-            "next_timestamp",
-            sfn.lead(df[ts_col]).over(
-                Window.partitionBy(*partition_cols).orderBy(ts_col)
-            ),
-        )
-
-    def __generate_column_time_fill(
-        self,
-        df: DataFrame,
-        partition_cols: Optional[List[str]],
-        ts_col: str,
-        target_col: str,
-    ) -> DataFrame:
-        """
-        Create timeseries columns for previous and next timestamps for a specific target column
-
-        :param df: input DataFrame
-        :param partition_cols: partition column names
-        :param ts_col: timestamp column name
-        :param target_col: target column name
-        """
-        window = Window
-        if partition_cols is not None:
-            window = Window.partitionBy(*partition_cols)
-
-        return df.withColumn(
-            f"previous_timestamp_{target_col}",
-            sfn.last(sfn.col(f"{ts_col}_{target_col}"), ignorenulls=True).over(
-                window.orderBy(ts_col).rowsBetween(Window.unboundedPreceding, 0)
-            ),
-        ).withColumn(
-            f"next_timestamp_{target_col}",
-            sfn.last(sfn.col(f"{ts_col}_{target_col}"), ignorenulls=True).over(
-                window.orderBy(sfn.col(ts_col).desc()).rowsBetween(
-                    Window.unboundedPreceding, 0
-                )
-            ),
-        )
-
-    def __generate_target_fill(
-        self,
-        df: DataFrame,
-        partition_cols: Optional[List[str]],
-        ts_col: str,
-        target_col: str,
-    ) -> DataFrame:
-        """
-        Create columns for previous and next value for a specific target column
-
-        :param df: input DataFrame
-        :param partition_cols: partition column names
-        :param ts_col: timestamp column name
-        :param target_col: target column name
-        """
-        window = Window
-
-        if partition_cols is not None:
-            window = Window.partitionBy(*partition_cols)
-        return (
-            df.withColumn(
-                f"previous_{target_col}",
-                sfn.last(df[target_col], ignorenulls=True).over(
-                    window.orderBy(ts_col).rowsBetween(Window.unboundedPreceding, 0)
-                ),
-            )
-            # Handle if subsequent value is null
-            .withColumn(
-                f"next_null_{target_col}",
-                sfn.last(df[target_col], ignorenulls=True).over(
-                    window.orderBy(sfn.col(ts_col).desc()).rowsBetween(
-                        Window.unboundedPreceding, 0
-                    )
-                ),
-            ).withColumn(
-                f"next_{target_col}",
-                sfn.lead(df[target_col]).over(window.orderBy(ts_col)),
-            )
-        )
-
-    def interpolate(
-        self,
-        tsdf: t_tsdf.TSDF,
-        ts_col: str,
-        partition_cols: Optional[List[str]],
-        target_cols: List[str],
-        freq: Optional[str],
-        func: Optional[Union[Callable | str]],
-        method: str,
-        show_interpolated: bool,
-        perform_checks: bool = True,
-    ) -> DataFrame:
-        """
-        Apply interpolation function.
-
-        :param tsdf: input TSDF
-        :param ts_col: timestamp column name
-        :param target_cols: numeric columns to interpolate
-        :param partition_cols: partition columns names
-        :param freq: frequency at which to sample
-        :param func: aggregate function used for sampling to the specified interval
-        :param method: interpolation function usded to fill missing values
-        :param show_interpolated: show if row is interpolated?
-        :param perform_checks: calculate time horizon and warnings if True (default is True)
-        :return: DataFrame containing interpolated data.
-        """
-        # Validate input parameters
-        self.__validate_fill(method)
-        self.__validate_col(tsdf.df, partition_cols, target_cols, ts_col)
-
-        if freq is None:
-            raise ValueError("freq cannot be None")
-
-        if func is None:
-            raise ValueError("func cannot be None")
-
-        if callable(func):
-            raise ValueError("func must be a string")
-
-        # Convert Frequency using resample dictionary
-        parsed_freq = t_resample.checkAllowableFreq(freq)
-        period, unit = parsed_freq[0], parsed_freq[1]
-        freq = f"{period} {t_resample.freq_dict[unit]}"  # type: ignore[literal-required]
-
-        # Throw warning for user to validate that the expected number of output rows is valid.
-        if perform_checks:
-            t_utils.calculate_time_horizon(tsdf.df, ts_col, freq, partition_cols)
-
-        # Only select required columns for interpolation
-        input_cols: List[str] = [ts_col, *target_cols]
-        if partition_cols is not None:
-            input_cols += [*partition_cols]
-
-        sampled_input: DataFrame = tsdf.df.select(*input_cols)
-
-        if self.is_resampled is False:
-            # Resample and Normalize Input
-            sampled_input = tsdf.resample(
-                freq=freq, func=func, metricCols=target_cols
-            ).df
-
-        # Fill timeseries for nearest values
-        time_series_filled = self.__generate_time_series_fill(
-            sampled_input, partition_cols, ts_col
-        )
-
-        # Generate surrogate timestamps for each target column
-        # This is required if multuple columns are being interpolated and may contain nulls
-        add_column_time: DataFrame = time_series_filled
-        for column in target_cols:
-            add_column_time = add_column_time.withColumn(
-                f"{ts_col}_{column}",
-                sfn.when(sfn.col(column).isNull(), None).otherwise(sfn.col(ts_col)),
-            )
-            add_column_time = self.__generate_column_time_fill(
-                add_column_time, partition_cols, ts_col, column
-            )
-
-        # Handle edge case if last value (latest) is null
-        edge_filled = add_column_time.withColumn(
-            "next_timestamp",
-            sfn.when(
-                sfn.col("next_timestamp").isNull(),
-                sfn.expr(f"{ts_col}+ interval {freq}"),
-            ).otherwise(sfn.col("next_timestamp")),
-        )
-
-        # Fill target column for nearest values
-        target_column_filled = edge_filled
-        for column in target_cols:
-            target_column_filled = self.__generate_target_fill(
-                target_column_filled, partition_cols, ts_col, column
-            )
-
-        # Generate missing timeseries values
-        exploded_series = target_column_filled.withColumn(
-            f"new_{ts_col}",
-            sfn.expr(
-                f"explode(sequence({ts_col}, next_timestamp - interval {freq}, interval {freq} )) as timestamp"
-            ),
-        )
-        # Mark rows that are interpolated if flag is set to True
-        flagged_series: DataFrame = exploded_series
-
-        flagged_series = (
-            exploded_series.withColumn(
-                "is_ts_interpolated",
-                sfn.when(sfn.col(f"new_{ts_col}") != sfn.col(ts_col), True).otherwise(
-                    False
-                ),
-            )
-            .withColumn(ts_col, sfn.col(f"new_{ts_col}"))
-            .drop(sfn.col(f"new_{ts_col}"))
-        )
-
-        # # Perform interpolation on each target column
-        interpolated_result: DataFrame = flagged_series
-        for target_col in target_cols:
-            # Interpolate target columns
-            interpolated_result = self.__interpolate_column(
-                interpolated_result, ts_col, target_col, method
-            )
-
-            interpolated_result = interpolated_result.drop(
-                f"previous_timestamp_{target_col}",
-                f"next_timestamp_{target_col}",
-                f"previous_{target_col}",
-                f"next_{target_col}",
-                f"next_null_{target_col}",
-                f"{ts_col}_{target_col}",
-            )
-
-        # Remove non-required columns
-        output: DataFrame = interpolated_result.drop(
-            "previous_timestamp", "next_timestamp"
-        )
-
-        # Hide is_interpolated columns based on flag
-        if show_interpolated is False:
-            interpolated_col_names = ["is_ts_interpolated"]
-            for column in target_cols:
-                interpolated_col_names.append(f"is_interpolated_{column}")
-            output = output.drop(*interpolated_col_names)
-
-        return output