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

wedata/tempo/io.py

@@ -1,61 +0,0 @@
-from __future__ import annotations
-
-import logging
-from collections import deque
-from typing import Optional
-
-import pyspark.sql.functions as sfn
-import wedata.tempo.tsdf as t_tsdf
-from pyspark.sql import SparkSession
-from pyspark.sql.utils import ParseException
-
-logger = logging.getLogger(__name__)
-
-
-def write(
-    tsdf: t_tsdf.TSDF,
-    spark: SparkSession,
-    tabName: str,
-    optimizationCols: Optional[list[str]] = None,
-) -> None:
-    """
-    param: tsdf: input TSDF object to write
-    param: tabName Delta output table name
-    param: optimizationCols list of columns to optimize on (time)
-    """
-    # hilbert curves more evenly distribute performance for querying multiple columns for Delta tables
-    spark.conf.set("spark.wedata.io.skipping.mdc.curve", "hilbert")
-
-    df = tsdf.df
-    ts_col = tsdf.ts_col
-    partitionCols = tsdf.partitionCols
-
-    view_df = df.withColumn("event_dt", sfn.to_date(sfn.col(ts_col))).withColumn(
-        "event_time",
-        sfn.translate(sfn.split(sfn.col(ts_col).cast("string"), " ")[1], ":", "").cast(
-            "double"
-        ),
-    )
-    view_cols = deque(view_df.columns)
-    view_cols.rotate(1)
-    view_df = view_df.select(*list(view_cols))
-
-    # Use replaceWhere instead of overwrite mode
-    writer = view_df.write.format("delta").partitionBy("event_dt")
-    writer = writer.option("replaceWhere", "true")
-    writer.saveAsTable(tabName)
-
-    if optimizationCols:
-        try:
-            spark.sql(
-                "optimize {} zorder by {}".format(
-                    tabName,
-                    "(" + ",".join(partitionCols + optimizationCols + [ts_col]) + ")",
-                )
-            )
-        except ParseException as e:
-            logger.error(
-                "Delta optimizations attempted, but was not successful.\nError: {}".format(
-                    e
-                )
-            )

wedata/tempo/ml.py

