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

replay/preprocessing/label_encoder.py

@@ -6,9 +6,11 @@ Contains classes for encoding categorical data
 ``LabelEncoder`` to apply multiple LabelEncodingRule to dataframe.
 """
 import abc
-import polars as pl
+import warnings
 from typing import Dict, List, Literal, Mapping, Optional, Sequence, Union
 
+import polars as pl
+
 from replay.utils import (
     PYSPARK_AVAILABLE,
     DataFrameLike,
@@ -20,13 +22,16 @@ from replay.utils import (
 
 if PYSPARK_AVAILABLE:
     from pyspark.sql import functions as sf
+    from pyspark.sql.types import LongType, StructType
     from pyspark.storagelevel import StorageLevel
-    from pyspark.sql.types import StructType, LongType
 
-HandleUnknownStrategies = Literal["error", "use_default_value"]
+HandleUnknownStrategies = Literal["error", "use_default_value", "drop"]
+
+
+class LabelEncoderTransformWarning(Warning):
+    """Label encoder transform warning."""
 
 
-# pylint: disable=missing-function-docstring
 class BaseLabelEncodingRule(abc.ABC):  # pragma: no cover
     """
     Interface of the label encoding rule
@@ -70,7 +75,6 @@ class BaseLabelEncodingRule(abc.ABC):  # pragma: no cover
         raise NotImplementedError()
 
 
-# pylint: disable=too-many-instance-attributes
 class LabelEncodingRule(BaseLabelEncodingRule):
     """
     Implementation of the encoding rule for categorical variables of PySpark and Pandas Data Frames.
@@ -79,7 +83,7 @@ class LabelEncodingRule(BaseLabelEncodingRule):
     """
 
     _ENCODED_COLUMN_SUFFIX: str = "_encoded"
