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

replay/splitters/cold_user_random_splitter.py

@@ -1,7 +1,7 @@
-from typing import Optional, Union
+from typing import Optional, Tuple
+
 import polars as pl
 
-from .base_splitter import Splitter, SplitterReturnType
 from replay.utils import (
     PYSPARK_AVAILABLE,
     DataFrameLike,
@@ -10,11 +10,12 @@ from replay.utils import (
     SparkDataFrame,
 )
 
+from .base_splitter import Splitter, SplitterReturnType
+
 if PYSPARK_AVAILABLE:
     import pyspark.sql.functions as sf
 
 
-# pylint: disable=too-few-public-methods, duplicate-code
 class ColdUserRandomSplitter(Splitter):
     """
     Test set consists of all actions of randomly chosen users.
@@ -28,7 +29,6 @@ class ColdUserRandomSplitter(Splitter):
         "item_column",
     ]
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         test_size: float,
@@ -52,14 +52,13 @@ class ColdUserRandomSplitter(Splitter):
         )
         self.seed = seed
         if test_size <= 0 or test_size >= 1:
-            raise ValueError("test_size must between 0 and 1")
+            msg = "test_size must between 0 and 1"
+            raise ValueError(msg)
         self.test_size = test_size
 
     def _core_split_pandas(
-        self,
-        interactions: PandasDataFrame,
-        threshold: float
-    ) -> Union[PandasDataFrame, PandasDataFrame]:
+        self, interactions: PandasDataFrame, threshold: float
+    ) -> Tuple[PandasDataFrame, PandasDataFrame]:
         users = PandasDataFrame(interactions[self.query_column].unique(), columns=[self.query_column])
         train_users = users.sample(frac=(1 - threshold), random_state=self.seed)
         train_users["is_test"] = False
@@ -74,19 +73,15 @@ class ColdUserRandomSplitter(Splitter):
         return train, test
 
     def _core_split_spark(
-        self,
-        interactions: SparkDataFrame,
-        threshold: float
-    ) -> Union[SparkDataFrame, SparkDataFrame]:
+        self, interactions: SparkDataFrame, threshold: float
+    ) -> Tuple[SparkDataFrame, SparkDataFrame]:
         users = interactions.select(self.query_column).distinct()
         train_users, _ = users.randomSplit(
             [1 - threshold, threshold],
             seed=self.seed,
         )
         interactions = interactions.join(
-            train_users.withColumn("is_test", sf.lit(False)),
-            on=self.query_column,
-            how="left"
+            train_users.withColumn("is_test", sf.lit(False)), on=self.query_column, how="left"
         ).na.fill({"is_test": True})
 
         train = interactions.filter(~sf.col("is_test")).drop("is_test")
@@ -95,27 +90,18 @@ class ColdUserRandomSplitter(Splitter):
         return train, test
 
     def _core_split_polars(
-        self,
-        interactions: PolarsDataFrame,
-        threshold: float
-    ) -> Union[PolarsDataFrame, PolarsDataFrame]:
+        self, interactions: PolarsDataFrame, threshold: float
+    ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
         train_users = (
-            interactions
-            .select(self.query_column)
+            interactions.select(self.query_column)
             .unique()
             .sample(fraction=(1 - threshold), seed=self.seed)
             .with_columns(pl.lit(False).alias("is_test"))
         )
 
-        interactions = (
-            interactions
-            .join(
-                train_users,
-                on=self.query_column, how="left")
-            .fill_null(True)
-        )
+        interactions = interactions.join(train_users, on=self.query_column, how="left").fill_null(True)
 
-        train = interactions.filter(~pl.col("is_test")).drop("is_test")  # pylint: disable=invalid-unary-operand-type
+        train = interactions.filter(~pl.col("is_test")).drop("is_test")
         test = interactions.filter(pl.col("is_test")).drop("is_test")
         return train, test
 
@@ -127,4 +113,5 @@ class ColdUserRandomSplitter(Splitter):
         if isinstance(interactions, PolarsDataFrame):
             return self._core_split_polars(interactions, self.test_size)
 
-        raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+        msg = f"{self} is not implemented for {type(interactions)}"
+        raise NotImplementedError(msg)

replay/splitters/k_folds.py

@@ -1,9 +1,11 @@
 from typing import Literal, Optional, Tuple
+
 import polars as pl
 
-from .base_splitter import Splitter, SplitterReturnType
 from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
+from .base_splitter import Splitter, SplitterReturnType
+
 if PYSPARK_AVAILABLE:
     import pyspark.sql.functions as sf
     from pyspark.sql import Window
@@ -11,11 +13,11 @@ if PYSPARK_AVAILABLE:
 StrategyName = Literal["query"]
 
 
-# pylint: disable=too-few-public-methods
 class KFolds(Splitter):
     """
     Splits interactions inside each query into folds at random.
     """
+
     _init_arg_names = [
         "n_folds",
         "strategy",
@@ -29,7 +31,6 @@ class KFolds(Splitter):
         "session_id_processing_strategy",
     ]
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         n_folds: Optional[int] = 5,
@@ -64,11 +65,12 @@ class KFolds(Splitter):
             item_column=item_column,
             timestamp_column=timestamp_column,
             session_id_column=session_id_column,
-            session_id_processing_strategy=session_id_processing_strategy
+            session_id_processing_strategy=session_id_processing_strategy,
         )
         self.n_folds = n_folds
         if strategy not in {"query"}:
-            raise ValueError(f"Wrong splitter parameter: {strategy}")
+            msg = f"Wrong splitter parameter: {strategy}"
+            raise ValueError(msg)
         self.strategy = strategy
         self.seed = seed
 
@@ -85,16 +87,10 @@ class KFolds(Splitter):
         dataframe = interactions.withColumn("_rand", sf.rand(self.seed))
         dataframe = dataframe.withColumn(
             "fold",
-            sf.row_number().over(
-                Window.partitionBy(self.query_column).orderBy("_rand")
-            )
-            % self.n_folds,
+            sf.row_number().over(Window.partitionBy(self.query_column).orderBy("_rand")) % self.n_folds,
         ).drop("_rand")
         for i in range(self.n_folds):
-            dataframe = dataframe.withColumn(
-                "is_test",
-                sf.when(sf.col("fold") == i, True).otherwise(False)
-            )
+            dataframe = dataframe.withColumn("is_test", sf.when(sf.col("fold") == i, True).otherwise(False))
             if self.session_id_column:
                 dataframe = self._recalculate_with_session_id_column(dataframe)
 
@@ -122,28 +118,21 @@ class KFolds(Splitter):
     def _query_split_polars(self, interactions: PolarsDataFrame) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
         dataframe = interactions.sample(fraction=1, shuffle=True, seed=self.seed).sort(self.query_column)
         dataframe = dataframe.with_columns(
-            (pl.cum_count(self.query_column).over(self.query_column) % self.n_folds)
-            .alias("fold")
+            (pl.cum_count(self.query_column).over(self.query_column) % self.n_folds).alias("fold")
         )
         for i in range(self.n_folds):
             dataframe = dataframe.with_columns(
-                pl.when(
-                    pl.col("fold") == i
-                )
-                .then(True)
-                .otherwise(False)
-                .alias("is_test")
+                pl.when(pl.col("fold") == i).then(True).otherwise(False).alias("is_test")
             )
             if self.session_id_column:
                 dataframe = self._recalculate_with_session_id_column(dataframe)
 
-            train = dataframe.filter(~pl.col("is_test")).drop("is_test", "fold")  # pylint: disable=invalid-unary-operand-type
+            train = dataframe.filter(~pl.col("is_test")).drop("is_test", "fold")
             test = dataframe.filter(pl.col("is_test")).drop("is_test", "fold")
 
             test = self._drop_cold_items_and_users(train, test)
             yield train, test
 
-    # pylint: disable=inconsistent-return-statements
     def _core_split(self, interactions: DataFrameLike) -> SplitterReturnType:
         if self.strategy == "query":
             if isinstance(interactions, SparkDataFrame):
@@ -153,4 +142,5 @@ class KFolds(Splitter):
             if isinstance(interactions, PolarsDataFrame):
                 return self._query_split_polars(interactions)
 
-            raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+            msg = f"{self} is not implemented for {type(interactions)}"
+            raise NotImplementedError(msg)

replay/splitters/last_n_splitter.py

@@ -4,9 +4,10 @@ import numpy as np
 import pandas as pd
 import polars as pl
 
-from .base_splitter import Splitter
 from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
+from .base_splitter import Splitter
+
 if PYSPARK_AVAILABLE:
     import pyspark.sql.functions as sf
     from pyspark.sql import Window
@@ -14,7 +15,6 @@ if PYSPARK_AVAILABLE:
 StrategyName = Literal["interactions", "timedelta"]
 
 
-# pylint: disable=too-few-public-methods
 class LastNSplitter(Splitter):
     """
     Split interactions by last N interactions/timedelta per user.
@@ -88,10 +88,11 @@ class LastNSplitter(Splitter):
     14        3        2 2020-01-05
     <BLANKLINE>
     """
+
     _init_arg_names = [
         "N",
         "divide_column",
-        "timestamp_col_format",
+        "time_column_format",
         "strategy",
         "drop_cold_users",
         "drop_cold_items",
@@ -102,10 +103,9 @@ class LastNSplitter(Splitter):
         "session_id_processing_strategy",
     ]
 
-    # pylint: disable=invalid-name, too-many-arguments
     def __init__(
         self,
-        N: int,
+        N: int,  # noqa: N803
         divide_column: str = "query_id",
         time_column_format: str = "yyyy-MM-dd HH:mm:ss",
         strategy: StrategyName = "interactions",
@@ -147,7 +147,8 @@ class LastNSplitter(Splitter):
             default: ``test``.
         """
         if strategy not in ["interactions", "timedelta"]:
-            raise ValueError("strategy must be equal 'interactions' or 'timedelta'")
+            msg = "strategy must be equal 'interactions' or 'timedelta'"
+            raise ValueError(msg)
         super().__init__(
             drop_cold_users=drop_cold_users,
             drop_cold_items=drop_cold_items,
@@ -160,9 +161,9 @@ class LastNSplitter(Splitter):
         self.N = N
         self.strategy = strategy
         self.divide_column = divide_column
-        self.timestamp_col_format = None
+        self.time_column_format = None
         if self.strategy == "timedelta":
-            self.timestamp_col_format = time_column_format
+            self.time_column_format = time_column_format
 
     def _add_time_partition(self, interactions: DataFrameLike) -> DataFrameLike:
         if isinstance(interactions, SparkDataFrame):
@@ -172,7 +173,8 @@ class LastNSplitter(Splitter):
         if isinstance(interactions, PolarsDataFrame):
             return self._add_time_partition_to_polars(interactions)
 
-        raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+        msg = f"{self} is not implemented for {type(interactions)}"
+        raise NotImplementedError(msg)
 
     def _add_time_partition_to_pandas(self, interactions: PandasDataFrame) -> PandasDataFrame:
         res = interactions.copy(deep=True)
@@ -191,8 +193,7 @@ class LastNSplitter(Splitter):
 
     def _add_time_partition_to_polars(self, interactions: PolarsDataFrame) -> PolarsDataFrame:
         res = interactions.sort(self.timestamp_column).with_columns(
-            pl.col(self.divide_column).cumcount().over(pl.col(self.divide_column))
-            .alias("row_num")
+            pl.col(self.divide_column).cumcount().over(pl.col(self.divide_column)).alias("row_num")
         )
 
         return res
@@ -205,7 +206,8 @@ class LastNSplitter(Splitter):
         if isinstance(interactions, PolarsDataFrame):
             return self._to_unix_timestamp_polars(interactions)
 
-        raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+        msg = f"{self} is not implemented for {type(interactions)}"
+        raise NotImplementedError(msg)
 
     def _to_unix_timestamp_pandas(self, interactions: PandasDataFrame) -> PandasDataFrame:
         time_column_type = dict(interactions.dtypes)[self.timestamp_column]
@@ -221,7 +223,7 @@ class LastNSplitter(Splitter):
         time_column_type = dict(interactions.dtypes)[self.timestamp_column]
         if time_column_type == "date":
             interactions = interactions.withColumn(
-                self.timestamp_column, sf.unix_timestamp(self.timestamp_column, self.timestamp_col_format)
+                self.timestamp_column, sf.unix_timestamp(self.timestamp_column, self.time_column_format)
             )
 
         return interactions
@@ -233,20 +235,19 @@ class LastNSplitter(Splitter):
 
         return interactions
 
-    # pylint: disable=invalid-name
-    def _partial_split_interactions(self, interactions: DataFrameLike, N: int) -> Tuple[DataFrameLike, DataFrameLike]:
+    def _partial_split_interactions(self, interactions: DataFrameLike, n: int) -> Tuple[DataFrameLike, DataFrameLike]:
         res = self._add_time_partition(interactions)
         if isinstance(interactions, SparkDataFrame):
-            return self._partial_split_interactions_spark(res, N)
+            return self._partial_split_interactions_spark(res, n)
         if isinstance(interactions, PandasDataFrame):
-            return self._partial_split_interactions_pandas(res, N)
-        return self._partial_split_interactions_polars(res, N)
+            return self._partial_split_interactions_pandas(res, n)
+        return self._partial_split_interactions_polars(res, n)
 
     def _partial_split_interactions_pandas(
-        self, interactions: PandasDataFrame, N: int
+        self, interactions: PandasDataFrame, n: int
     ) -> Tuple[PandasDataFrame, PandasDataFrame]:
         interactions["count"] = interactions.groupby(self.divide_column, sort=False)[self.divide_column].transform(len)
-        interactions["is_test"] = interactions["row_num"] > (interactions["count"] - float(N))
+        interactions["is_test"] = interactions["row_num"] > (interactions["count"] - float(n))
         if self.session_id_column:
             interactions = self._recalculate_with_session_id_column(interactions)
 
@@ -256,14 +257,14 @@ class LastNSplitter(Splitter):
         return train, test
 
     def _partial_split_interactions_spark(
-        self, interactions: SparkDataFrame, N: int
+        self, interactions: SparkDataFrame, n: int
     ) -> Tuple[SparkDataFrame, SparkDataFrame]:
         interactions = interactions.withColumn(
             "count", sf.count(self.timestamp_column).over(Window.partitionBy(self.divide_column))
         )
         # float(n) - because DataFrame.filter is changing order
         # of sorted DataFrame to descending
-        interactions = interactions.withColumn("is_test", sf.col("row_num") > sf.col("count") - sf.lit(float(N)))
+        interactions = interactions.withColumn("is_test", sf.col("row_num") > sf.col("count") - sf.lit(float(n)))
         if self.session_id_column:
             interactions = self._recalculate_with_session_id_column(interactions)
 
@@ -273,27 +274,22 @@ class LastNSplitter(Splitter):
         return train, test
 
     def _partial_split_interactions_polars(
-        self, interactions: PolarsDataFrame, N: int
+        self, interactions: PolarsDataFrame, n: int
     ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
         interactions = interactions.with_columns(
-            pl.col(self.timestamp_column).count().over(self.divide_column)
-            .alias("count")
-        )
-        interactions = interactions.with_columns(
-            (pl.col("row_num") > (pl.col("count") - N))
-            .alias("is_test")
+            pl.col(self.timestamp_column).count().over(self.divide_column).alias("count")
         )
+        interactions = interactions.with_columns((pl.col("row_num") > (pl.col("count") - n)).alias("is_test"))
         if self.session_id_column:
             interactions = self._recalculate_with_session_id_column(interactions)
 
-        train = interactions.filter(~pl.col("is_test")).drop("row_num", "count", "is_test")  # pylint: disable=invalid-unary-operand-type
+        train = interactions.filter(~pl.col("is_test")).drop("row_num", "count", "is_test")
         test = interactions.filter(pl.col("is_test")).drop("row_num", "count", "is_test")
 
         return train, test
 
     def _partial_split_timedelta(
-        self,
-        interactions: DataFrameLike, timedelta: int
+        self, interactions: DataFrameLike, timedelta: int
     ) -> Tuple[DataFrameLike, DataFrameLike]:
         if isinstance(interactions, SparkDataFrame):
             return self._partial_split_timedelta_spark(interactions, timedelta)
@@ -341,22 +337,16 @@ class LastNSplitter(Splitter):
     def _partial_split_timedelta_polars(
         self, interactions: PolarsDataFrame, timedelta: int
     ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
-        res = (
-            interactions
-            .with_columns(
-                (pl.col(self.timestamp_column).max().over(self.divide_column) - pl.col(self.timestamp_column))
-                .alias("diff_timestamp")
+        res = interactions.with_columns(
+            (pl.col(self.timestamp_column).max().over(self.divide_column) - pl.col(self.timestamp_column)).alias(
+                "diff_timestamp"
             )
-            .with_columns(
-                (pl.col("diff_timestamp") < timedelta)
-                .alias("is_test")
-            )
-        )
+        ).with_columns((pl.col("diff_timestamp") < timedelta).alias("is_test"))
 
         if self.session_id_column:
             res = self._recalculate_with_session_id_column(res)
 
-        train = res.filter(~pl.col("is_test")).drop("diff_timestamp", "is_test")  # pylint: disable=invalid-unary-operand-type
+        train = res.filter(~pl.col("is_test")).drop("diff_timestamp", "is_test")
         test = res.filter(pl.col("is_test")).drop("diff_timestamp", "is_test")
 
         return train, test

replay/splitters/new_users_splitter.py

@@ -1,15 +1,16 @@
-from typing import Optional, Union
+from typing import Optional, Tuple
+
 import polars as pl
 
-from .base_splitter import Splitter, SplitterReturnType
 from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
+from .base_splitter import Splitter, SplitterReturnType
+
 if PYSPARK_AVAILABLE:
     import pyspark.sql.functions as sf
     from pyspark.sql import Window
 
 
-# pylint: disable=too-few-public-methods, duplicate-code
 class NewUsersSplitter(Splitter):
     """
     Only new users will be assigned to test set.
@@ -63,7 +64,6 @@ class NewUsersSplitter(Splitter):
         "session_id_processing_strategy",
     ]
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         test_size: float,
@@ -91,24 +91,23 @@ class NewUsersSplitter(Splitter):
             item_column=item_column,
             timestamp_column=timestamp_column,
             session_id_column=session_id_column,
-            session_id_processing_strategy=session_id_processing_strategy
+            session_id_processing_strategy=session_id_processing_strategy,
         )
         if test_size < 0 or test_size > 1:
-            raise ValueError("test_size must between 0 and 1")
+            msg = "test_size must between 0 and 1"
+            raise ValueError(msg)
         self.test_size = test_size
 
     def _core_split_pandas(
-        self,
-        interactions: PandasDataFrame,
-        threshold: float
-    ) -> Union[PandasDataFrame, PandasDataFrame]:
-        start_date_by_user = interactions.groupby(self.query_column).agg(
-            _start_dt_by_user=(self.timestamp_column, "min")
-        ).reset_index()
+        self, interactions: PandasDataFrame, threshold: float
+    ) -> Tuple[PandasDataFrame, PandasDataFrame]:
+        start_date_by_user = (
+            interactions.groupby(self.query_column).agg(_start_dt_by_user=(self.timestamp_column, "min")).reset_index()
+        )
         test_start_date = (
-            start_date_by_user
-            .groupby("_start_dt_by_user")
-            .agg(_num_users_by_start_date=(self.query_column, "count")).reset_index()
+            start_date_by_user.groupby("_start_dt_by_user")
+            .agg(_num_users_by_start_date=(self.query_column, "count"))
+            .reset_index()
             .sort_values(by="_start_dt_by_user", ascending=False)
         )
         test_start_date["_cum_num_users_to_dt"] = test_start_date["_num_users_by_start_date"].cumsum()
@@ -120,9 +119,7 @@ class NewUsersSplitter(Splitter):
 
         train = interactions[interactions[self.timestamp_column] < test_start]
         test = interactions.merge(
-            start_date_by_user[start_date_by_user["_start_dt_by_user"] >= test_start],
-            how="inner",
-            on=self.query_column
+            start_date_by_user[start_date_by_user["_start_dt_by_user"] >= test_start], how="inner", on=self.query_column
         ).drop(columns=["_start_dt_by_user"])
 
         if self.session_id_column:
@@ -136,10 +133,8 @@ class NewUsersSplitter(Splitter):
         return train, test
 
     def _core_split_spark(
-        self,
-        interactions: SparkDataFrame,
-        threshold: float
-    ) -> Union[SparkDataFrame, SparkDataFrame]:
+        self, interactions: SparkDataFrame, threshold: float
+    ) -> Tuple[SparkDataFrame, SparkDataFrame]:
         start_date_by_user = interactions.groupby(self.query_column).agg(
             sf.min(self.timestamp_column).alias("_start_dt_by_user")
         )
@@ -175,53 +170,33 @@ class NewUsersSplitter(Splitter):
         return train, test
 
     def _core_split_polars(
-        self,
-        interactions: PolarsDataFrame,
-        threshold: float
-    ) -> Union[PolarsDataFrame, PolarsDataFrame]:
-        start_date_by_user = (
-            interactions
-            .group_by(self.query_column).agg(
-                pl.col(self.timestamp_column).min()
-                .alias("_start_dt_by_user")
-            )
+        self, interactions: PolarsDataFrame, threshold: float
+    ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
+        start_date_by_user = interactions.group_by(self.query_column).agg(
+            pl.col(self.timestamp_column).min().alias("_start_dt_by_user")
         )
         test_start_date = (
-            start_date_by_user
-            .group_by("_start_dt_by_user").agg(
-                pl.col(self.query_column).count()
-                .alias("_num_users_by_start_date")
-            )
+            start_date_by_user.group_by("_start_dt_by_user")
+            .agg(pl.col(self.query_column).count().alias("_num_users_by_start_date"))
             .sort("_start_dt_by_user", descending=True)
             .with_columns(
-                pl.col("_num_users_by_start_date").cum_sum()
-                .alias("cum_sum_users"),
+                pl.col("_num_users_by_start_date").cum_sum().alias("cum_sum_users"),
             )
-            .filter(
-                pl.col("cum_sum_users") >= pl.col("cum_sum_users").max() * threshold
-            )
-            ["_start_dt_by_user"]
+            .filter(pl.col("cum_sum_users") >= pl.col("cum_sum_users").max() * threshold)["_start_dt_by_user"]
             .max()
         )
 
         train = interactions.filter(pl.col(self.timestamp_column) < test_start_date)
         test = interactions.join(
-            start_date_by_user.filter(pl.col("_start_dt_by_user") >= test_start_date),
-            on=self.query_column,
-            how="inner"
+            start_date_by_user.filter(pl.col("_start_dt_by_user") >= test_start_date), on=self.query_column, how="inner"
         ).drop("_start_dt_by_user")
 
         if self.session_id_column:
             interactions = interactions.with_columns(
-                pl.when(
-                    pl.col(self.timestamp_column) < test_start_date
-                )
-                .then(False)
-                .otherwise(True)
-                .alias("is_test")
+                pl.when(pl.col(self.timestamp_column) < test_start_date).then(False).otherwise(True).alias("is_test")
             )
             interactions = self._recalculate_with_session_id_column(interactions)
-            train = interactions.filter(~pl.col("is_test")).drop("is_test")  # pylint: disable=invalid-unary-operand-type
+            train = interactions.filter(~pl.col("is_test")).drop("is_test")
             test = interactions.filter(pl.col("is_test")).drop("is_test")
 
         return train, test
@@ -234,4 +209,5 @@ class NewUsersSplitter(Splitter):
         if isinstance(interactions, PolarsDataFrame):
             return self._core_split_polars(interactions, self.test_size)
 
-        raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+        msg = f"{self} is not implemented for {type(interactions)}"
+        raise NotImplementedError(msg)

replay/splitters/random_splitter.py

@@ -1,10 +1,10 @@
-from typing import Optional, Union
+from typing import Optional, Tuple
 
-from .base_splitter import Splitter, SplitterReturnType
 from replay.utils import DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
+from .base_splitter import Splitter, SplitterReturnType
+
 
-# pylint: disable=too-few-public-methods, duplicate-code
 class RandomSplitter(Splitter):
     """Assign records into train and test at random."""
 
@@ -17,7 +17,6 @@ class RandomSplitter(Splitter):
         "item_column",
     ]
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         test_size: float,
@@ -25,7 +24,7 @@ class RandomSplitter(Splitter):
         drop_cold_users: bool = False,
         seed: Optional[int] = None,
         query_column: str = "query_id",
-        item_column: str = "item_id"
+        item_column: str = "item_id",
     ):
         """
         :param test_size: test size 0 to 1
@@ -39,37 +38,30 @@ class RandomSplitter(Splitter):
             drop_cold_items=drop_cold_items,
             drop_cold_users=drop_cold_users,
             query_column=query_column,
-            item_column=item_column
+            item_column=item_column,
         )
         self.seed = seed
         if test_size < 0 or test_size > 1:
-            raise ValueError("test_size must between 0 and 1")
+            msg = "test_size must between 0 and 1"
+            raise ValueError(msg)
         self.test_size = test_size
 
     def _random_split_spark(
-        self,
-        interactions: SparkDataFrame,
-        threshold: float
-    ) -> Union[SparkDataFrame, SparkDataFrame]:
-        train, test = interactions.randomSplit(
-            [1 - threshold, threshold], self.seed
-        )
+        self, interactions: SparkDataFrame, threshold: float
+    ) -> Tuple[SparkDataFrame, SparkDataFrame]:
+        train, test = interactions.randomSplit([1 - threshold, threshold], self.seed)
         return train, test
 
     def _random_split_pandas(
-        self,
-        interactions: PandasDataFrame,
-        threshold: float
-    ) -> Union[PandasDataFrame, PandasDataFrame]:
+        self, interactions: PandasDataFrame, threshold: float
+    ) -> Tuple[PandasDataFrame, PandasDataFrame]:
         train = interactions.sample(frac=(1 - threshold), random_state=self.seed)
         test = interactions.drop(train.index)
         return train, test
 
     def _random_split_polars(
-        self,
-        interactions: PolarsDataFrame,
-        threshold: float
-    ) -> Union[PolarsDataFrame, PolarsDataFrame]:
+        self, interactions: PolarsDataFrame, threshold: float
+    ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
         train_size = int(len(interactions) * (1 - threshold)) + 1
         shuffled_interactions = interactions.sample(fraction=1, shuffle=True, seed=self.seed)
         train = shuffled_interactions[:train_size]
@@ -84,4 +76,5 @@ class RandomSplitter(Splitter):
         if isinstance(interactions, PolarsDataFrame):
             return self._random_split_polars(interactions, self.test_size)
 
-        raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+        msg = f"{self} is not implemented for {type(interactions)}"
+        raise NotImplementedError(msg)