replay-rec 0.16.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

replay/utils/dataframe_bucketizer.py

@@ -9,39 +9,37 @@ if PYSPARK_AVAILABLE:
     from replay.utils.session_handler import State
 
 
-class DataframeBucketizer(
-    Transformer, DefaultParamsWritable, DefaultParamsReadable
-):  # pylint: disable=R0901
+class DataframeBucketizer(Transformer, DefaultParamsWritable, DefaultParamsReadable):
     """
     Buckets the input dataframe, dumps it to spark warehouse directory,
     and returns a bucketed dataframe.
     """
 
-    bucketingKey = Param(
+    bucketing_key = Param(
         Params._dummy(),
-        "bucketingKey",
+        "bucketing_key",
         "bucketing key (also used as sort key)",
         typeConverter=TypeConverters.toString,
     )
 
-    partitionNum = Param(
+    partition_num = Param(
         Params._dummy(),
-        "partitionNum",
+        "partition_num",
         "number of buckets",
         typeConverter=TypeConverters.toInt,
     )
 
-    tableName = Param(
+    table_name = Param(
         Params._dummy(),
-        "tableName",
+        "table_name",
         "parquet file name (for storage  in 'spark-warehouse') and spark table name",
         typeConverter=TypeConverters.toString,
     )
 
-    sparkWarehouseDir = Param(
+    spark_warehouse_dir = Param(
         Params._dummy(),
-        "sparkWarehouseDir",
-        "sparkWarehouseDir",
+        "spark_warehouse_dir",
+        "spark_warehouse_dir",
         typeConverter=TypeConverters.toString,
     )
 
@@ -62,10 +60,10 @@ class DataframeBucketizer(
                 i.e. value of 'spark.sql.warehouse.dir' property
         """
         super().__init__()
-        self.set(self.bucketingKey, bucketing_key)
-        self.set(self.partitionNum, partition_num)
-        self.set(self.tableName, table_name)
-        self.set(self.sparkWarehouseDir, spark_warehouse_dir)
+        self.set(self.bucketing_key, bucketing_key)
+        self.set(self.partition_num, partition_num)
+        self.set(self.table_name, table_name)
+        self.set(self.spark_warehouse_dir, spark_warehouse_dir)
 
     def __enter__(self):
         return self
@@ -76,31 +74,27 @@ class DataframeBucketizer(
     def remove_parquet(self):
         """Removes parquets where bucketed dataset is stored"""
         spark = State().session
-        spark_warehouse_dir = self.getOrDefault(self.sparkWarehouseDir)
-        table_name = self.getOrDefault(self.tableName)
-        fs = get_fs(spark)  # pylint: disable=invalid-name
-        fs_path = spark._jvm.org.apache.hadoop.fs.Path(
-            f"{spark_warehouse_dir}/{table_name}"
-        )
+        spark_warehouse_dir = self.getOrDefault(self.spark_warehouse_dir)
+        table_name = self.getOrDefault(self.table_name)
+        fs = get_fs(spark)
+        fs_path = spark._jvm.org.apache.hadoop.fs.Path(f"{spark_warehouse_dir}/{table_name}")
         is_exists = fs.exists(fs_path)
         if is_exists:
             fs.delete(fs_path, True)
 
     def set_table_name(self, table_name: str):
         """Sets table name"""
-        self.set(self.tableName, table_name)
+        self.set(self.table_name, table_name)
 
     def _transform(self, dataset: SparkDataFrame):
-        bucketing_key = self.getOrDefault(self.bucketingKey)
-        partition_num = self.getOrDefault(self.partitionNum)
-        table_name = self.getOrDefault(self.tableName)
-        spark_warehouse_dir = self.getOrDefault(self.sparkWarehouseDir)
+        bucketing_key = self.getOrDefault(self.bucketing_key)
+        partition_num = self.getOrDefault(self.partition_num)
+        table_name = self.getOrDefault(self.table_name)
+        spark_warehouse_dir = self.getOrDefault(self.spark_warehouse_dir)
 
         if not table_name:
-            raise ValueError(
-                "Parameter 'table_name' is not set! "
-                "Please set it via method 'set_table_name'."
-            )
+            msg = "Parameter 'table_name' is not set! Please set it via method 'set_table_name'."
+            raise ValueError(msg)
 
         (
             dataset.repartition(partition_num, bucketing_key)

replay/utils/distributions.py

@@ -22,23 +22,11 @@ def item_distribution(
     :return: DataFrame with results
     """
     log = convert2spark(log)
-    res = (
-        log.groupBy("item_idx")
-        .agg(sf.countDistinct("user_idx").alias("user_count"))
-        .select("item_idx", "user_count")
-    )
+    res = log.groupBy("item_idx").agg(sf.countDistinct("user_idx").alias("user_count")).select("item_idx", "user_count")
 
     rec = convert2spark(recommendations)
     rec = get_top_k_recs(rec, k)
-    rec = (
-        rec.groupBy("item_idx")
-        .agg(sf.countDistinct("user_idx").alias("rec_count"))
-        .select("item_idx", "rec_count")
-    )
+    rec = rec.groupBy("item_idx").agg(sf.countDistinct("user_idx").alias("rec_count")).select("item_idx", "rec_count")
 
-    res = (
-        res.join(rec, on="item_idx", how="outer")
-        .fillna(0)
-        .orderBy(["user_count", "item_idx"])
-    )
+    res = res.join(rec, on="item_idx", how="outer").fillna(0).orderBy(["user_count", "item_idx"])
     return spark_to_pandas(res, allow_collect_to_master)

replay/utils/model_handler.py

@@ -1,17 +1,18 @@
-# pylint: disable=wildcard-import,invalid-name,unused-wildcard-import,unspecified-encoding
+import functools
 import json
 import os
 import pickle
+import warnings
 from os.path import join
 from pathlib import Path
-from typing import Union
+from typing import Any, Callable, Optional, Union
 
 from replay.data.dataset_utils import DatasetLabelEncoder
 from replay.models import *
 from replay.models.base_rec import BaseRecommender
 from replay.splitters import *
-from .session_handler import State
 
+from .session_handler import State
 from .types import PYSPARK_AVAILABLE
 
 if PYSPARK_AVAILABLE:
@@ -26,9 +27,7 @@ if PYSPARK_AVAILABLE:
         :param spark: spark session
         :return:
         """
-        fs = spark._jvm.org.apache.hadoop.fs.FileSystem.get(
-            spark._jsc.hadoopConfiguration()
-        )
+        fs = spark._jvm.org.apache.hadoop.fs.FileSystem.get(spark._jsc.hadoopConfiguration())
         return fs
 
     def get_list_of_paths(spark: SparkSession, dir_path: str):
@@ -44,9 +43,7 @@ if PYSPARK_AVAILABLE:
         return [str(f.getPath()) for f in statuses]
 
 
-def save(
-    model: BaseRecommender, path: Union[str, Path], overwrite: bool = False
-):
+def save(model: BaseRecommender, path: Union[str, Path], overwrite: bool = False):
     """
     Save fitted model to disk as a folder
 
@@ -63,9 +60,8 @@ def save(
     if not overwrite:
         is_exists = fs.exists(spark._jvm.org.apache.hadoop.fs.Path(path))
         if is_exists:
-            raise FileExistsError(
-                f"Path '{path}' already exists. Mode is 'overwrite = False'."
-            )
+            msg = f"Path '{path}' already exists. Mode is 'overwrite = False'."
+            raise FileExistsError(msg)
 
     fs.mkdirs(spark._jvm.org.apache.hadoop.fs.Path(path))
     model._save_model(join(path, "model"))
@@ -74,9 +70,7 @@ def save(
     init_args["_model_name"] = str(model)
     sc = spark.sparkContext
     df = spark.read.json(sc.parallelize([json.dumps(init_args)]))
-    df.coalesce(1).write.mode("overwrite").option(
-        "ignoreNullFields", "false"
-    ).json(join(path, "init_args.json"))
+    df.coalesce(1).write.mode("overwrite").option("ignoreNullFields", "false").json(join(path, "init_args.json"))
 
     dataframes = model._dataframes
     df_path = join(path, "dataframes")
@@ -85,13 +79,9 @@ def save(
             df.write.mode("overwrite").parquet(join(df_path, name))
 
     if hasattr(model, "fit_queries"):
-        model.fit_queries.write.mode("overwrite").parquet(
-            join(df_path, "fit_queries")
-        )
+        model.fit_queries.write.mode("overwrite").parquet(join(df_path, "fit_queries"))
     if hasattr(model, "fit_items"):
-        model.fit_items.write.mode("overwrite").parquet(
-            join(df_path, "fit_items")
-        )
+        model.fit_items.write.mode("overwrite").parquet(join(df_path, "fit_items"))
     if hasattr(model, "study"):
         save_picklable_to_parquet(model.study, join(path, "study"))
 
@@ -104,18 +94,11 @@ def load(path: str, model_type=None) -> BaseRecommender:
     :return: Restored trained model
     """
     spark = State().session
-    args = (
-        spark.read.json(join(path, "init_args.json"))
-        .first()
-        .asDict(recursive=True)
-    )
+    args = spark.read.json(join(path, "init_args.json")).first().asDict(recursive=True)
     name = args["_model_name"]
     del args["_model_name"]
 
-    if model_type is not None:
-        model_class = model_type
-    else:
-        model_class = globals()[name]
+    model_class = model_type if model_type is not None else globals()[name]
 
     model = model_class(**args)
 
@@ -180,9 +163,7 @@ def save_splitter(splitter: Splitter, path: str, overwrite: bool = False):
     sc = spark.sparkContext
     df = spark.read.json(sc.parallelize([json.dumps(init_args)]))
     if overwrite:
-        df.coalesce(1).write.mode("overwrite").json(
-            join(path, "init_args.json")
-        )
+        df.coalesce(1).write.mode("overwrite").json(join(path, "init_args.json"))
     else:
         df.coalesce(1).write.json(join(path, "init_args.json"))
 
@@ -200,3 +181,25 @@ def load_splitter(path: str) -> Splitter:
     del args["_splitter_name"]
     splitter = globals()[name]
     return splitter(**args)
+
+
+def deprecation_warning(message: Optional[str] = None) -> Callable[..., Any]:
+    """
+    Decorator that throws deprecation warnings.
+
+    :param message: message to deprecation warning without func name.
+    """
+    base_msg = "will be deprecated in future versions."
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            msg = f"{func.__qualname__} {message if message else base_msg}"
+            warnings.simplefilter("always", DeprecationWarning)  # turn off filter
+            warnings.warn(msg, category=DeprecationWarning, stacklevel=2)
+            warnings.simplefilter("default", DeprecationWarning)  # reset filter
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

replay/utils/session_handler.py

@@ -36,7 +36,6 @@ def get_spark_session(
         Default: ``None``.
     """
     if os.environ.get("SCRIPT_ENV", None) == "cluster":  # pragma: no cover
-        # pylint: disable=no-member
         return SparkSession.builder.getOrCreate()
 
     os.environ["PYSPARK_PYTHON"] = sys.executable
@@ -46,33 +45,32 @@ def get_spark_session(
         path_to_replay_jar = os.environ.get("REPLAY_JAR_PATH")
     else:
         if pyspark_version.startswith("3.1"):  # pragma: no cover
-            path_to_replay_jar = "https://repo1.maven.org/maven2/io/github/sb-ai-lab/replay_2.12/3.1.3/replay_2.12-3.1.3.jar"
-        elif pyspark_version.startswith("3.2") or pyspark_version.startswith(
-            "3.3"
-        ):
+            path_to_replay_jar = (
+                "https://repo1.maven.org/maven2/io/github/sb-ai-lab/replay_2.12/3.1.3/replay_2.12-3.1.3.jar"
+            )
+        elif pyspark_version.startswith(("3.2", "3.3")):
             path_to_replay_jar = "https://repo1.maven.org/maven2/io/github/sb-ai-lab/replay_2.12/3.2.0_als_metrics/replay_2.12-3.2.0_als_metrics.jar"
         elif pyspark_version.startswith("3.4"):  # pragma: no cover
             path_to_replay_jar = "https://repo1.maven.org/maven2/io/github/sb-ai-lab/replay_2.12/3.4.0_als_metrics/replay_2.12-3.4.0_als_metrics.jar"
         else:  # pragma: no cover
-            path_to_replay_jar = "https://repo1.maven.org/maven2/io/github/sb-ai-lab/replay_2.12/3.1.3/replay_2.12-3.1.3.jar"
+            path_to_replay_jar = (
+                "https://repo1.maven.org/maven2/io/github/sb-ai-lab/replay_2.12/3.1.3/replay_2.12-3.1.3.jar"
+            )
             logging.warning(
-                "Replay ALS model support only spark 3.1-3.4 versions! "
-                "Replay will use 'https://repo1.maven.org/maven2/io/github/sb-ai-lab/replay_2.12/3.1.3/replay_2.12-3.1.3.jar' in 'spark.jars' property."
+                "Replay ALS model support only spark 3.1-3.4 versions! Replay will use "
+                "'https://repo1.maven.org/maven2/io/github/sb-ai-lab/replay_2.12/3.1.3/replay_2.12-3.1.3.jar' "
+                "in 'spark.jars' property."
             )
 
     if core_count is None:  # checking out env variable
         core_count = int(os.environ.get("REPLAY_SPARK_CORE_COUNT", "-1"))
     if spark_memory is None:
         env_var = os.environ.get("REPLAY_SPARK_MEMORY")
-        if env_var is not None:  # pragma: no cover
-            spark_memory = int(env_var)
-        else:  # pragma: no cover
-            spark_memory = floor(psutil.virtual_memory().total / 1024**3 * 0.7)
+        spark_memory = int(env_var) if env_var is not None else floor(psutil.virtual_memory().total / 1024**3 * 0.7)
     if shuffle_partitions is None:
         shuffle_partitions = os.cpu_count() * 3
     driver_memory = f"{spark_memory}g"
     user_home = os.environ["HOME"]
-    # pylint: disable=no-member
     spark = (
         SparkSession.builder.config("spark.driver.memory", driver_memory)
         .config(
@@ -111,7 +109,6 @@ def logger_with_settings() -> logging.Logger:
     return logger
 
 
-# pylint: disable=too-few-public-methods
 class Borg:
     """
     This class allows to share objects between instances.
@@ -123,7 +120,6 @@ class Borg:
         self.__dict__ = self._shared_state
 
 
-# pylint: disable=too-few-public-methods
 class State(Borg):
     """
     All modules look for Spark session via this class. You can put your own session here.

replay/utils/spark_utils.py

@@ -10,14 +10,17 @@ import pandas as pd
 from numpy.random import default_rng
 
 from .session_handler import State
-
 from .types import PYSPARK_AVAILABLE, DataFrameLike, MissingImportType, NumType, SparkDataFrame
 
 if PYSPARK_AVAILABLE:
     import pyspark.sql.types as st
     from pyspark.ml.linalg import DenseVector, Vectors, VectorUDT
-    from pyspark.sql import Column, SparkSession, Window
-    from pyspark.sql import functions as sf
+    from pyspark.sql import (
+        Column,
+        SparkSession,
+        Window,
+        functions as sf,
+    )
     from pyspark.sql.column import _to_java_column, _to_seq
     from pyspark.sql.types import DoubleType, IntegerType, StructField, StructType
 else:
@@ -48,7 +51,6 @@ def spark_to_pandas(data: SparkDataFrame, allow_collect_to_master: bool = False)
     return data.toPandas()
 
 
-# pylint: disable=invalid-name
 def convert2spark(data_frame: Optional[DataFrameLike]) -> Optional[SparkDataFrame]:
     """
     Converts Pandas DataFrame to Spark DataFrame
@@ -61,7 +63,7 @@ def convert2spark(data_frame: Optional[DataFrameLike]) -> Optional[SparkDataFram
     if isinstance(data_frame, SparkDataFrame):
         return data_frame
     spark = State().session
-    return spark.createDataFrame(data_frame)  # type: ignore
+    return spark.createDataFrame(data_frame)
 
 
 def get_top_k(
@@ -76,7 +78,11 @@ def get_top_k(
 
     >>> from replay.utils.session_handler import State
     >>> spark = State().session
-    >>> log = spark.createDataFrame([(1, 2, 1.), (1, 3, 1.), (1, 4, 0.5), (2, 1, 1.)]).toDF("user_id", "item_id", "relevance")
+    >>> log = (
+    ...    spark
+    ...    .createDataFrame([(1, 2, 1.), (1, 3, 1.), (1, 4, 0.5), (2, 1, 1.)])
+    ...    .toDF("user_id", "item_id", "relevance")
+    ... )
     >>> log.show()
     +-------+-------+---------+
     |user_id|item_id|relevance|
@@ -108,9 +114,7 @@ def get_top_k(
     return (
         dataframe.withColumn(
             "temp_rank",
-            sf.row_number().over(
-                Window.partitionBy(partition_by_col).orderBy(*order_by_col)
-            ),
+            sf.row_number().over(Window.partitionBy(partition_by_col).orderBy(*order_by_col)),
         )
         .filter(sf.col("temp_rank") <= k)
         .drop("temp_rank")
@@ -141,6 +145,7 @@ def get_top_k_recs(
 
 
 if PYSPARK_AVAILABLE:
+
     @sf.udf(returnType=st.DoubleType())
     def vector_dot(one: DenseVector, two: DenseVector) -> float:  # pragma: no cover
         """
@@ -179,10 +184,8 @@ if PYSPARK_AVAILABLE:
         """
         return float(one.dot(two))
 
-    @sf.udf(returnType=VectorUDT())  # type: ignore
-    def vector_mult(
-        one: Union[DenseVector, NumType], two: DenseVector
-    ) -> DenseVector:  # pragma: no cover
+    @sf.udf(returnType=VectorUDT())
+    def vector_mult(one: Union[DenseVector, NumType], two: DenseVector) -> DenseVector:  # pragma: no cover
         """
         elementwise vector multiplication
 
@@ -271,9 +274,7 @@ def multiply_scala_udf(scalar, vector):
     return Column(_f.apply(_to_seq(sc, [scalar, vector], _to_java_column)))
 
 
-def get_log_info(
-    log: SparkDataFrame, user_col="user_idx", item_col="item_idx"
-) -> str:
+def get_log_info(log: SparkDataFrame, user_col="user_idx", item_col="item_idx") -> str:
     """
     Basic log statistics
 
@@ -310,9 +311,7 @@ def get_log_info(
     )
 
 
-def get_stats(
-    log: SparkDataFrame, group_by: str = "user_id", target_column: str = "relevance"
-) -> SparkDataFrame:
+def get_stats(log: SparkDataFrame, group_by: str = "user_id", target_column: str = "relevance") -> SparkDataFrame:
     """
     Calculate log statistics: min, max, mean, median ratings, number of ratings.
     >>> from replay.utils.session_handler import get_spark_session, State
@@ -351,14 +350,9 @@ def get_stats(
         "count": sf.count,
     }
     agg_functions_list = [
-        func(target_column).alias(str(name + "_" + target_column))
-        for name, func in agg_functions.items()
+        func(target_column).alias(str(name + "_" + target_column)) for name, func in agg_functions.items()
     ]
-    agg_functions_list.append(
-        sf.expr(f"percentile_approx({target_column}, 0.5)").alias(
-            "median_" + target_column
-        )
-    )
+    agg_functions_list.append(sf.expr(f"percentile_approx({target_column}, 0.5)").alias("median_" + target_column))
 
     return log.groupBy(group_by).agg(*agg_functions_list)
 
@@ -369,13 +363,9 @@ def check_numeric(feature_table: SparkDataFrame) -> None:
     :param feature_table: spark DataFrame
     """
     for column in feature_table.columns:
-        if not isinstance(
-            feature_table.schema[column].dataType, st.NumericType
-        ):
-            raise ValueError(
-                f"""Column {column} has type {feature_table.schema[
-            column].dataType}, that is not numeric."""
-            )
+        if not isinstance(feature_table.schema[column].dataType, st.NumericType):
+            msg = f"Column {column} has type {feature_table.schema[column].dataType}, that is not numeric."
+            raise ValueError(msg)
 
 
 def horizontal_explode(
@@ -420,10 +410,7 @@ def horizontal_explode(
     num_columns = len(data_frame.select(column_to_explode).head()[0])
     return data_frame.select(
         *other_columns,
-        *[
-            sf.element_at(column_to_explode, i + 1).alias(f"{prefix}_{i}")
-            for i in range(num_columns)
-        ],
+        *[sf.element_at(column_to_explode, i + 1).alias(f"{prefix}_{i}") for i in range(num_columns)],
     )
 
 
@@ -442,7 +429,6 @@ def join_or_return(first, second, on, how):
     return first.join(second, on=on, how=how)
 
 
-# pylint: disable=too-many-arguments
 def fallback(
     base: SparkDataFrame,
     fill: SparkDataFrame,
@@ -471,15 +457,11 @@ def fallback(
     diff = max_in_fill - min_in_base
     fill = fill.withColumnRenamed(rating_column, "relevance_fallback")
     if diff >= 0:
-        fill = fill.withColumn(
-            "relevance_fallback", sf.col("relevance_fallback") - diff - margin
-        )
-    recs = base.join(
-        fill, on=[query_column, item_column], how="full_outer"
+        fill = fill.withColumn("relevance_fallback", sf.col("relevance_fallback") - diff - margin)
+    recs = base.join(fill, on=[query_column, item_column], how="full_outer")
+    recs = recs.withColumn(rating_column, sf.coalesce(rating_column, "relevance_fallback")).select(
+        query_column, item_column, rating_column
     )
-    recs = recs.withColumn(
-        rating_column, sf.coalesce(rating_column, "relevance_fallback")
-    ).select(query_column, item_column, rating_column)
     recs = get_top_k_recs(recs, k, query_column=query_column, rating_column=rating_column)
     return recs
 
@@ -537,9 +519,7 @@ def join_with_col_renaming(
             right = right.withColumnRenamed(name, f"{name}_{suffix}")
         on_condition &= sf.col(name) == sf.col(f"{name}_{suffix}")
 
-    return (left.join(right, on=on_condition, how=how)).drop(
-        *[f"{name}_{suffix}" for name in on_col_name]
-    )
+    return (left.join(right, on=on_condition, how=how)).drop(*[f"{name}_{suffix}" for name in on_col_name])
 
 
 def process_timestamp_column(
@@ -562,7 +542,8 @@ def process_timestamp_column(
     :return: dataframe with updated column ``column_name``
     """
     if column_name not in dataframe.columns:
-        raise ValueError(f"Column {column_name} not found")
+        msg = f"Column {column_name} not found"
+        raise ValueError(msg)
 
     # no conversion needed
     if isinstance(dataframe.schema[column_name].dataType, st.TimestampType):
@@ -570,9 +551,7 @@ def process_timestamp_column(
 
     # unix timestamp
     if isinstance(dataframe.schema[column_name].dataType, st.NumericType):
-        return dataframe.withColumn(
-            column_name, sf.to_timestamp(sf.from_unixtime(sf.col(column_name)))
-        )
+        return dataframe.withColumn(column_name, sf.to_timestamp(sf.from_unixtime(sf.col(column_name))))
 
     # datetime in string format
     dataframe = dataframe.withColumn(
@@ -583,6 +562,7 @@ def process_timestamp_column(
 
 
 if PYSPARK_AVAILABLE:
+
     @sf.udf(returnType=VectorUDT())
     def list_to_vector_udf(array: st.ArrayType) -> DenseVector:  # pragma: no cover
         """
@@ -603,9 +583,7 @@ if PYSPARK_AVAILABLE:
         return float(first.squared_distance(second))
 
     @sf.udf(returnType=st.FloatType())
-    def vector_euclidean_distance_similarity(
-        first: DenseVector, second: DenseVector
-    ) -> float:  # pragma: no cover
+    def vector_euclidean_distance_similarity(first: DenseVector, second: DenseVector) -> float:  # pragma: no cover
         """
         :param first: first vector
         :param second: second vector
@@ -642,7 +620,7 @@ def drop_temp_view(temp_view_name: str) -> None:
     spark.catalog.dropTempView(temp_view_name)
 
 
-def sample_top_k_recs(pairs: SparkDataFrame, k: int, seed: int = None):
+def sample_top_k_recs(pairs: SparkDataFrame, k: int, seed: Optional[int] = None):
     """
     Sample k items for each user with probability proportional to the relevance score.
 
@@ -660,17 +638,13 @@ def sample_top_k_recs(pairs: SparkDataFrame, k: int, seed: int = None):
     """
     pairs = pairs.withColumn(
         "probability",
-        sf.col("relevance")
-        / sf.sum("relevance").over(Window.partitionBy("user_idx")),
+        sf.col("relevance") / sf.sum("relevance").over(Window.partitionBy("user_idx")),
     )
 
     def grouped_map(pandas_df: pd.DataFrame) -> pd.DataFrame:  # pragma: no cover
         user_idx = pandas_df["user_idx"][0]
 
-        if seed is not None:
-            local_rng = default_rng(seed + user_idx)
-        else:
-            local_rng = default_rng()
+        local_rng = default_rng(seed + user_idx) if seed is not None else default_rng()
 
         items_positions = local_rng.choice(
             np.arange(pandas_df.shape[0]),
@@ -686,6 +660,7 @@ def sample_top_k_recs(pairs: SparkDataFrame, k: int, seed: int = None):
                 "relevance": pandas_df["relevance"].values[items_positions],
             }
         )
+
     rec_schema = StructType(
         [
             StructField("user_idx", IntegerType()),
@@ -716,19 +691,12 @@ def filter_cold(
     if df is None:
         return 0, df
 
-    num_cold = (
-        df.select(col_name)
-        .distinct()
-        .join(warm_df, on=col_name, how="anti")
-        .count()
-    )
+    num_cold = df.select(col_name).distinct().join(warm_df, on=col_name, how="anti").count()
 
     if num_cold == 0:
         return 0, df
 
-    return num_cold, df.join(
-        warm_df.select(col_name), on=col_name, how="inner"
-    )
+    return num_cold, df.join(warm_df.select(col_name), on=col_name, how="inner")
 
 
 def get_unique_entities(
@@ -745,17 +713,14 @@ def get_unique_entities(
     if isinstance(df, SparkDataFrame):
         unique = df.select(column).distinct()
     elif isinstance(df, collections.abc.Iterable):
-        unique = spark.createDataFrame(
-            data=pd.DataFrame(pd.unique(list(df)), columns=[column])
-        )
+        unique = spark.createDataFrame(data=pd.DataFrame(pd.unique(list(df)), columns=[column]))
     else:
-        raise ValueError(f"Wrong type {type(df)}")
+        msg = f"Wrong type {type(df)}"
+        raise ValueError(msg)
     return unique
 
 
-def return_recs(
-    recs: SparkDataFrame, recs_file_path: Optional[str] = None
-) -> Optional[SparkDataFrame]:
+def return_recs(recs: SparkDataFrame, recs_file_path: Optional[str] = None) -> Optional[SparkDataFrame]:
     """
     Save dataframe `recs` to `recs_file_path` if presents otherwise cache
     and materialize the dataframe.
@@ -785,7 +750,7 @@ def save_picklable_to_parquet(obj: Any, path: str) -> None:
     sc = State().session.sparkContext
     # We can use `RDD.saveAsPickleFile`, but it has no "overwrite" parameter
     pickled_instance = pickle.dumps(obj)
-    Record = collections.namedtuple("Record", ["data"])
+    Record = collections.namedtuple("Record", ["data"])  # noqa: PYI024
     rdd = sc.parallelize([Record(pickled_instance)])
     instance_df = rdd.map(lambda rec: Record(bytearray(rec.data))).toDF()
     instance_df.write.mode("overwrite").parquet(path)
@@ -812,9 +777,10 @@ def assert_omp_single_thread():
     PyTorch uses multithreading for cpu math operations via OpenMP library. Sometimes this
     leads to failures when OpenMP multithreading is mixed with multiprocessing.
     """
-    omp_num_threads = os.environ.get('OMP_NUM_THREADS', None)
-    if omp_num_threads != '1':
-        logging.getLogger("replay").warning(
-            'Environment variable "OMP_NUM_THREADS" is set to "%s". '
-            'Set it to 1 if the working process freezes.', omp_num_threads
+    omp_num_threads = os.environ.get("OMP_NUM_THREADS", None)
+    if omp_num_threads != "1":
+        msg = (
+            f'Environment variable "OMP_NUM_THREADS" is set to "{omp_num_threads}". '
+            f"Set it to 1 if the working process freezes."
         )
+        logging.getLogger("replay").warning(msg)