-    _HANDLE_UNKNOWN_STRATEGIES = ("error", "use_default_value")
+    _HANDLE_UNKNOWN_STRATEGIES = ("error", "use_default_value", "drop")
     _TRANSFORM_PERFORMANCE_THRESHOLD_FOR_PANDAS = 100_000
 
     def __init__(
@@ -99,6 +103,7 @@ class LabelEncodingRule(BaseLabelEncodingRule):
             When set to ``error`` an error will be raised in case an unknown label is present during transform.
             When set to ``use_default_value``, the encoded value of unknown label will be set
             to the value given for the parameter default_value.
+            When set to ``drop``, the unknown labels will be dropped.
             Default: ``error``.
         :param default_value: Default value that will fill the unknown labels after transform.
             When the parameter handle_unknown is set to ``use_default_value``,
@@ -110,7 +115,8 @@ class LabelEncodingRule(BaseLabelEncodingRule):
             Default: ``None``.
         """
         if handle_unknown not in self._HANDLE_UNKNOWN_STRATEGIES:
-            raise ValueError(f"handle_unknown should be either 'error' or 'use_default_value', got {handle_unknown}.")
+            msg = f"handle_unknown should be either 'error' or 'use_default_value', got {handle_unknown}."
+            raise ValueError(msg)
         self._handle_unknown = handle_unknown
         if (
             self._handle_unknown == "use_default_value"
@@ -118,7 +124,8 @@ class LabelEncodingRule(BaseLabelEncodingRule):
             and not isinstance(default_value, int)
             and default_value != "last"
         ):
-            raise ValueError("Default value should be None, int or 'last'")
+            msg = "Default value should be None, int or 'last'"
+            raise ValueError(msg)
 
         self._default_value = default_value
         self._col = column
@@ -135,12 +142,14 @@ class LabelEncodingRule(BaseLabelEncodingRule):
 
     def get_mapping(self) -> Mapping:
         if self._mapping is None:
-            raise RuntimeError("Label encoder is not fitted")
+            msg = "Label encoder is not fitted"
+            raise RuntimeError(msg)
         return self._mapping
 
     def get_inverse_mapping(self) -> Mapping:
         if self._mapping is None:
-            raise RuntimeError("Label encoder is not fitted")
+            msg = "Label encoder is not fitted"
+            raise RuntimeError(msg)
         return self._inverse_mapping
 
     def _make_inverse_mapping(self) -> Mapping:
@@ -159,17 +168,14 @@ class LabelEncodingRule(BaseLabelEncodingRule):
             unique_col_values.rdd.zipWithIndex()
             .toDF(
                 StructType()
-                .add("_1",
-                     StructType()
-                     .add(self._col, df.schema[self._col].dataType, True),
-                     True)
+                .add("_1", StructType().add(self._col, df.schema[self._col].dataType, True), True)
                 .add("_2", LongType(), True)
             )
             .select(sf.col(f"_1.{self._col}").alias(self._col), sf.col("_2").alias(self._target_col))
             .persist(StorageLevel.MEMORY_ONLY)
         )
 
-        self._mapping = mapping_on_spark.rdd.collectAsMap()  # type: ignore
+        self._mapping = mapping_on_spark.rdd.collectAsMap()
         mapping_on_spark.unpersist()
         unique_col_values.unpersist()
 
@@ -198,17 +204,18 @@ class LabelEncodingRule(BaseLabelEncodingRule):
         elif isinstance(df, PolarsDataFrame):
             self._fit_polars(df)
         else:
-            raise NotImplementedError(f"{self.__class__.__name__} is not implemented for {type(df)}")
+            msg = f"{self.__class__.__name__} is not implemented for {type(df)}"
+            raise NotImplementedError(msg)
         self._inverse_mapping = self._make_inverse_mapping()
         self._inverse_mapping_list = self._make_inverse_mapping_list()
-        if self._handle_unknown == "use_default_value":
-            if self._default_value in self._inverse_mapping:
-                raise ValueError(
-                    "The used value for default_value "
-                    f"{self._default_value} is one of the "
-                    "values already used for encoding the "
-                    "seen labels."
-                )
+        if self._handle_unknown == "use_default_value" and self._default_value in self._inverse_mapping:
+            msg = (
+                "The used value for default_value "
+                f"{self._default_value} is one of the "
+                "values already used for encoding the "
+                "seen labels."
+            )
+            raise ValueError(msg)
         self._is_fitted = True
         return self
 
@@ -226,18 +233,15 @@ class LabelEncodingRule(BaseLabelEncodingRule):
             new_unique_values.rdd.zipWithIndex()
             .toDF(
                 StructType()
-                .add("_1",
-                     StructType()
-                     .add(self._col, df.schema[self._col].dataType),
-                     True)
+                .add("_1", StructType().add(self._col, df.schema[self._col].dataType), True)
                 .add("_2", LongType(), True)
             )
             .select(sf.col(f"_1.{self._col}").alias(self._col), sf.col("_2").alias(self._target_col))
             .withColumn(self._target_col, sf.col(self._target_col) + max_value)
-            .rdd.collectAsMap()  # type: ignore
+            .rdd.collectAsMap()
         )
-        self._mapping.update(new_data)  # type: ignore
-        self._inverse_mapping.update({v: k for k, v in new_data.items()})  # type: ignore
+        self._mapping.update(new_data)
+        self._inverse_mapping.update({v: k for k, v in new_data.items()})
         self._inverse_mapping_list.extend(new_data.keys())
         new_unique_values.unpersist()
 
@@ -245,9 +249,10 @@ class LabelEncodingRule(BaseLabelEncodingRule):
         assert self._mapping is not None
 
         new_unique_values = set(df[self._col].tolist()) - set(self._mapping)
-        new_data: dict = {value: max(self._mapping.values()) + i for i, value in enumerate(new_unique_values, start=1)}
-        self._mapping.update(new_data)  # type: ignore
-        self._inverse_mapping.update({v: k for k, v in new_data.items()})  # type: ignore
+        last_mapping_value = max(self._mapping.values())
+        new_data: dict = {value: last_mapping_value + i for i, value in enumerate(new_unique_values, start=1)}
+        self._mapping.update(new_data)
+        self._inverse_mapping.update({v: k for k, v in new_data.items()})
         self._inverse_mapping_list.extend(new_data.keys())
 
     def _partial_fit_polars(self, df: PolarsDataFrame) -> None:
@@ -255,8 +260,8 @@ class LabelEncodingRule(BaseLabelEncodingRule):
 
         new_unique_values = set(df.select(self._col).unique().to_series().to_list()) - set(self._mapping)
         new_data: dict = {value: max(self._mapping.values()) + i for i, value in enumerate(new_unique_values, start=1)}
-        self._mapping.update(new_data)  # type: ignore
-        self._inverse_mapping.update({v: k for k, v in new_data.items()})  # type: ignore
+        self._mapping.update(new_data)
+        self._inverse_mapping.update({v: k for k, v in new_data.items()})
         self._inverse_mapping_list.extend(new_data.keys())
 
     def partial_fit(self, df: DataFrameLike) -> "LabelEncodingRule":
@@ -276,7 +281,8 @@ class LabelEncodingRule(BaseLabelEncodingRule):
         elif isinstance(df, PolarsDataFrame):
             self._partial_fit_polars(df)
         else:
-            raise NotImplementedError(f"{self.__class__.__name__} is not implemented for {type(df)}")
+            msg = f"{self.__class__.__name__} is not implemented for {type(df)}"
+            raise NotImplementedError(msg)
 
         self._is_fitted = True
         return self
@@ -299,14 +305,24 @@ class LabelEncodingRule(BaseLabelEncodingRule):
             joined_df.loc[unknown_mask, self._target_col] = -1
             is_unknown_label |= unknown_mask.sum() > 0
 
-        if is_unknown_label and default_value != -1:
+        if is_unknown_label:
             unknown_mask = joined_df[self._target_col] == -1
-            if self._handle_unknown == "error":
+            if self._handle_unknown == "drop":
+                joined_df.drop(joined_df[unknown_mask].index, inplace=True)
+                if joined_df.empty:
+                    warnings.warn(
+                        f"You are trying to transform dataframe with all values are unknown for {self._col}, "
+                        "with `handle_unknown_strategy=drop` leads to empty dataframe",
+                        LabelEncoderTransformWarning,
+                    )
+            elif self._handle_unknown == "error":
                 unknown_unique_labels = joined_df[self._col][unknown_mask].unique().tolist()
                 msg = f"Found unknown labels {unknown_unique_labels} in column {self._col} during transform"
                 raise ValueError(msg)
-            joined_df[self._target_col] = joined_df[self._target_col].astype("int")
-            joined_df[self._target_col] = joined_df[self._target_col].replace({-1: default_value})
+            else:
+                if default_value != -1:
+                    joined_df[self._target_col] = joined_df[self._target_col].astype("int")
+                    joined_df[self._target_col] = joined_df[self._target_col].replace({-1: default_value})
 
         result_df = joined_df.drop(self._col, axis=1).rename(columns={self._target_col: self._col})
         return result_df
@@ -318,17 +334,24 @@ class LabelEncodingRule(BaseLabelEncodingRule):
         transformed_df = df.join(mapping_on_spark, on=self._col, how="left").withColumn(
             "unknown_mask", sf.isnull(self._target_col)
         )
-        unknown_label_count = transformed_df.select(sf.sum(sf.col("unknown_mask").cast("long"))).first()[
-            0
-        ]  # type: ignore
+        unknown_label_count = transformed_df.select(sf.sum(sf.col("unknown_mask").cast("long"))).first()[0]
         if unknown_label_count > 0:
-            if self._handle_unknown == "error":
+            if self._handle_unknown == "drop":
+                transformed_df = transformed_df.filter("unknown_mask == False")
+                if transformed_df.rdd.isEmpty():
+                    warnings.warn(
+                        f"You are trying to transform dataframe with all values are unknown for {self._col}, "
+                        "with `handle_unknown_strategy=drop` leads to empty dataframe",
+                        LabelEncoderTransformWarning,
+                    )
+            elif self._handle_unknown == "error":
                 collected_list = transformed_df.filter("unknown_mask == True").select(self._col).distinct().collect()
                 unique_labels = [row[self._col] for row in collected_list]
                 msg = f"Found unknown labels {unique_labels} in column {self._col} during transform"
                 raise ValueError(msg)
-            if default_value:
-                transformed_df = transformed_df.fillna({self._target_col: default_value})
+            else:
+                if default_value:
+                    transformed_df = transformed_df.fillna({self._target_col: default_value})
 
         result_df = transformed_df.drop(self._col, "unknown_mask").withColumnRenamed(self._target_col, self._col)
         return result_df
@@ -338,20 +361,27 @@ class LabelEncodingRule(BaseLabelEncodingRule):
             [list(self.get_mapping().keys()), list(self.get_mapping().values())],
             schema=[self._col, self._target_col],
         )
-        mapping_on_polars = mapping_on_polars.with_columns(
-            pl.col(self._col).cast(df.get_column(self._col).dtype)
-        )
+        mapping_on_polars = mapping_on_polars.with_columns(pl.col(self._col).cast(df.get_column(self._col).dtype))
         transformed_df = df.join(mapping_on_polars, on=self._col, how="left").with_columns(
             pl.col(self._target_col).is_null().alias("unknown_mask")
         )
         unknown_df = transformed_df.filter(pl.col("unknown_mask"))
         if not unknown_df.is_empty():
-            if self._handle_unknown == "error":
+            if self._handle_unknown == "drop":
+                transformed_df = transformed_df.filter(pl.col("unknown_mask") == "false")
+                if transformed_df.is_empty():
+                    warnings.warn(
+                        f"You are trying to transform dataframe with all values are unknown for {self._col}, "
+                        "with `handle_unknown_strategy=drop` leads to empty dataframe",
+                        LabelEncoderTransformWarning,
+                    )
+            elif self._handle_unknown == "error":
                 unique_labels = unknown_df.select(self._col).unique().to_series().to_list()
                 msg = f"Found unknown labels {unique_labels} in column {self._col} during transform"
                 raise ValueError(msg)
-            if default_value:
-                transformed_df = transformed_df.with_columns(pl.col(self._target_col).fill_null(default_value))
+            else:
+                if default_value:
+                    transformed_df = transformed_df.with_columns(pl.col(self._target_col).fill_null(default_value))
 
         result_df = transformed_df.drop([self._col, "unknown_mask"]).rename({self._target_col: self._col})
         return result_df
@@ -364,18 +394,20 @@ class LabelEncodingRule(BaseLabelEncodingRule):
         :returns: transformed dataframe.
         """
         if self._mapping is None:
-            raise RuntimeError("Label encoder is not fitted")
+            msg = "Label encoder is not fitted"
+            raise RuntimeError(msg)
 
         default_value = len(self._mapping) if self._default_value == "last" else self._default_value
 
         if isinstance(df, PandasDataFrame):
-            transformed_df = self._transform_pandas(df, default_value)  # type: ignore
+            transformed_df = self._transform_pandas(df, default_value)
         elif isinstance(df, SparkDataFrame):
-            transformed_df = self._transform_spark(df, default_value)  # type: ignore
+            transformed_df = self._transform_spark(df, default_value)
         elif isinstance(df, PolarsDataFrame):
-            transformed_df = self._transform_polars(df, default_value)  # type: ignore
+            transformed_df = self._transform_polars(df, default_value)
         else:
-            raise NotImplementedError(f"{self.__class__.__name__} is not implemented for {type(df)}")
+            msg = f"{self.__class__.__name__} is not implemented for {type(df)}"
+            raise NotImplementedError(msg)
         return transformed_df
 
     def _inverse_transform_pandas(self, df: PandasDataFrame) -> PandasDataFrame:
@@ -414,7 +446,8 @@ class LabelEncodingRule(BaseLabelEncodingRule):
         :returns: initial dataframe.
         """
         if self._mapping is None:
-            raise RuntimeError("Label encoder is not fitted")
+            msg = "Label encoder is not fitted"
+            raise RuntimeError(msg)
 
         if isinstance(df, PandasDataFrame):
             transformed_df = self._inverse_transform_pandas(df)
@@ -423,7 +456,8 @@ class LabelEncodingRule(BaseLabelEncodingRule):
         elif isinstance(df, PolarsDataFrame):
             transformed_df = self._inverse_transform_polars(df)
         else:
-            raise NotImplementedError(f"{self.__class__.__name__} is not implemented for {type(df)}")
+            msg = f"{self.__class__.__name__} is not implemented for {type(df)}"
+            raise NotImplementedError(msg)
         return transformed_df
 
     def set_default_value(self, default_value: Optional[Union[int, str]]) -> None:
@@ -434,7 +468,8 @@ class LabelEncodingRule(BaseLabelEncodingRule):
         :param default_value: default value.
         """
         if default_value is not None and not isinstance(default_value, int) and default_value != "last":
-            raise ValueError("Default value should be None, int or 'last'")
+            msg = "Default value should be None, int or 'last'"
+            raise ValueError(msg)
         self._default_value = default_value
 
     def set_handle_unknown(self, handle_unknown: HandleUnknownStrategies) -> None:
@@ -444,7 +479,8 @@ class LabelEncodingRule(BaseLabelEncodingRule):
         :param handle_unknown: handle unknown strategy.
         """
         if handle_unknown not in self._HANDLE_UNKNOWN_STRATEGIES:
-            raise ValueError(f"handle_unknown should be either 'error' or 'use_default_value', got {handle_unknown}.")
+            msg = f"handle_unknown should be either 'error' or 'use_default_value', got {handle_unknown}."
+            raise ValueError(msg)
         self._handle_unknown = handle_unknown
 
 
@@ -582,11 +618,12 @@ class LabelEncoder:
             If ``str`` value, should be \"last\" only, then fill by n_classes number.
             Default ``None``.
         """
-        columns = [i.column for i in self.rules]  # pylint: disable=W0212
+        columns = [i.column for i in self.rules]
         for column, handle_unknown in handle_unknown_rules.items():
             if column not in columns:
-                raise ValueError(f"Column {column} not found.")
-            rule = list(filter(lambda x: x.column == column, self.rules))  # pylint: disable = W0212, W0640
+                msg = f"Column {column} not found."
+                raise ValueError(msg)
+            rule = list(filter(lambda x: x.column == column, self.rules))
             rule[0].set_handle_unknown(handle_unknown)
 
     def set_default_values(self, default_value_rules: Dict[str, Optional[Union[int, str]]]) -> None:
@@ -605,9 +642,10 @@ class LabelEncoder:
             to the value given for the parameter default_value.
             Default: ``error``.
         """
-        columns = [i.column for i in self.rules]  # pylint: disable=W0212
+        columns = [i.column for i in self.rules]
         for column, default_value in default_value_rules.items():
             if column not in columns:
-                raise ValueError(f"Column {column} not found.")
-            rule = list(filter(lambda x: x.column == column, self.rules))  # pylint: disable = W0212, W0640
+                msg = f"Column {column} not found."
+                raise ValueError(msg)
+            rule = list(filter(lambda x: x.column == column, self.rules))
             rule[0].set_default_value(default_value)

replay/preprocessing/sessionizer.py

@@ -10,7 +10,6 @@ if PYSPARK_AVAILABLE:
     from pyspark.sql.window import Window
 
 
-# pylint: disable=too-many-instance-attributes, too-few-public-methods
 class Sessionizer:
     """
     Create and filter sessions from given interactions.
@@ -51,7 +50,6 @@ class Sessionizer:
     <BLANKLINE>
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         user_column: str = "user_id",
@@ -191,7 +189,6 @@ class Sessionizer:
                 Window.partitionBy(self.user_column).orderBy(sf.col(self.time_column), sf.col("timestamp_diff").desc())
             ),
         )
