replay-rec 0.17.0rc0__py3-none-any.whl → 0.17.1__py3-none-any.whl

replay/__init__.py

@@ -1,2 +1,2 @@
 """ RecSys library """
-__version__ = "0.17.0.preview"
+__version__ = "0.17.1"

replay/data/dataset.py

@@ -3,11 +3,22 @@
 """
 from __future__ import annotations
 
-from typing import Callable, Dict, Iterable, List, Optional, Sequence
+import json
+from pathlib import Path
+from typing import Callable, Dict, Iterable, List, Optional, Sequence, Union
 
 import numpy as np
-
-from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
+from pandas import read_parquet as pd_read_parquet
+from polars import read_parquet as pl_read_parquet
+
+from replay.utils import (
+    PYSPARK_AVAILABLE,
+    DataFrameLike,
+    PandasDataFrame,
+    PolarsDataFrame,
+    SparkDataFrame,
+)
+from replay.utils.session_handler import get_spark_session
 
 from .schema import FeatureHint, FeatureInfo, FeatureSchema, FeatureSource, FeatureType
 
@@ -47,9 +58,7 @@ class Dataset:
         self._query_features = query_features
         self._item_features = item_features
 
-        self.is_pandas = isinstance(interactions, PandasDataFrame)
-        self.is_spark = isinstance(interactions, SparkDataFrame)
-        self.is_polars = isinstance(interactions, PolarsDataFrame)
+        self._assign_df_type()
 
         self._categorical_encoded = categorical_encoded
 
@@ -74,16 +83,8 @@ class Dataset:
             msg = "Interactions and query features should have the same type."
             raise TypeError(msg)
 
-        self._feature_source_map: Dict[FeatureSource, DataFrameLike] = {
-            FeatureSource.INTERACTIONS: self.interactions,
-            FeatureSource.QUERY_FEATURES: self.query_features,
-            FeatureSource.ITEM_FEATURES: self.item_features,
-        }
-
-        self._ids_feature_map: Dict[FeatureHint, DataFrameLike] = {
-            FeatureHint.QUERY_ID: self.query_features if self.query_features is not None else self.interactions,
-            FeatureHint.ITEM_ID: self.item_features if self.item_features is not None else self.interactions,
-        }
+        self._get_feature_source_map()
+        self._get_ids_source_map()
 
         self._feature_schema = self._fill_feature_schema(feature_schema)
 
@@ -92,7 +93,6 @@ class Dataset:
                 self._check_ids_consistency(hint=FeatureHint.QUERY_ID)
             if self.item_features is not None:
                 self._check_ids_consistency(hint=FeatureHint.ITEM_ID)
-
             if self._categorical_encoded:
                 self._check_encoded()
 
@@ -189,6 +189,157 @@ class Dataset:
         """
         return self._feature_schema
 