@@ -1,129 +0,0 @@
-from typing import Any, List, Tuple
-from functools import reduce
-
-from pyspark.sql import DataFrame
-from pyspark.sql.window import Window, WindowSpec
-from pyspark.sql import functions as sfn
-
-from pyspark.ml.param import Param, Params, TypeConverters
-from pyspark.ml.tuning import CrossValidator
-
-
-TMP_SPLIT_COL = "__tmp_split_col"
-TMP_GAP_COL = "__tmp_gap_row"
-
-
-def _get_parent_params():
-    dummy = Params()
-    dummy.uid = "undefined"
-    return dummy
-
-
-class TimeSeriesCrossValidator(CrossValidator):
-    # some additional parameters
-
-    timeSeriesCol: Param[str] = Param(
-        _get_parent_params(),
-        "timeSeriesCol",
-        "The name of the time series column",
-        typeConverter=TypeConverters.toString,
-    )
-    seriesIdCols: Param[List[str]] = Param(
-        _get_parent_params(),
-        "seriesIdCols",
-        "The name of the series id columns",
-        typeConverter=TypeConverters.toListString,
-    )
-    gap: Param[int] = Param(
-        _get_parent_params(),
-        "gap",
-        "The gap between training and test set",
-        typeConverter=TypeConverters.toInt,
-    )
-
-    def __init__(
-        self,
-        timeSeriesCol: str = "event_ts",
-        seriesIdCols: List[str] = [],
-        gap: int = 0,
-        **other_kwargs: Any,
-    ) -> None:
-        super(TimeSeriesCrossValidator, self).__init__(**other_kwargs)
-        self._setDefault(timeSeriesCol="event_ts", seriesIdCols=[], gap=0)
-        self._set(timeSeriesCol=timeSeriesCol, seriesIdCols=seriesIdCols, gap=gap)
-
-    def getTimeSeriesCol(self) -> str:
-        return self.getOrDefault(self.timeSeriesCol)
-
-    def getSeriesIdCols(self) -> List[str]:
-        return self.getOrDefault(self.seriesIdCols)
-
-    def getGap(self) -> int:
-        return self.getOrDefault(self.gap)
-
-    def setTimeSeriesCol(self, value: str) -> "TimeSeriesCrossValidator":
-        return self._set(timeSeriesCol=value)
-
-    def setSeriesIdCols(self, value: List[str]) -> "TimeSeriesCrossValidator":
-        return self._set(seriesIdCols=value)
-
-    def setGap(self, value: int) -> "TimeSeriesCrossValidator":
-        return self._set(gap=value)
-
-    def _get_split_win(self, desc: bool = False) -> WindowSpec:
-        ts_col_expr = sfn.col(self.getTimeSeriesCol())
-        if desc:
-            ts_col_expr = ts_col_expr.desc()
-        win = Window.orderBy(ts_col_expr)
-        series_id_cols = self.getSeriesIdCols()
-        if series_id_cols and len(series_id_cols) > 0:
-            win = win.partitionBy(*series_id_cols)
-        return win
-
-    def _kFold(self, dataset: DataFrame) -> List[Tuple[DataFrame, DataFrame]]:
-        nFolds = self.getOrDefault(self.numFolds)
-        nSplits = nFolds + 1
-
-        # split the data into nSplits subsets by timeseries order
-        split_df = dataset.withColumn(
-            TMP_SPLIT_COL, sfn.ntile(nSplits).over(self._get_split_win())
-        )
-        all_splits = [
-            split_df.filter(sfn.col(TMP_SPLIT_COL) == i).drop(TMP_SPLIT_COL)
-            for i in range(1, nSplits + 1)
-        ]
-        assert len(all_splits) == nSplits
-
-        # compose the k folds by including all previous splits in the training set,
-        # and the next split in the test set
-        kFolds = [
-            (reduce(lambda a, b: a.union(b), all_splits[: i + 1]), all_splits[i + 1])
-            for i in range(nFolds)
-        ]
-        assert len(kFolds) == nFolds
-        for tv in kFolds:
-            assert len(tv) == 2
-
-        # trim out a gap from the training datasets, if specified
-        gap = self.getOrDefault(self.gap)
-        if gap > 0:
-            order_cols = self.getSeriesIdCols() + [self.getTimeSeriesCol()]
-            # trim each training dataset by the specified gap
-            kFolds = [
-                (
-                    (
-                        train_df.withColumn(
-                            TMP_GAP_COL,
-                            sfn.row_number().over(self._get_split_win(desc=True)),
-                        )
-                        .where(sfn.col(TMP_GAP_COL) > gap)
-                        .drop(TMP_GAP_COL)
-                        .orderBy(*order_cols)
-                    ),
-                    test_df,
-                )
-                for (train_df, test_df) in kFolds
-            ]
-
-        # return the k folds (training, test) datasets
-        return kFolds

wedata/tempo/resample.py

