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

replay/data/nn/utils.py

@@ -2,11 +2,11 @@ from typing import Optional
 
 import polars as pl
 
-from replay.utils.spark_utils import spark_to_pandas
 from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame
+from replay.utils.spark_utils import spark_to_pandas
 
 if PYSPARK_AVAILABLE:  # pragma: no cover
-    import pyspark.sql.functions as F
+    import pyspark.sql.functions as sf
 
 
 def groupby_sequences(events: DataFrameLike, groupby_col: str, sort_col: Optional[str] = None) -> DataFrameLike:
@@ -38,9 +38,7 @@ def groupby_sequences(events: DataFrameLike, groupby_col: str, sort_col: Optiona
             event_cols_without_groupby.insert(0, sort_col)
             events = events.sort(event_cols_without_groupby)
 
-        grouped_sequences = events.group_by(groupby_col).agg(
-            *[pl.col(x) for x in event_cols_without_groupby]
-        )
+        grouped_sequences = events.group_by(groupby_col).agg(*[pl.col(x) for x in event_cols_without_groupby])
     else:
         event_cols_without_groupby = events.columns.copy()
         event_cols_without_groupby.remove(groupby_col)
@@ -49,16 +47,16 @@ def groupby_sequences(events: DataFrameLike, groupby_col: str, sort_col: Optiona
             event_cols_without_groupby.remove(sort_col)
             event_cols_without_groupby.insert(0, sort_col)
 
-        all_cols_struct = F.struct(event_cols_without_groupby)  # type: ignore
+        all_cols_struct = sf.struct(event_cols_without_groupby)
 
-        collect_fn = F.collect_list(all_cols_struct)
+        collect_fn = sf.collect_list(all_cols_struct)
         if sort_col:
-            collect_fn = F.sort_array(collect_fn)
+            collect_fn = sf.sort_array(collect_fn)
 
         grouped_sequences = (
             events.groupby(groupby_col)
             .agg(collect_fn.alias("_"))
-            .select([F.col(groupby_col)] + [F.col(f"_.{col}").alias(col) for col in event_cols_without_groupby])
+            .select([sf.col(groupby_col)] + [sf.col(f"_.{col}").alias(col) for col in event_cols_without_groupby])
             .drop("_")
         )
 

replay/data/schema.py

@@ -45,7 +45,6 @@ class FeatureInfo:
     Information about a feature.
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         column: str,
@@ -72,7 +71,8 @@ class FeatureInfo:
         self._feature_hint = feature_hint
 
         if feature_type == FeatureType.NUMERICAL and cardinality:
-            raise ValueError("Cardinality is needed only with categorical feature_type.")
+            msg = "Cardinality is needed only with categorical feature_type."
+            raise ValueError(msg)
         self._cardinality = cardinality
 
     @property
@@ -112,14 +112,12 @@ class FeatureInfo:
         :returns: cardinality of the feature.
         """
         if self.feature_type != FeatureType.CATEGORICAL:
-            raise RuntimeError(
-                f"Can not get cardinality because feature_type of {self.column} column is not categorical."
-            )
+            msg = f"Can not get cardinality because feature_type of {self.column} column is not categorical."
+            raise RuntimeError(msg)
         if hasattr(self, "_cardinality_callback") and self._cardinality is None:
             self._cardinality = self._cardinality_callback(self._column)
         return self._cardinality
 
-    # pylint: disable=attribute-defined-outside-init
     def _set_cardinality_callback(self, callback: Callable) -> None:
         self._cardinality_callback = callback
 
@@ -130,7 +128,6 @@ class FeatureInfo:
         self._cardinality = None
 
 
-# pylint: disable=too-many-public-methods
 class FeatureSchema(Mapping[str, FeatureInfo]):
     """
     Key-value like collection with information about all dataset features.
@@ -174,8 +171,9 @@ class FeatureSchema(Mapping[str, FeatureInfo]):
         :returns: extract a feature information from a schema.
         """
         if len(self._features_schema) > 1:
-            raise ValueError("Only one element feature schema can be converted to single feature")
-        return list(self._features_schema.values())[0]
+            msg = "Only one element feature schema can be converted to single feature"
+            raise ValueError(msg)
+        return next(iter(self._features_schema.values()))
 
     def items(self) -> ItemsView[str, FeatureInfo]:
         return self._features_schema.items()
@@ -186,7 +184,7 @@ class FeatureSchema(Mapping[str, FeatureInfo]):
     def values(self) -> ValuesView[FeatureInfo]:
         return self._features_schema.values()
 
-    def get(  # type: ignore
+    def get(
         self,
         key: str,
         default: Optional[FeatureInfo] = None,
@@ -358,7 +356,7 @@ class FeatureSchema(Mapping[str, FeatureInfo]):
         for filtration_func, filtration_param in zip(filter_functions, filter_parameters):
             filtered_features = list(
                 filter(
-                    lambda x: filtration_func(x, filtration_param),  # type: ignore  # pylint: disable=W0640
+                    lambda x: filtration_func(x, filtration_param),
                     filtered_features,
                 )
             )
@@ -391,7 +389,7 @@ class FeatureSchema(Mapping[str, FeatureInfo]):
         for filtration_func, filtration_param in zip(filter_functions, filter_parameters):
             filtered_features = list(
                 filter(
-                    lambda x: filtration_func(x, filtration_param),  # type: ignore  # pylint: disable=W0640
+                    lambda x: filtration_func(x, filtration_param),
                     filtered_features,
                 )
             )
@@ -426,7 +424,6 @@ class FeatureSchema(Mapping[str, FeatureInfo]):
     def _type_drop(value: FeatureInfo, feature_type: FeatureType) -> bool:
         return value.feature_type != feature_type if feature_type else True
 
-    # pylint: disable=no-self-use
     @staticmethod
     def _hint_drop(value: FeatureInfo, feature_hint: FeatureHint) -> bool:
         return value.feature_hint != feature_hint if feature_hint else True
@@ -451,13 +448,16 @@ class FeatureSchema(Mapping[str, FeatureInfo]):
                 item_query_names[feature.feature_hint] += [feature.column]
 
         if len(duplicates) > 0:
-            raise ValueError(
+            msg = (
                 "Features column names should be unique, exept ITEM_ID and QUERY_ID columns. "
-                + f"{duplicates} columns are not unique."
+                f"{duplicates} columns are not unique."
             )
+            raise ValueError(msg)
 
         if len(item_query_names[FeatureHint.ITEM_ID]) > 1:
-            raise ValueError(f"ITEM_ID must be present only once. Rename {item_query_names[FeatureHint.ITEM_ID]}")
+            msg = f"ITEM_ID must be present only once. Rename {item_query_names[FeatureHint.ITEM_ID]}"
+            raise ValueError(msg)
 
         if len(item_query_names[FeatureHint.QUERY_ID]) > 1:
-            raise ValueError(f"QUERY_ID must be present only once. Rename {item_query_names[FeatureHint.QUERY_ID]}")
+            msg = f"QUERY_ID must be present only once. Rename {item_query_names[FeatureHint.QUERY_ID]}"
+            raise ValueError(msg)

replay/data/spark_schema.py

@@ -4,7 +4,6 @@ if PYSPARK_AVAILABLE:
     from pyspark.sql.types import DoubleType, IntegerType, StructField, StructType, TimestampType
 
 
-# pylint: disable=too-many-arguments
 def get_schema(
     query_column: str = "query_id",
     item_column: str = "item_id",

replay/metrics/base_metric.py

@@ -1,11 +1,11 @@
 import warnings
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Mapping, Union
+from typing import Any, Dict, List, Mapping, Optional, Union
 
 import numpy as np
 import polars as pl
 
-from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, SparkDataFrame, PolarsDataFrame
+from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
 from .descriptors import CalculationDescriptor, Mean
 
@@ -27,7 +27,7 @@ class MetricDuplicatesWarning(Warning):
 class Metric(ABC):
     """Base metric class"""
 
-    def __init__(  # pylint: disable=too-many-arguments
+    def __init__(
         self,
         topk: Union[List[int], int],
         query_column: str = "query_id",
@@ -46,11 +46,13 @@ class Metric(ABC):
         if isinstance(topk, list):
             for item in topk:
                 if not isinstance(item, int):
-                    raise ValueError(f"{item} is not int")
+                    msg = f"{item} is not int"
+                    raise ValueError(msg)
         elif isinstance(topk, int):
             topk = [topk]
         else:
-            raise ValueError("topk not list or int")
+            msg = "topk not list or int"
+            raise ValueError(msg)
         self.topk = sorted(topk)
         self.query_column = query_column
         self.item_column = item_column
@@ -60,11 +62,8 @@ class Metric(ABC):
     @property
     def __name__(self) -> str:
         mode_name = self._mode.__name__
-        return str(type(self).__name__) + (
-            f"-{mode_name}" if mode_name != "Mean" else ""
-        )
+        return str(type(self).__name__) + (f"-{mode_name}" if mode_name != "Mean" else "")
 
-    # pylint: disable=no-self-use
     def _check_dataframes_equal_types(
         self,
         recommendations: MetricsDataFrameLike,
@@ -74,39 +73,31 @@ class Metric(ABC):
         Types of all data frames must be the same.
         """
         if not isinstance(recommendations, type(ground_truth)):
-            raise ValueError("All given data frames must have the same type")
+            msg = "All given data frames must have the same type"
+            raise ValueError(msg)
 
     def _duplicate_warn(self):
         warnings.warn(
-            "The recommendations contain duplicated users and items."
-            "The metrics may be higher than the actual ones.",
+            "The recommendations contain duplicated users and items.The metrics may be higher than the actual ones.",
             MetricDuplicatesWarning,
         )
 
     def _check_duplicates_spark(self, recommendations: SparkDataFrame) -> None:
         duplicates_count = (
-            recommendations.groupBy(self.query_column, self.item_column)
-            .count()
-            .filter("count >= 2")
-            .count()
+            recommendations.groupBy(self.query_column, self.item_column).count().filter("count >= 2").count()
         )
         if duplicates_count:
             self._duplicate_warn()
 
     def _check_duplicates_dict(self, recommendations: Dict) -> None:
-        for _, items in recommendations.items():
+        for items in recommendations.values():
             items_set = set(items)
             if len(items) != len(items_set):
                 self._duplicate_warn()
                 return
 
     def _check_duplicates_polars(self, recommendations: PolarsDataFrame) -> None:
-        duplicates_count = (
-            recommendations
-            .group_by(self.query_column, self.item_column)
-            .len()
-            .filter(pl.col("len") > 1)
-        )
+        duplicates_count = recommendations.group_by(self.query_column, self.item_column).len().filter(pl.col("len") > 1)
         if not duplicates_count.is_empty():
             self._duplicate_warn()
 
@@ -144,11 +135,7 @@ class Metric(ABC):
             else self._convert_dict_to_dict_with_score(recommendations)
         )
         self._check_duplicates_dict(recommendations)
-        ground_truth = (
-            self._convert_pandas_to_dict_without_score(ground_truth)
-            if is_pandas
-            else ground_truth
-        )
+        ground_truth = self._convert_pandas_to_dict_without_score(ground_truth) if is_pandas else ground_truth
         assert isinstance(ground_truth, dict)
         return self._dict_call(
             list(ground_truth),
@@ -164,7 +151,6 @@ class Metric(ABC):
             .to_dict()
         )
 
-    # pylint: disable=no-self-use
     def _convert_dict_to_dict_with_score(self, data: Dict) -> Dict:
         converted_data = {}
         for user, items in data.items():
@@ -191,31 +177,21 @@ class Metric(ABC):
         distribution_per_user = {}
         for user in users:
             args = [kwargs[key].get(user, None) for key in keys_list]
-            distribution_per_user[user] = self._get_metric_value_by_user(
-                self.topk, *args
-            )  # pylint: disable=protected-access
+            distribution_per_user[user] = self._get_metric_value_by_user(self.topk, *args)
         if self._mode.__name__ == "PerUser":
             return self._aggregate_results_per_user(distribution_per_user)
         distribution = np.stack(list(distribution_per_user.values()))
         assert distribution.shape[1] == len(self.topk)
-        metrics = []
-        for k in range(distribution.shape[1]):
-            metrics.append(self._mode.cpu(distribution[:, k]))
+        metrics = [self._mode.cpu(distribution[:, k]) for k in range(distribution.shape[1])]
         return self._aggregate_results(metrics)
 
     def _get_items_list_per_user_spark(
-        self, recommendations: SparkDataFrame, extra_column: str = None
+        self, recommendations: SparkDataFrame, extra_column: Optional[str] = None
     ) -> SparkDataFrame:
         recommendations = recommendations.groupby(self.query_column).agg(
             sf.sort_array(
                 sf.collect_list(
-                    sf.struct(
-                        *[
-                            c
-                            for c in [self.rating_column, self.item_column, extra_column]
-                            if c is not None
-                        ]
-                    )
+                    sf.struct(*[c for c in [self.rating_column, self.item_column, extra_column] if c is not None])
                 ),
                 False,
             ).alias("pred")
@@ -231,7 +207,7 @@ class Metric(ABC):
         return recommendations
 
     def _get_items_list_per_user_polars(
-        self, recommendations: PolarsDataFrame, extra_column: str = None
+        self, recommendations: PolarsDataFrame, extra_column: Optional[str] = None
     ) -> PolarsDataFrame:
         selection = [self.query_column, "pred_item_id"]
         sorting = [self.rating_column, self.item_column]
@@ -242,8 +218,7 @@ class Metric(ABC):
             selection.append(extra_column)
 
         recommendations = (
-            recommendations
-            .sort(sorting, descending=True)
+            recommendations.sort(sorting, descending=True)
             .group_by(self.query_column)
             .agg(*agg)
             .rename({self.item_column: "pred_item_id"})
@@ -253,7 +228,7 @@ class Metric(ABC):
         return recommendations
 
     def _get_items_list_per_user(
-        self, recommendations: Union[SparkDataFrame, PolarsDataFrame], extra_column: str = None
+        self, recommendations: Union[SparkDataFrame, PolarsDataFrame], extra_column: Optional[str] = None
     ) -> Union[SparkDataFrame, PolarsDataFrame]:
         if isinstance(recommendations, SparkDataFrame):
             return self._get_items_list_per_user_spark(recommendations, extra_column)
@@ -265,7 +240,7 @@ class Metric(ABC):
     ) -> Union[SparkDataFrame, PolarsDataFrame]:
         cols = data.columns
         cols.remove(self.query_column)
-        cols = [self.query_column] + sorted(cols)
+        cols = [self.query_column, *sorted(cols)]
         return data.select(*cols)
 
     def _get_enriched_recommendations(
@@ -300,8 +275,7 @@ class Metric(ABC):
         ground_truth: PolarsDataFrame,
     ) -> PolarsDataFrame:
         true_items_by_users = (
-            ground_truth
-            .group_by(self.query_column)
+            ground_truth.group_by(self.query_column)
             .agg(pl.col(self.item_column))
             .rename({self.item_column: "ground_truth"})
         )
@@ -313,9 +287,7 @@ class Metric(ABC):
         )
         return self._rearrange_columns(enriched_recommendations)
 
-    def _aggregate_results_per_user(
-        self, distribution_per_user: Dict[Any, List[float]]
-    ) -> MetricsPerUserReturnType:
+    def _aggregate_results_per_user(self, distribution_per_user: Dict[Any, List[float]]) -> MetricsPerUserReturnType:
         res: MetricsPerUserReturnType = {}
         for index, val in enumerate(self.topk):
             metric_name = f"{self.__name__}@{val}"
@@ -335,18 +307,12 @@ class Metric(ABC):
         """
         Calculating metrics for PySpark DataFrame.
         """
-        recs_with_topk_list = recs.withColumn(
-            "k", sf.array(*[sf.lit(x) for x in self.topk])
-        )
+        recs_with_topk_list = recs.withColumn("k", sf.array(*[sf.lit(x) for x in self.topk]))
         distribution = self._get_metric_distribution(recs_with_topk_list)
         if self._mode.__name__ == "PerUser":
             return self._aggregate_results_per_user(distribution.rdd.collectAsMap())
         metrics = [
-            self._mode.spark(
-                distribution.select(sf.col("value").getItem(i)).withColumnRenamed(
-                    f"value[{i}]", "val"
-                )
-            )
+            self._mode.spark(distribution.select(sf.col("value").getItem(i)).withColumnRenamed(f"value[{i}]", "val"))
             for i in range(len(self.topk))
         ]
         return self._aggregate_results(metrics)
@@ -355,27 +321,23 @@ class Metric(ABC):
         distribution = self._get_metric_distribution(recs)
         if self._mode.__name__ == "PerUser":
             return self._aggregate_results_per_user(
-                dict(distribution.select(
-                    self.query_column,
-                    value=pl.concat_list(pl.exclude(self.query_column))
-                ).iter_rows())
+                dict(
+                    distribution.select(
+                        self.query_column, value=pl.concat_list(pl.exclude(self.query_column))
+                    ).iter_rows()
+                )
             )
-        metrics = [self._mode.cpu(distribution.select(column))
-                   for column in distribution.columns[1:]]
+        metrics = [self._mode.cpu(distribution.select(column)) for column in distribution.columns[1:]]
         return self._aggregate_results(metrics)
 
-    def _spark_call(
-        self, recommendations: SparkDataFrame, ground_truth: SparkDataFrame
-    ) -> MetricsReturnType:
+    def _spark_call(self, recommendations: SparkDataFrame, ground_truth: SparkDataFrame) -> MetricsReturnType:
         """
         Implementation for PySpark DataFrame.
         """
         recs = self._get_enriched_recommendations(recommendations, ground_truth)
         return self._spark_compute(recs)
 
-    def _polars_call(
-        self, recommendations: PolarsDataFrame, ground_truth: PolarsDataFrame
-    ) -> MetricsReturnType:
+    def _polars_call(self, recommendations: PolarsDataFrame, ground_truth: PolarsDataFrame) -> MetricsReturnType:
         """
         Implementation for Polars DataFrame.
         """
@@ -383,7 +345,7 @@ class Metric(ABC):
         return self._polars_compute(recs)
 
     def _get_metric_distribution(
-            self, recs: Union[PolarsDataFrame, SparkDataFrame]
+        self, recs: Union[PolarsDataFrame, SparkDataFrame]
     ) -> Union[PolarsDataFrame, SparkDataFrame]:
         if isinstance(recs, SparkDataFrame):
             return self._get_metric_distribution_spark(recs)
@@ -406,16 +368,13 @@ class Metric(ABC):
         distribution = recs.map_rows(lambda x: (x[0], *cur_class._get_metric_value_by_user(self.topk, *x[1:])))
         distribution = distribution.rename({"column_0": self.query_column})
         distribution = distribution.rename(
-            {distribution.columns[x + 1]: f"value_{self.topk[x]}"
-             for x in range(len(self.topk))}
+            {distribution.columns[x + 1]: f"value_{self.topk[x]}" for x in range(len(self.topk))}
         )
         return distribution
 
     @staticmethod
     @abstractmethod
-    def _get_metric_value_by_user(  # pylint: disable=invalid-name
-        ks: List[int], *args: List
-    ) -> List[float]:  # pragma: no cover
+    def _get_metric_value_by_user(ks: List[int], *args: List) -> List[float]:  # pragma: no cover
         """
         Metric calculation for one user.
 

replay/metrics/categorical_diversity.py

@@ -4,7 +4,7 @@ from typing import Dict, List, Union
 import numpy as np
 import polars as pl
 
-from replay.utils import PYSPARK_AVAILABLE, PandasDataFrame, SparkDataFrame, PolarsDataFrame
+from replay.utils import PYSPARK_AVAILABLE, PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
 from .base_metric import (
     Metric,
@@ -16,11 +16,12 @@ from .base_metric import (
 from .descriptors import CalculationDescriptor, Mean
 
 if PYSPARK_AVAILABLE:
-    from pyspark.sql import Window
-    from pyspark.sql import functions as F
+    from pyspark.sql import (
+        Window,
+        functions as sf,
+    )
 
 
-# pylint: disable=too-few-public-methods
 class CategoricalDiversity(Metric):
     """
     Metric calculation is as follows:
@@ -59,7 +60,6 @@ class CategoricalDiversity(Metric):
     <BLANKLINE>
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         topk: Union[List, int],
@@ -108,31 +108,21 @@ class CategoricalDiversity(Metric):
         precalculated_answer = self._precalculate_unique_cats(recommendations)
         return self._dict_call(precalculated_answer)
 
-    # pylint: disable=arguments-differ
     def _get_enriched_recommendations(
-        self, recommendations: Union[PolarsDataFrame, SparkDataFrame],
+        self,
+        recommendations: Union[PolarsDataFrame, SparkDataFrame],
     ) -> Union[PolarsDataFrame, SparkDataFrame]:
         if isinstance(recommendations, SparkDataFrame):
             return self._get_enriched_recommendations_spark(recommendations)
         else:
             return self._get_enriched_recommendations_polars(recommendations)
 
-    # pylint: disable=arguments-differ
-    def _get_enriched_recommendations_spark(
-        self, recommendations: SparkDataFrame
-    ) -> SparkDataFrame:
-        window = Window.partitionBy(self.query_column).orderBy(
-            F.col(self.rating_column).desc()
-        )
-        sorted_by_score_recommendations = recommendations.withColumn(
-            "rank", F.row_number().over(window)
-        )
+    def _get_enriched_recommendations_spark(self, recommendations: SparkDataFrame) -> SparkDataFrame:
+        window = Window.partitionBy(self.query_column).orderBy(sf.col(self.rating_column).desc())
+        sorted_by_score_recommendations = recommendations.withColumn("rank", sf.row_number().over(window))
         return sorted_by_score_recommendations
 
-    # pylint: disable=arguments-differ
-    def _get_enriched_recommendations_polars(
-        self, recommendations: PolarsDataFrame
-    ) -> PolarsDataFrame:
+    def _get_enriched_recommendations_polars(self, recommendations: PolarsDataFrame) -> PolarsDataFrame:
         sorted_by_score_recommendations = recommendations.select(
             pl.all().sort_by(self.rating_column, descending=True).over(self.query_column)
         )
@@ -146,13 +136,9 @@ class CategoricalDiversity(Metric):
     def _spark_compute_per_user(self, recs: SparkDataFrame) -> MetricsPerUserReturnType:
         distribution_per_user = defaultdict(list)
         for k in self.topk:
-            filtered_recs = recs.filter(F.col("rank") <= k)
-            aggreagated_by_user = filtered_recs.groupBy(self.query_column).agg(
-                F.countDistinct(self.category_column)
-            )
-            aggreagated_by_user_dict = (
-                aggreagated_by_user.rdd.collectAsMap()
-            )  # type:ignore
+            filtered_recs = recs.filter(sf.col("rank") <= k)
+            aggreagated_by_user = filtered_recs.groupBy(self.query_column).agg(sf.countDistinct(self.category_column))
+            aggreagated_by_user_dict = aggreagated_by_user.rdd.collectAsMap()
             for user, metric in aggreagated_by_user_dict.items():
                 distribution_per_user[user].append(metric / k)
         return self._aggregate_results_per_user(dict(distribution_per_user))
@@ -161,12 +147,8 @@ class CategoricalDiversity(Metric):
         distribution_per_user = defaultdict(list)
         for k in self.topk:
             filtered_recs = recs.filter(pl.col("rank") <= k)
-            aggreagated_by_user = filtered_recs.group_by(self.query_column).agg(
-                pl.col(self.category_column).n_unique()
-            )
-            aggreagated_by_user_dict = (
-                dict(aggreagated_by_user.iter_rows())
-            )  # type:ignore
+            aggreagated_by_user = filtered_recs.group_by(self.query_column).agg(pl.col(self.category_column).n_unique())
+            aggreagated_by_user_dict = dict(aggreagated_by_user.iter_rows())
             for user, metric in aggreagated_by_user_dict.items():
                 distribution_per_user[user].append(metric / k)
         return self._aggregate_results_per_user(dict(distribution_per_user))
@@ -174,10 +156,10 @@ class CategoricalDiversity(Metric):
     def _spark_compute_agg(self, recs: SparkDataFrame) -> MetricsMeanReturnType:
         metrics = []
         for k in self.topk:
-            filtered_recs = recs.filter(F.col("rank") <= k)
+            filtered_recs = recs.filter(sf.col("rank") <= k)
             aggregated_by_user = (
                 filtered_recs.groupBy(self.query_column)
-                .agg(F.countDistinct(self.category_column))
+                .agg(sf.countDistinct(self.category_column))
                 .drop(self.query_column)
             )
             metrics.append(self._mode.spark(aggregated_by_user) / k)
@@ -195,7 +177,6 @@ class CategoricalDiversity(Metric):
             metrics.append(self._mode.cpu(aggregated_by_user) / k)
         return self._aggregate_results(metrics)
 
-    # pylint: disable=arguments-differ
     def _spark_call(self, recommendations: SparkDataFrame) -> MetricsReturnType:
         """
         Implementation for Pyspark DataFrame.
@@ -205,7 +186,6 @@ class CategoricalDiversity(Metric):
             return self._spark_compute_per_user(recs)
         return self._spark_compute_agg(recs)
 
-    # pylint: disable=arguments-differ
     def _polars_call(self, recommendations: PolarsDataFrame) -> MetricsReturnType:
         """
         Implementation for Polars DataFrame.
@@ -223,7 +203,6 @@ class CategoricalDiversity(Metric):
             .to_dict()
         )
 