+    def _get_df_type(self) -> str:
+        """
+        :returns: Stored dataframe type.
+        """
+        if self.is_spark:
+            return "spark"
+        if self.is_pandas:
+            return "pandas"
+        if self.is_polars:
+            return "polars"
+        msg = "No known dataframe types are provided"
+        raise ValueError(msg)
+
+    def _to_parquet(self, df: DataFrameLike, path: Path) -> None:
+        """
+        Save the content of the dataframe in parquet format to the provided path.
+
+        :param df: Dataframe to save.
+        :param path: Path to save the dataframe to.
+        """
+        if self.is_spark:
+            path = str(path)
+            df = df.withColumn("idx", sf.monotonically_increasing_id())
+            df.write.mode("overwrite").parquet(path)
+        elif self.is_pandas:
+            df.to_parquet(path)
+        elif self.is_polars:
+            df.write_parquet(path)
+        else:
+            msg = """
+            _to_parquet() can only be used to save polars|pandas|spark dataframes;
+            No known dataframe types are provided
+            """
+            raise TypeError(msg)
+
+    @staticmethod
+    def _read_parquet(path: Path, mode: str) -> Union[SparkDataFrame, PandasDataFrame, PolarsDataFrame]:
+        """
+        Read the parquet file as dataframe.
+
+        :param path: The parquet file path.
+        :param mode: Dataframe type. Can be spark|pandas|polars.
+        :returns: The dataframe read from the file.
+        """
+        if mode == "spark":
+            path = str(path)
+            spark_session = get_spark_session()
+            df = spark_session.read.parquet(path)
+            if "idx" in df.columns:
+                df = df.orderBy("idx").drop("idx")
+            return df
+        if mode == "pandas":
+            df = pd_read_parquet(path)
+            if "idx" in df.columns:
+                df = df.set_index("idx").reset_index(drop=True)
+            return df
+        if mode == "polars":
+            df = pl_read_parquet(path, use_pyarrow=True)
+            if "idx" in df.columns:
+                df = df.sort("idx").drop("idx")
+            return df
+        msg = f"_read_parquet() can only be used to read polars|pandas|spark dataframes, not {mode}"
+        raise TypeError(msg)
+
+    def save(self, path: str) -> None:
+        """
+        Save the Dataset to the provided path.
+
+        :param path: Path to save the Dataset to.
+        """
+        dataset_dict = {}
+        dataset_dict["_class_name"] = self.__class__.__name__
+
+        interactions_type = self._get_df_type()
+        dataset_dict["init_args"] = {
+            "feature_schema": [],
+            "interactions": interactions_type,
+            "item_features": (interactions_type if self.item_features is not None else None),
+            "query_features": (interactions_type if self.query_features is not None else None),
+            "check_consistency": False,
+            "categorical_encoded": self._categorical_encoded,
+        }
+
+        for feature in self.feature_schema.all_features:
+            dataset_dict["init_args"]["feature_schema"].append(
+                {
+                    "column": feature.column,
+                    "feature_type": feature.feature_type.name,
+                    "feature_hint": (feature.feature_hint.name if feature.feature_hint else None),
+                }
+            )
+
+        base_path = Path(path).with_suffix(".replay").resolve()
+        base_path.mkdir(parents=True, exist_ok=True)
+
+        with open(base_path / "init_args.json", "w+") as file:
+            json.dump(dataset_dict, file)
+
+        df_data = {
+            "interactions": self.interactions,
+            "item_features": self.item_features,
+            "query_features": self.query_features,
+        }
+
+        for df_name, df in df_data.items():
+            if df is not None:
+                df_path = base_path / f"{df_name}.parquet"
+                self._to_parquet(df, df_path)
+
+    @classmethod
+    def load(
+        cls,
+        path: str,
+        dataframe_type: Optional[str] = None,
+    ) -> Dataset:
+        """
+        Load the Dataset from the provided path.
+
+        :param path: The file path
+        :dataframe_type: Dataframe type to use to store internal data.
+            Can be spark|pandas|polars|None.
+            If not provided automatically sets to the one used when the Dataset was saved.
+        :returns: Loaded Dataset.
+        """
+        base_path = Path(path).with_suffix(".replay").resolve()
+        with open(base_path / "init_args.json", "r") as file:
+            dataset_dict = json.loads(file.read())
+
+        if dataframe_type not in ["pandas", "spark", "polars", None]:
+            msg = f"Argument dataframe_type can be spark|pandas|polars|None, not {dataframe_type}"
+            raise ValueError(msg)
+
+        feature_schema_data = dataset_dict["init_args"]["feature_schema"]
+        features_list = []
+        for feature_data in feature_schema_data:
+            f_type = feature_data["feature_type"]
+            f_hint = feature_data["feature_hint"]
+            feature_data["feature_type"] = FeatureType[f_type] if f_type else None
+            feature_data["feature_hint"] = FeatureHint[f_hint] if f_hint else None
+            features_list.append(FeatureInfo(**feature_data))
+        dataset_dict["init_args"]["feature_schema"] = FeatureSchema(features_list)
+
+        for df_name in ["interactions", "query_features", "item_features"]:
+            df_type = dataset_dict["init_args"][df_name]
+            if df_type:
+                df_type = dataframe_type or df_type
+                load_path = base_path / f"{df_name}.parquet"
+                dataset_dict["init_args"][df_name] = cls._read_parquet(load_path, df_type)
+        dataset = cls(**dataset_dict["init_args"])
+        return dataset
+
     if PYSPARK_AVAILABLE:
 
         def persist(self, storage_level: StorageLevel = StorageLevel(True, True, False, True, 1)) -> None:
@@ -283,6 +434,24 @@ class Dataset:
             categorical_encoded=self._categorical_encoded,
         )
 
+    def _get_feature_source_map(self):
+        self._feature_source_map: Dict[FeatureSource, DataFrameLike] = {
+            FeatureSource.INTERACTIONS: self.interactions,
+            FeatureSource.QUERY_FEATURES: self.query_features,
+            FeatureSource.ITEM_FEATURES: self.item_features,
+        }
+
+    def _get_ids_source_map(self):
+        self._ids_feature_map: Dict[FeatureHint, DataFrameLike] = {
+            FeatureHint.QUERY_ID: self.query_features if self.query_features is not None else self.interactions,
+            FeatureHint.ITEM_ID: self.item_features if self.item_features is not None else self.interactions,
+        }
+
+    def _assign_df_type(self):
+        self.is_pandas = isinstance(self.interactions, PandasDataFrame)
+        self.is_spark = isinstance(self.interactions, SparkDataFrame)
+        self.is_polars = isinstance(self.interactions, PolarsDataFrame)
+
     def _get_cardinality(self, feature: FeatureInfo) -> Callable:
         def callback(column: str) -> int:
             if feature.feature_hint in [FeatureHint.ITEM_ID, FeatureHint.QUERY_ID]:
@@ -381,7 +550,11 @@ class Dataset:
             is_consistent = (
                 self.interactions.select(ids_column)
                 .distinct()
-                .join(features_df.select(ids_column).distinct(), on=[ids_column], how="leftanti")
+                .join(
+                    features_df.select(ids_column).distinct(),
+                    on=[ids_column],
+                    how="leftanti",
+                )
                 .count()
             ) == 0
         else:
@@ -389,7 +562,11 @@ class Dataset:
                 len(
                     self.interactions.select(ids_column)
                     .unique()
-                    .join(features_df.select(ids_column).unique(), on=ids_column, how="anti")
+                    .join(
+                        features_df.select(ids_column).unique(),
+                        on=ids_column,
+                        how="anti",
+                    )
                 )
                 == 0
             )
@@ -399,7 +576,11 @@ class Dataset:
             raise ValueError(msg)
 
     def _check_column_encoded(
-        self, data: DataFrameLike, column: str, source: FeatureSource, cardinality: Optional[int]
+        self,
+        data: DataFrameLike,
+        column: str,
+        source: FeatureSource,
+        cardinality: Optional[int],
     ) -> None:
         """
         Checks that IDs are encoded:
@@ -482,6 +663,51 @@ class Dataset:
                     feature.cardinality,
                 )
 
+    def to_pandas(self) -> None:
+        """
+        Convert internally stored dataframes to pandas.DataFrame.
+        """
+        from replay.utils.common import convert2pandas
+
+        self._interactions = convert2pandas(self._interactions)
+        if self._query_features is not None:
+            self._query_features = convert2pandas(self._query_features)
+        if self._item_features is not None:
+            self._item_features = convert2pandas(self.item_features)
+        self._get_feature_source_map()
+        self._get_ids_source_map()
+        self._assign_df_type()
+
+    def to_spark(self):
+        """
+        Convert internally stored dataframes to pyspark.sql.DataFrame.
+        """
+        from replay.utils.common import convert2spark
+
+        self._interactions = convert2spark(self._interactions)
+        if self._query_features is not None:
+            self._query_features = convert2spark(self._query_features)
+        if self._item_features is not None:
+            self._item_features = convert2spark(self._item_features)
+        self._get_feature_source_map()
+        self._get_ids_source_map()
+        self._assign_df_type()
+
+    def to_polars(self):
+        """
+        Convert internally stored dataframes to polars.DataFrame.
+        """
+        from replay.utils.common import convert2polars
+
+        self._interactions = convert2polars(self._interactions)
+        if self._query_features is not None:
+            self._query_features = convert2polars(self._query_features)
+        if self._item_features is not None:
+            self._item_features = convert2polars(self._item_features)
+        self._get_feature_source_map()
+        self._get_ids_source_map()
+        self._assign_df_type()
+
 
 def nunique(data: DataFrameLike, column: str) -> int:
     """