@@ -1,318 +0,0 @@
-from __future__ import annotations
-
-from typing import (
-    Any,
-    Callable,
-    List,
-    Optional,
-    Tuple,
-    TypedDict,
-    Union,
-    get_type_hints,
-)
-
-import pyspark.sql.functions as sfn
-from pyspark.sql import DataFrame
-from pyspark.sql.window import Window
-
-import wedata.tempo.tsdf as t_tsdf
-
-# define global frequency options
-MUSEC = "microsec"
-MS = "ms"
-SEC = "sec"
-MIN = "min"
-HR = "hr"
-DAY = "day"
-
-# define global aggregate function options for downsampling
-floor = "floor"
-min = "min"
-max = "max"
-average = "mean"
-ceiling = "ceil"
-
-
-class FreqDict(TypedDict):
-    musec: str
-    microsec: str
-    microsecond: str
-    microseconds: str
-    ms: str
-    millisecond: str
-    milliseconds: str
-    sec: str
-    second: str
-    seconds: str
-    min: str
-    minute: str
-    minutes: str
-    hr: str
-    hour: str
-    hours: str
-    day: str
-    days: str
-
-
-freq_dict: FreqDict = {
-    "musec": "microseconds",
-    "microsec": "microseconds",
-    "microsecond": "microseconds",
-    "microseconds": "microseconds",
-    "ms": "milliseconds",
-    "millisecond": "milliseconds",
-    "milliseconds": "milliseconds",
-    "sec": "seconds",
-    "second": "seconds",
-    "seconds": "seconds",
-    "min": "minutes",
-    "minute": "minutes",
-    "minutes": "minutes",
-    "hr": "hours",
-    "hour": "hours",
-    "hours": "hours",
-    "day": "days",
-    "days": "days",
-}
-
-ALLOWED_FREQ_KEYS: List[str] = list(get_type_hints(FreqDict).keys())
-
-
-def is_valid_allowed_freq_keys(val: str, literal_constant: List[str]) -> bool:
-    return val in literal_constant
-
-
-allowableFreqs = [MUSEC, MS, SEC, MIN, HR, DAY]
-allowableFuncs = [floor, min, max, average, ceiling]
-
-
-def _appendAggKey(
-    tsdf: t_tsdf.TSDF, freq: Optional[str] = None
-) -> Tuple[t_tsdf.TSDF, int | str, Any]:
-    """
-    :param tsdf: TSDF object as input
-    :param freq: frequency at which to upsample
-    :return: triple - 1) return a TSDF with a new aggregate key (called agg_key) 2) return the period for use in interpolation, 3) return the time increment (also necessary for interpolation)
-    """
-    df = tsdf.df
-    parsed_freq = checkAllowableFreq(freq)
-    period, unit = parsed_freq[0], parsed_freq[1]
-
-    agg_window = sfn.window(
-        sfn.col(tsdf.ts_col), "{} {}".format(period, freq_dict[unit])  # type: ignore[literal-required]
-    )
-
-    df = df.withColumn("agg_key", agg_window)
-
-    return (
-        t_tsdf.TSDF(df, tsdf.ts_col, partition_cols=tsdf.partitionCols),
-        period,
-        freq_dict[unit],  # type: ignore[literal-required]
-    )
-
-
-def aggregate(
-    tsdf: t_tsdf.TSDF,
-    freq: str,
-    func: Union[Callable, str],
-    metricCols: Optional[List[str]] = None,
-    prefix: Optional[str] = None,
-    fill: Optional[bool] = None,
-) -> DataFrame:
-    """
-    aggregate a data frame by a coarser timestamp than the initial TSDF ts_col
-    :param tsdf: input TSDF object
-    :param func: aggregate function
-    :param metricCols: columns used for aggregates
-    :param prefix: the metric columns with the aggregate named function
-    :param fill: upsample based on the time increment for 0s in numeric columns
-    :return: TSDF object with newly aggregated timestamp as ts_col with aggregated values
-    """
-    tsdf, period, unit = _appendAggKey(tsdf, freq)
-
-    df = tsdf.df
-
-    groupingCols = tsdf.partitionCols + ["agg_key"]
-
-    if metricCols is None:
-        metricCols = list(set(df.columns).difference(set(groupingCols + [tsdf.ts_col])))
-
-    if prefix is None:
-        prefix = ""
-    else:
-        prefix = prefix + "_"
-
-    groupingCols = [sfn.col(column) for column in groupingCols]
-
-    if func == floor:
-        metricCol = sfn.struct([tsdf.ts_col] + metricCols)
-        res = df.withColumn("struct_cols", metricCol).groupBy(groupingCols)
-        res = res.agg(sfn.min("struct_cols").alias("closest_data")).select(
-            *groupingCols, sfn.col("closest_data.*")
-        )
-        new_cols = [sfn.col(tsdf.ts_col)] + [
-            sfn.col(c).alias("{}".format(prefix) + c) for c in metricCols
-        ]
-        res = res.select(*groupingCols, *new_cols)
-    elif func == average:
-        exprs = {x: "avg" for x in metricCols}
-        res = df.groupBy(groupingCols).agg(exprs)
-        agg_metric_cls = list(
-            set(res.columns).difference(
-                set(tsdf.partitionCols + [tsdf.ts_col, "agg_key"])
-            )
-        )
-        new_cols = [
-            sfn.col(c).alias(
-                "{}".format(prefix) + (c.split("avg(")[1]).replace(")", "")
-            )
-            for c in agg_metric_cls
-        ]
-        res = res.select(*groupingCols, *new_cols)
-    elif func == min:
-        exprs = {x: "min" for x in metricCols}
-        res = df.groupBy(groupingCols).agg(exprs)
-        agg_metric_cls = list(
-            set(res.columns).difference(
-                set(tsdf.partitionCols + [tsdf.ts_col, "agg_key"])
-            )
-        )
-        new_cols = [
-            sfn.col(c).alias(
-                "{}".format(prefix) + (c.split("min(")[1]).replace(")", "")
-            )
-            for c in agg_metric_cls
-        ]
-        res = res.select(*groupingCols, *new_cols)
-    elif func == max:
-        exprs = {x: "max" for x in metricCols}
-        res = df.groupBy(groupingCols).agg(exprs)
-        agg_metric_cls = list(
-            set(res.columns).difference(
-                set(tsdf.partitionCols + [tsdf.ts_col, "agg_key"])
-            )
-        )
-        new_cols = [
-            sfn.col(c).alias(
-                "{}".format(prefix) + (c.split("max(")[1]).replace(")", "")
-            )
-            for c in agg_metric_cls
-        ]
-        res = res.select(*groupingCols, *new_cols)
-    elif func == ceiling:
-        metricCol = sfn.struct([tsdf.ts_col] + metricCols)
-        res = df.withColumn("struct_cols", metricCol).groupBy(groupingCols)
-        res = res.agg(sfn.max("struct_cols").alias("ceil_data")).select(
-            *groupingCols, sfn.col("ceil_data.*")
-        )
-        new_cols = [sfn.col(tsdf.ts_col)] + [
-            sfn.col(c).alias("{}".format(prefix) + c) for c in metricCols
-        ]
-        res = res.select(*groupingCols, *new_cols)
-
-    # aggregate by the window and drop the end time (use start time as new ts_col)
-    res = (
-        res.drop(tsdf.ts_col)
-        .withColumnRenamed("agg_key", tsdf.ts_col)
-        .withColumn(tsdf.ts_col, sfn.col(tsdf.ts_col).start)
-    )
-
-    # sort columns so they are consistent
-    non_part_cols = set(set(res.columns) - set(tsdf.partitionCols)) - set([tsdf.ts_col])
-    sel_and_sort = tsdf.partitionCols + [tsdf.ts_col] + sorted(non_part_cols)
-    res = res.select(sel_and_sort)
-
-    fillW = Window.partitionBy(tsdf.partitionCols)
-
-    imputes = (
-        res.select(
-            *tsdf.partitionCols,
-            sfn.min(tsdf.ts_col).over(fillW).alias("from"),
-            sfn.max(tsdf.ts_col).over(fillW).alias("until"),
-        )
-        .distinct()
-        .withColumn(
-            tsdf.ts_col,
-            sfn.explode(
-                sfn.expr("sequence(from, until, interval {} {})".format(period, unit))
-            ),
-        )
-        .drop("from", "until")
-    )
-
-    metrics = []
-    for col in res.dtypes:
-        if col[1] in ["long", "double", "decimal", "integer", "float", "int"]:
-            metrics.append(col[0])
-
-    if fill:
-        res = imputes.join(
-            res, tsdf.partitionCols + [tsdf.ts_col], "leftouter"
-        ).na.fill(0, metrics)
-
-    return res
-
-
-def checkAllowableFreq(freq: Optional[str]) -> Tuple[Union[int | str], str]:
-    """
-    Parses frequency and checks against allowable frequencies
-    :param freq: frequncy at which to upsample/downsample, declared in resample function
-    :return: list of parsed frequency value and time suffix
-    """
-    if not isinstance(freq, str):
-        raise TypeError(f"Invalid type for `freq` argument: {freq}.")
-
-    # TODO - return either int OR str for first argument
-    allowable_freq: Tuple[Union[int | str], str] = (
-        0,
-        "will_always_fail_if_not_overwritten",
-    )
-
-    if is_valid_allowed_freq_keys(
-        freq.lower(),
-        ALLOWED_FREQ_KEYS,
-    ):
-        allowable_freq = 1, freq
-        return allowable_freq
-
-    try:
-        periods = freq.lower().split(" ")[0].strip()
-        units = freq.lower().split(" ")[1].strip()
-    except IndexError:
-        raise ValueError(
-            "Allowable grouping frequencies are microsecond (musec), millisecond (ms), sec (second), min (minute), hr (hour), day. Reformat your frequency as <integer> <day/hour/minute/second>"
-        )
-
-    if is_valid_allowed_freq_keys(
-        units.lower(),
-        ALLOWED_FREQ_KEYS,
-    ):
-        if units.startswith(MUSEC):
-            allowable_freq = periods, MUSEC
-        elif units.startswith(MS) | units.startswith("millis"):
-            allowable_freq = periods, MS
-        elif units.startswith(SEC):
-            allowable_freq = periods, SEC
-        elif units.startswith(MIN):
-            allowable_freq = periods, MIN
-        elif units.startswith("hour") | units.startswith(HR):
-            allowable_freq = periods, "hour"
-        elif units.startswith(DAY):
-            allowable_freq = periods, DAY
-    else:
-        raise ValueError(f"Invalid value for `freq` argument: {freq}.")
-
-    return allowable_freq
-
-
-def validateFuncExists(func: Union[Callable | str]) -> None:
-    if func is None:
-        raise TypeError(
-            "Aggregate function missing. Provide one of the allowable functions: "
-            + ", ".join(allowableFuncs)
-        )
-    elif func not in allowableFuncs:
-        raise ValueError(
-            "Aggregate function is not in the valid list. Provide one of the allowable functions: "
-            + ", ".join(allowableFuncs)
-        )