-        # data_with_sum_timediff.cache()
 
         grouped_users = data_with_sum_timediff.groupBy(self.user_column).count()
         grouped_users_with_cumsum = grouped_users.withColumn(
@@ -212,11 +209,9 @@ class Sessionizer:
             )
         )
 
-        # data_with_sum_timediff.unpersist()
         return result
 
     def _filter_sessions(self, interactions: DataFrameLike) -> DataFrameLike:
-        # interactions.cache()
         if isinstance(interactions, SparkDataFrame):
             return self._filter_sessions_spark(interactions)
 
@@ -254,8 +249,6 @@ class Sessionizer:
             entries_counter.select(self.session_column), self.session_column, how="right"
         )
 
-        # filtered_interactions.cache()
-
         nunique = filtered_interactions.groupby(self.user_column).agg(
             sf.expr("count(distinct session_id)").alias("nunique")
         )
@@ -284,9 +277,6 @@ class Sessionizer:
         result = self._filter_sessions(result)
         columns_order += [self.session_column]
 
-        if isinstance(result, SparkDataFrame):
-            result = result.select(*columns_order)
-        else:
-            result = result[columns_order]
+        result = result.select(*columns_order) if isinstance(result, SparkDataFrame) else result[columns_order]
 
         return result

replay/scenarios/fallback.py