replay/data/nn/schema.py

@@ -408,6 +408,48 @@ class TensorSchema(Mapping[str, TensorFeatureInfo]):
             return None
         return rating_features.item().name
 
+    def _get_object_args(self) -> Dict:
+        """
+        Returns list of features represented as dictionaries.
+        """
+        features = [
+            {
+                "name": feature.name,
+                "feature_type": feature.feature_type.name,
+                "is_seq": feature.is_seq,
+                "feature_hint": feature.feature_hint.name if feature.feature_hint else None,
+                "feature_sources": [
+                    {"source": x.source.name, "column": x.column, "index": x.index} for x in feature.feature_sources
+                ]
+                if feature.feature_sources
+                else None,
+                "cardinality": feature.cardinality if feature.feature_type == FeatureType.CATEGORICAL else None,
+                "embedding_dim": feature.embedding_dim if feature.feature_type == FeatureType.CATEGORICAL else None,
+                "tensor_dim": feature.tensor_dim if feature.feature_type == FeatureType.NUMERICAL else None,
+            }
+            for feature in self.all_features
+        ]
+        return features
+
+    @classmethod
+    def _create_object_by_args(cls, args: Dict) -> "TensorSchema":
+        features_list = []
+        for feature_data in args:
+            feature_data["feature_sources"] = (
+                [
+                    TensorFeatureSource(source=FeatureSource[x["source"]], column=x["column"], index=x["index"])
+                    for x in feature_data["feature_sources"]
+                ]
+                if feature_data["feature_sources"]
+                else None
+            )
+            f_type = feature_data["feature_type"]
+            f_hint = feature_data["feature_hint"]
+            feature_data["feature_type"] = FeatureType[f_type] if f_type else None
+            feature_data["feature_hint"] = FeatureHint[f_hint] if f_hint else None
+            features_list.append(TensorFeatureInfo(**feature_data))
+        return TensorSchema(features_list)
+
     def filter(
         self,
         name: Optional[str] = None,

replay/data/nn/sequence_tokenizer.py

@@ -24,7 +24,10 @@ SequenceDataFrameLike = Union[PandasDataFrame, PolarsDataFrame]
 
 class SequenceTokenizer:
     """
-    Data tokenizer for transformers
+    Data tokenizer for transformers;
+    Encodes all categorical features (the ones marked as FeatureType.CATEGORICAL in
+    the FeatureSchema) and stores all data as items sequences (sorted by time if a
+    feature of type FeatureHint.TIMESTAMP is provided, unsorted otherwise).
     """
 
     def __init__(
@@ -278,17 +281,17 @@ class SequenceTokenizer:
         ]
 
         for tensor_feature in tensor_schema.values():
-            source = tensor_feature.feature_source
-            assert source is not None
+            for source in tensor_feature.feature_sources:
+                assert source is not None
 
-            # Some columns already added to encoder, skip them
-            if source.column in features_subset:
-                continue
+                # Some columns already added to encoder, skip them
+                if source.column in features_subset:
+                    continue
 
-            if isinstance(source.source, FeatureSource):
-                features_subset.append(source.column)
-            else:
-                assert False, "Unknown tensor feature source"
+                if isinstance(source.source, FeatureSource):
+                    features_subset.append(source.column)
+                else:
+                    assert False, "Unknown tensor feature source"
 
         return set(features_subset)
 
@@ -404,7 +407,7 @@ class SequenceTokenizer:
 
     @classmethod
     @deprecation_warning("with `use_pickle` equals to `True` will be deprecated in future versions")
-    def load(cls, path: str, use_pickle: bool = False) -> "SequenceTokenizer":
+    def load(cls, path: str, use_pickle: bool = False, **kwargs) -> "SequenceTokenizer":
         """
         Load tokenizer object from the given path.
 
@@ -422,18 +425,7 @@ class SequenceTokenizer:
 
             # load tensor_schema, tensor_features
             tensor_schema_data = tokenizer_dict["init_args"]["tensor_schema"]
-            features_list = []
-            for feature_data in tensor_schema_data:
-                feature_data["feature_sources"] = [
-                    TensorFeatureSource(source=FeatureSource[x["source"]], column=x["column"], index=x["index"])
-                    for x in feature_data["feature_sources"]
-                ]
-                f_type = feature_data["feature_type"]
-                f_hint = feature_data["feature_hint"]
-                feature_data["feature_type"] = FeatureType[f_type] if f_type else None
-                feature_data["feature_hint"] = FeatureHint[f_hint] if f_hint else None
-                features_list.append(TensorFeatureInfo(**feature_data))
-            tokenizer_dict["init_args"]["tensor_schema"] = TensorSchema(features_list)
+            tokenizer_dict["init_args"]["tensor_schema"] = TensorSchema._create_object_by_args(tensor_schema_data)
 
             # Load encoder columns and rules
             types = list(FeatureHint) + list(FeatureSource)
@@ -447,7 +439,7 @@ class SequenceTokenizer:
                 rule_data = rules_dict[rule]
                 if rule_data["mapping"] and rule_data["is_int"]:
                     rule_data["mapping"] = {int(key): value for key, value in rule_data["mapping"].items()}
-                    del rule_data["is_int"]
+                del rule_data["is_int"]
 
                 tokenizer_dict["encoder"]["encoding_rules"][rule] = LabelEncodingRule(**rule_data)
 
@@ -478,31 +470,9 @@ class SequenceTokenizer:
                 "allow_collect_to_master": self._allow_collect_to_master,
                 "handle_unknown_rule": self._encoder._handle_unknown_rule,
                 "default_value_rule": self._encoder._default_value_rule,
-                "tensor_schema": [],
+                "tensor_schema": self._tensor_schema._get_object_args(),
             }
 
-            # save tensor schema
-            for feature in list(self._tensor_schema.values()):
-                tokenizer_dict["init_args"]["tensor_schema"].append(
-                    {
-                        "name": feature.name,
-                        "feature_type": feature.feature_type.name,
-                        "is_seq": feature.is_seq,
-                        "feature_hint": feature.feature_hint.name if feature.feature_hint else None,
-                        "feature_sources": [
-                            {"source": x.source.name, "column": x.column, "index": x.index}
-                            for x in feature.feature_sources
-                        ]
-                        if feature.feature_sources
-                        else None,
-                        "cardinality": feature.cardinality if feature.feature_type == FeatureType.CATEGORICAL else None,
-                        "embedding_dim": feature.embedding_dim
-                        if feature.feature_type == FeatureType.CATEGORICAL
-                        else None,
-                        "tensor_dim": feature.tensor_dim if feature.feature_type == FeatureType.NUMERICAL else None,
-                    }
-                )
-
             # save DatasetLabelEncoder
             tokenizer_dict["encoder"] = {
                 "features_columns": {key.name: value for key, value in self._encoder._features_columns.items()},

replay/data/nn/sequential_dataset.py

@@ -1,7 +1,10 @@
 import abc
+import json
+from pathlib import Path
 from typing import Tuple, Union
 
 import numpy as np
+import pandas as pd
 import polars as pl
 from pandas import DataFrame as PandasDataFrame
 from polars import DataFrame as PolarsDataFrame
@@ -100,6 +103,23 @@ class SequentialDataset(abc.ABC):
         rhs_filtered = rhs.filter_by_query_id(common_queries)
         return lhs_filtered, rhs_filtered
 
+    def save(self, path: str) -> None:
+        base_path = Path(path).with_suffix(".replay").resolve()
+        base_path.mkdir(parents=True, exist_ok=True)
+
+        sequential_dict = {}
+        sequential_dict["_class_name"] = self.__class__.__name__
+        self._sequences.reset_index().to_json(base_path / "sequences.json")
+        sequential_dict["init_args"] = {
+            "tensor_schema": self._tensor_schema._get_object_args(),
+            "query_id_column": self._query_id_column,
+            "item_id_column": self._item_id_column,
+            "sequences_path": "sequences.json",
+        }
+
+        with open(base_path / "init_args.json", "w+") as file:
+            json.dump(sequential_dict, file)
+
 
 class PandasSequentialDataset(SequentialDataset):
     """
@@ -174,6 +194,25 @@ class PandasSequentialDataset(SequentialDataset):
                 msg = "Tensor schema does not match with provided data frame"
                 raise ValueError(msg)
 
+    @classmethod
+    def load(cls, path: str, **kwargs) -> "PandasSequentialDataset":
+        """
+        Method for loading PandasSequentialDataset object from `.replay` directory.
+        """
+        base_path = Path(path).with_suffix(".replay").resolve()
+        with open(base_path / "init_args.json", "r") as file:
+            sequential_dict = json.loads(file.read())
+
+        sequences = pd.read_json(base_path / sequential_dict["init_args"]["sequences_path"])
+        dataset = cls(
+            tensor_schema=TensorSchema._create_object_by_args(sequential_dict["init_args"]["tensor_schema"]),
+            query_id_column=sequential_dict["init_args"]["query_id_column"],
+            item_id_column=sequential_dict["init_args"]["item_id_column"],
+            sequences=sequences,
+        )
+
+        return dataset
+
 
 class PolarsSequentialDataset(PandasSequentialDataset):
     """
@@ -199,7 +238,7 @@ class PolarsSequentialDataset(PandasSequentialDataset):
         self._query_id_column = query_id_column
         self._item_id_column = item_id_column
 
-        self._sequences = sequences.to_pandas()
+        self._sequences = self._convert_polars_to_pandas(sequences)
         if self._sequences.index.name != query_id_column:
             self._sequences = self._sequences.set_index(query_id_column)
 
@@ -211,12 +250,47 @@ class PolarsSequentialDataset(PandasSequentialDataset):
             tensor_schema=self._tensor_schema,
             query_id_column=self._query_id_column,
             item_id_column=self._item_id_column,
-            sequences=pl.from_pandas(filtered_sequences),
+            sequences=self._convert_pandas_to_polars(filtered_sequences),
         )
 
+    def _convert_polars_to_pandas(self, df: PolarsDataFrame) -> PandasDataFrame:
+        pandas_df = PandasDataFrame(df.to_dict(as_series=False))
+
+        for column in pandas_df.select_dtypes(include="object").columns:
+            if isinstance(pandas_df[column].iloc[0], list):
+                pandas_df[column] = pandas_df[column].apply(lambda x: np.array(x))
+
+        return pandas_df
+
+    def _convert_pandas_to_polars(self, df: PandasDataFrame) -> PolarsDataFrame:
+        for column in df.select_dtypes(include="object").columns:
+            if isinstance(df[column].iloc[0], np.ndarray):
+                df[column] = df[column].apply(lambda x: x.tolist())
+
+        return pl.from_dict(df.to_dict("list"))
+
     @classmethod
     def _check_if_schema_matches_data(cls, tensor_schema: TensorSchema, data: PolarsDataFrame) -> None:
         for tensor_feature_name in tensor_schema:
             if tensor_feature_name not in data:
                 msg = "Tensor schema does not match with provided data frame"
                 raise ValueError(msg)
+
+    @classmethod
+    def load(cls, path: str, **kwargs) -> "PandasSequentialDataset":
+        """
+        Method for loading PandasSequentialDataset object from `.replay` directory.
+        """
+        base_path = Path(path).with_suffix(".replay").resolve()
+        with open(base_path / "init_args.json", "r") as file:
+            sequential_dict = json.loads(file.read())
+
+        sequences = pl.DataFrame(pd.read_json(base_path / sequential_dict["init_args"]["sequences_path"]))
+        dataset = cls(
+            tensor_schema=TensorSchema._create_object_by_args(sequential_dict["init_args"]["tensor_schema"]),
+            query_id_column=sequential_dict["init_args"]["query_id_column"],
+            item_id_column=sequential_dict["init_args"]["item_id_column"],
+            sequences=sequences,
+        )
+
+        return dataset