-    # pylint: disable=no-self-use
     def _precalculate_unique_cats(self, recommendations: Dict) -> Dict:
         """
         Precalculate unique categories for each prefix for each user.
@@ -238,24 +217,16 @@ class CategoricalDiversity(Metric):
             answer[user] = unique_len
         return answer
 
-    # pylint: disable=arguments-renamed,arguments-differ
-    def _dict_compute_per_user(
-        self, precalculated_answer: Dict
-    ) -> MetricsPerUserReturnType:  # type:ignore
+    def _dict_compute_per_user(self, precalculated_answer: Dict) -> MetricsPerUserReturnType:
         distribution_per_user = defaultdict(list)
         for k in self.topk:
             for user, unique_cats in precalculated_answer.items():
-                distribution_per_user[user].append(
-                    unique_cats[min(len(unique_cats), k) - 1] / k
-                )
+                distribution_per_user[user].append(unique_cats[min(len(unique_cats), k) - 1] / k)
         return self._aggregate_results_per_user(distribution_per_user)
 
-    # pylint: disable=arguments-renamed
-    def _dict_compute_mean(
-        self, precalculated_answer: Dict
-    ) -> MetricsMeanReturnType:  # type:ignore
+    def _dict_compute_mean(self, precalculated_answer: Dict) -> MetricsMeanReturnType:
         distribution_list = []
-        for _, unique_cats in precalculated_answer.items():
+        for unique_cats in precalculated_answer.values():
             metrics_per_user = []
             for k in self.topk:
                 metric = unique_cats[min(len(unique_cats), k) - 1] / k
@@ -264,12 +235,9 @@ class CategoricalDiversity(Metric):
 
         distribution = np.stack(distribution_list)
         assert distribution.shape[1] == len(self.topk)
-        metrics = []
-        for k in range(distribution.shape[1]):
-            metrics.append(self._mode.cpu(distribution[:, k]))
+        metrics = [self._mode.cpu(distribution[:, k]) for k in range(distribution.shape[1])]
         return self._aggregate_results(metrics)
 
-    # pylint: disable=arguments-differ
     def _dict_call(self, precalculated_answer: Dict) -> MetricsReturnType:
         """
         Calculating metrics in dict format.
@@ -279,7 +247,5 @@ class CategoricalDiversity(Metric):
         return self._dict_compute_mean(precalculated_answer)
 
     @staticmethod
-    def _get_metric_value_by_user(
-        ks: List[int], *args: List
-    ) -> List[float]:  # pragma: no cover
+    def _get_metric_value_by_user(ks: List[int], *args: List) -> List[float]:  # pragma: no cover
         pass