@@ -1,16 +1,14 @@
-# pylint: disable=protected-access
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from replay.data import Dataset
-from replay.preprocessing.filters import MinCountFilter
 from replay.metrics import NDCG, Metric
 from replay.models import PopRec
 from replay.models.base_rec import BaseRecommender
+from replay.preprocessing.filters import MinCountFilter
 from replay.utils import SparkDataFrame
 from replay.utils.spark_utils import fallback, get_unique_entities
 
 
-# pylint: disable=too-many-instance-attributes
 class Fallback(BaseRecommender):
     """Fill missing recommendations using fallback model.
     Behaves like a recommender and have the same interface."""
@@ -33,16 +31,15 @@ class Fallback(BaseRecommender):
         self.threshold = threshold
         self.hot_queries = None
         self.main_model = main_model
-        # pylint: disable=invalid-name
         self.fb_model = fallback_model
 
-    # TO DO: add save/load for scenarios
+    # TODO: add save/load for scenarios
     @property
     def _init_args(self):
         return {"threshold": self.threshold}
 
     def __str__(self):
-        return f"Fallback_{str(self.main_model)}_{str(self.fb_model)}"
+        return f"Fallback_{self.main_model!s}_{self.fb_model!s}"
 
     def fit(
         self,
@@ -67,7 +64,6 @@ class Fallback(BaseRecommender):
         self._fit_wrap(hot_dataset)
         self.fb_model._fit_wrap(dataset)
 
-    # pylint: disable=too-many-arguments
     def predict(
         self,
         dataset: Dataset,
@@ -125,7 +121,6 @@ class Fallback(BaseRecommender):
         pred = fallback(hot_pred, cold_pred, k)
         return pred
 
-    # pylint: disable=too-many-arguments, too-many-locals
     def optimize(
         self,
         train_dataset: Dataset,

replay/splitters/base_splitter.py

@@ -1,4 +1,6 @@
+import json
 from abc import ABC, abstractmethod
+from pathlib import Path
 from typing import Optional, Tuple
 
 import polars as pl
@@ -7,19 +9,20 @@ from replay.utils import (
     PYSPARK_AVAILABLE,
     DataFrameLike,
     PandasDataFrame,
-    SparkDataFrame,
     PolarsDataFrame,
+    SparkDataFrame,
 )
 
 if PYSPARK_AVAILABLE:
-    from pyspark.sql import Window
-    from pyspark.sql import functions as sf
+    from pyspark.sql import (
+        Window,
+        functions as sf,
+    )
 
 
 SplitterReturnType = Tuple[DataFrameLike, DataFrameLike]
 
 
-# pylint: disable=too-few-public-methods, too-many-instance-attributes
 class Splitter(ABC):
     """Base class"""
 
@@ -33,7 +36,6 @@ class Splitter(ABC):
         "session_id_processing_strategy",
     ]
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         drop_cold_items: bool = False,
@@ -68,17 +70,43 @@ class Splitter(ABC):
     def _init_args(self):
         return {name: getattr(self, name) for name in self._init_arg_names}
 
+    def save(self, path: str) -> None:
+        """
+        Method for saving splitter in `.replay` directory.
+        """
+        base_path = Path(path).with_suffix(".replay").resolve()
+        base_path.mkdir(parents=True, exist_ok=True)
+
+        splitter_dict = {}
+        splitter_dict["init_args"] = self._init_args
+        splitter_dict["_class_name"] = str(self)
+
+        with open(base_path / "init_args.json", "w+") as file:
+            json.dump(splitter_dict, file)
+
+    @classmethod
+    def load(cls, path: str) -> "Splitter":
+        """
+        Method for loading splitter from `.replay` directory.
+        """
+        base_path = Path(path).with_suffix(".replay").resolve()
+        with open(base_path / "init_args.json", "r") as file:
+            splitter_dict = json.loads(file.read())
+        splitter = cls(**splitter_dict["init_args"])
+
+        return splitter
+
     def __str__(self):
         return type(self).__name__
 
-    # pylint: disable=too-many-arguments
     def _drop_cold_items_and_users(
         self,
         train: DataFrameLike,
         test: DataFrameLike,
     ) -> DataFrameLike:
         if isinstance(train, type(test)) is False:
-            raise TypeError("Train and test dataframes must have consistent types")
+            msg = "Train and test dataframes must have consistent types"
+            raise TypeError(msg)
 
         if isinstance(test, SparkDataFrame):
             return self._drop_cold_items_and_users_from_spark(train, test)
@@ -105,7 +133,6 @@ class Splitter(ABC):
         train: SparkDataFrame,
         test: SparkDataFrame,
     ) -> SparkDataFrame:
-
         if self.drop_cold_items:
             train_tmp = train.select(sf.col(self.item_column).alias("item")).distinct()
             test = test.join(train_tmp, train_tmp["item"] == test[self.item_column]).drop("item")
@@ -121,7 +148,6 @@ class Splitter(ABC):
         train: PolarsDataFrame,
         test: PolarsDataFrame,
     ) -> PolarsDataFrame:
-
         if self.drop_cold_items:
             train_tmp = train.select(self.item_column).unique()
             test = test.join(train_tmp, on=self.item_column)
@@ -164,9 +190,9 @@ class Splitter(ABC):
     def _recalculate_with_session_id_column_pandas(self, data: PandasDataFrame) -> PandasDataFrame:
         agg_function_name = "first" if self.session_id_processing_strategy == "train" else "last"
         res = data.copy()
-        res["is_test"] = res.groupby(
-            [self.query_column, self.session_id_column]
-        )["is_test"].transform(agg_function_name)
+        res["is_test"] = res.groupby([self.query_column, self.session_id_column])["is_test"].transform(
+            agg_function_name
+        )
 
         return res
 
@@ -176,7 +202,7 @@ class Splitter(ABC):
             "is_test",
             agg_function("is_test").over(
                 Window.orderBy(self.timestamp_column)
-                .partitionBy(self.query_column, self.session_id_column)  # type: ignore
+                .partitionBy(self.query_column, self.session_id_column)
                 .rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
             ),
         )
@@ -186,7 +212,9 @@ class Splitter(ABC):
     def _recalculate_with_session_id_column_polars(self, data: PolarsDataFrame) -> PolarsDataFrame:
         agg_function = pl.Expr.first if self.session_id_processing_strategy == "train" else pl.Expr.last
         res = data.with_columns(
-            agg_function(pl.col("is_test").sort_by(self.timestamp_column))
-            .over([self.query_column, self.session_id_column]))
+            agg_function(pl.col("is_test").sort_by(self.timestamp_column)).over(
+                [self.query_column, self.session_id_column]
+            )
+        )
 
         return res