keep_extra_cols argument for dataset and interactions

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Optional `epochs` argument to `ImplicitALSWrapperModel.fit` method ([#203](https://github.com/MobileTeleSystems/RecTools/pull/203))
 - `save` and `load` methods to all of the models ([#206](https://github.com/MobileTeleSystems/RecTools/pull/206))
 - Model configs example ([#207](https://github.com/MobileTeleSystems/RecTools/pull/207))
+- `keep_extra_cols` argument to `Dataset.construct` and `Interactions.from_raw` methods ([#208](https://github.com/MobileTeleSystems/RecTools/pull/208))
 
 
 ## [0.8.0] - 28.08.2024

rectools/dataset/dataset.py

@@ -102,6 +102,7 @@ def construct(
         item_features_df: tp.Optional[pd.DataFrame] = None,
         cat_item_features: tp.Iterable[str] = (),
         make_dense_item_features: bool = False,
+        keep_extra_cols: bool = False,
     ) -> "Dataset":
         """Class method for convenient `Dataset` creation.
 
@@ -133,6 +134,8 @@ def construct(
             Used only if `user_features_df` (`item_features_df`) is not ``None``.
             - if ``False``, `SparseFeatures.from_flatten` method will be used;
             - if ``True``,  `DenseFeatures.from_dataframe` method will be used.
+        keep_extra_cols: bool, default ``False``
+            Flag to keep all columns from interactions besides the default ones.
 
         Returns
         -------
@@ -144,7 +147,7 @@ def construct(
                 raise KeyError(f"Column '{col}' must be present in `interactions_df`")
         user_id_map = IdMap.from_values(interactions_df[Columns.User].values)
         item_id_map = IdMap.from_values(interactions_df[Columns.Item].values)
-        interactions = Interactions.from_raw(interactions_df, user_id_map, item_id_map)
+        interactions = Interactions.from_raw(interactions_df, user_id_map, item_id_map, keep_extra_cols)
 
         user_features, user_id_map = cls._make_features(
             user_features_df,

rectools/dataset/interactions.py

@@ -42,6 +42,7 @@ class Interactions:
                 - `Columns.Weight` - weight of interaction, float, use ``1`` if interactions have no weight;
                 - `Columns.Datetime` - timestamp of interactions,
                   assign random value if you're not going to use it later.
+            Extra columns can also be present.
     """
 
     df: pd.DataFrame = attr.ib()
@@ -81,12 +82,15 @@ def __attrs_post_init__(self) -> None:
         """Convert datetime and weight columns to the right data types."""
         self._convert_weight_and_datetime_types(self.df)
 
+    @staticmethod
+    def _add_extra_cols(df: pd.DataFrame, interactions: pd.DataFrame) -> None:
+        extra_cols = [col for col in interactions.columns if col not in df.columns]
+        for extra_col in extra_cols:
+            df[extra_col] = interactions[extra_col].values
+
     @classmethod
     def from_raw(
-        cls,
-        interactions: pd.DataFrame,
-        user_id_map: IdMap,
-        item_id_map: IdMap,
+        cls, interactions: pd.DataFrame, user_id_map: IdMap, item_id_map: IdMap, keep_extra_cols: bool = False
     ) -> "Interactions":
         """
         Create `Interactions` from dataset with external ids and id mappings.
@@ -104,6 +108,8 @@ def from_raw(
             User identifiers mapping.
         item_id_map : IdMap
             Item identifiers mapping.
+        keep_extra_cols: bool, default ``False``
+            Flag to keep all columns from interactions besides the default ones.
 
         Returns
         -------
@@ -120,6 +126,8 @@ def from_raw(
         df[Columns.Weight] = interactions[Columns.Weight].values
         df[Columns.Datetime] = interactions[Columns.Datetime].values
         cls._convert_weight_and_datetime_types(df)
+        if keep_extra_cols:
+            cls._add_extra_cols(df, interactions)
 
         return cls(df)
 
@@ -159,6 +167,7 @@ def to_external(
         item_id_map: IdMap,
         include_weight: bool = True,
         include_datetime: bool = True,
+        include_extra_cols: bool = True,
     ) -> pd.DataFrame:
         """
         Convert itself to `pd.DataFrame` with replacing internal user and item ids to external ones.
@@ -173,6 +182,8 @@ def to_external(
             Whether to include weight column into resulting table or not
         include_datetime : bool, default ``True``
             Whether to include datetime column into resulting table or not.
+        include_extra_cols: bool, default ``True``
+            Whether to include extra columns into resulting table or not.
 
         Returns
         -------
@@ -184,10 +195,17 @@ def to_external(
                 Columns.Item: item_id_map.convert_to_external(self.df[Columns.Item].values),
             }
         )
+        cols_to_add = []
 
         if include_weight:
-            res[Columns.Weight] = self.df[Columns.Weight]
+            cols_to_add.append(Columns.Weight)
         if include_datetime:
-            res[Columns.Datetime] = self.df[Columns.Datetime]
-
+            cols_to_add.append(Columns.Datetime)
+        if include_extra_cols:
+            cols_not_to_add = [Columns.User, Columns.Item, Columns.Weight, Columns.Datetime]
+            extra_cols = [col for col in self.df if col not in cols_not_to_add]
+            cols_to_add.extend(extra_cols)
+
+        for col in cols_to_add:
+            res[col] = self.df[col]
         return res

tests/dataset/test_dataset.py

@@ -36,14 +36,14 @@ class TestDataset:
     def setup_method(self) -> None:
         self.interactions_df = pd.DataFrame(
             [
-                ["u1", "i1", 2, "2021-09-09"],
-                ["u1", "i2", 2, "2021-09-05"],
-                ["u1", "i1", 6, "2021-08-09"],
-                ["u2", "i1", 7, "2020-09-09"],
-                ["u2", "i5", 9, "2021-09-03"],
-                ["u3", "i1", 2, "2021-09-09"],
+                ["u1", "i1", 2, "2021-09-09", 5],
+                ["u1", "i2", 2, "2021-09-05", 6],
+                ["u1", "i1", 6, "2021-08-09", 7],
+                ["u2", "i1", 7, "2020-09-09", 8],
+                ["u2", "i5", 9, "2021-09-03", 9],
+                ["u3", "i1", 2, "2021-09-09", 10],
             ],
-            columns=[Columns.User, Columns.Item, Columns.Weight, Columns.Datetime],
+            columns=[Columns.User, Columns.Item, Columns.Weight, Columns.Datetime, "extra_col"],
         )
         self.expected_user_id_map = IdMap.from_values(["u1", "u2", "u3"])
         self.expected_item_id_map = IdMap.from_values(["i1", "i2", "i5"])
@@ -78,6 +78,14 @@ def assert_dataset_equal_to_expected(
         assert_feature_set_equal(actual.user_features, expected_user_features)
         assert_feature_set_equal(actual.item_features, expected_item_features)
 
+    def test_construct_with_extra_cols(self) -> None:
+
+        dataset = Dataset.construct(self.interactions_df, keep_extra_cols=True)
+        actual = dataset.interactions
+        expected = self.expected_interactions
+        expected.df["extra_col"] = self.interactions_df["extra_col"]
+        assert_interactions_set_equal(actual, expected)
+
     def test_construct_without_features(self) -> None:
         dataset = Dataset.construct(self.interactions_df)
         self.assert_dataset_equal_to_expected(dataset, None, None)
@@ -276,14 +284,17 @@ def test_raises_when_in_dense_features_absent_some_ids_that_present_in_interacti
 
     @pytest.mark.parametrize("include_weight", (True, False))
     @pytest.mark.parametrize("include_datetime", (True, False))
-    def test_get_raw_interactions(self, include_weight: bool, include_datetime: bool) -> None:
-        dataset = Dataset.construct(self.interactions_df)
+    @pytest.mark.parametrize("keep_extra_cols", (True, False))
+    def test_get_raw_interactions(self, include_weight: bool, include_datetime: bool, keep_extra_cols: bool) -> None:
+        dataset = Dataset.construct(self.interactions_df, keep_extra_cols=keep_extra_cols)
         actual = dataset.get_raw_interactions(include_weight, include_datetime)
         expected = self.interactions_df.astype({Columns.Weight: "float64", Columns.Datetime: "datetime64[ns]"})
         if not include_weight:
             expected.drop(columns=Columns.Weight, inplace=True)
         if not include_datetime:
             expected.drop(columns=Columns.Datetime, inplace=True)
+        if not keep_extra_cols:
+            expected.drop(columns="extra_col", inplace=True)
         pd.testing.assert_frame_equal(actual, expected)
 
     @pytest.fixture

tests/dataset/test_interactions.py

@@ -36,6 +36,16 @@ def setup_method(self) -> None:
                 Columns.Item: [0, 1, 0, 1],
                 Columns.Weight: [5, 7.0, 4, 1],
                 Columns.Datetime: [datetime(2021, 9, 8)] * 4,
+                "extra_col": [1, 2, 3, 4],
+            }
+        )
+        self.raw_df = pd.DataFrame(
+            {
+                Columns.User: ["u1", "u2", "u1", "u1"],
+                Columns.Item: ["i1", "i2", "i1", "i2"],
+                Columns.Weight: [5, 7, 4, 1],
+                Columns.Datetime: ["2021-09-08"] * 4,
+                "extra_col": [1, 2, 3, 4],
             }
         )
 
@@ -46,8 +56,9 @@ def test_creation(self) -> None:
     def test_missing_columns_validation(self, subtests: SubTests) -> None:
         for col in self.df.columns:
             with subtests.test(f"drop {col} column"):
-                with pytest.raises(KeyError):
-                    Interactions(self.df.drop(columns=col))
+                if col != "extra_col":
+                    with pytest.raises(KeyError):
+                        Interactions(self.df.drop(columns=col))
 
     @pytest.mark.parametrize("column", (Columns.User, Columns.Item))
     def test_types_validation(self, column: str) -> None:
@@ -60,19 +71,16 @@ def test_positivity_validation(self, column: str) -> None:
             self.df.at[0, column] = -1
             Interactions(self.df)
 
-    def test_from_raw_creation(self) -> None:
-        raw_df = pd.DataFrame(
-            {
-                Columns.User: ["u1", "u2", "u1", "u1"],
-                Columns.Item: ["i1", "i2", "i1", "i2"],
-                Columns.Weight: [5, 7, 4, 1],
-                Columns.Datetime: ["2021-09-08"] * 4,
-            }
-        )
+    @pytest.mark.parametrize("keep_extra_cols", (True, False))
+    def test_from_raw_creation(self, keep_extra_cols: bool) -> None:
+        raw_df = self.raw_df
         user_id_map = IdMap(np.array(["u0", "u1", "u2"]))
         item_id_map = IdMap.from_values(["i1", "i2"])
-        interactions = Interactions.from_raw(raw_df, user_id_map, item_id_map)
-        pd.testing.assert_frame_equal(interactions.df, self.df)
+        interactions = Interactions.from_raw(raw_df, user_id_map, item_id_map, keep_extra_cols=keep_extra_cols)
+        excepted = self.df
+        if not keep_extra_cols:
+            excepted.drop(columns="extra_col", inplace=True)
+        pd.testing.assert_frame_equal(interactions.df, excepted)
 
     @pytest.mark.parametrize(
         "with_weights,expected_data",
@@ -105,12 +113,15 @@ def test_raises_when_datetime_type_incorrect(self) -> None:
 
     @pytest.mark.parametrize("include_weight", (True, False))
     @pytest.mark.parametrize("include_datetime", (True, False))
-    def test_to_external(self, include_weight: bool, include_datetime: bool) -> None:
+    @pytest.mark.parametrize("include_extra_cols", (True, False))
+    def test_to_external(self, include_weight: bool, include_datetime: bool, include_extra_cols: bool) -> None:
         user_id_map = IdMap(np.array([10, 20, 30]))
         item_id_map = IdMap(np.array(["i1", "i2"]))
         interactions = Interactions(self.df)
 
-        actual = interactions.to_external(user_id_map, item_id_map, include_weight, include_datetime)
+        actual = interactions.to_external(
+            user_id_map, item_id_map, include_weight, include_datetime, include_extra_cols
+        )
         expected = pd.DataFrame(
             [
                 [20, "i1"],
@@ -124,6 +135,8 @@ def test_to_external(self, include_weight: bool, include_datetime: bool) -> None
             expected[Columns.Weight] = self.df[Columns.Weight]
         if include_datetime:
             expected[Columns.Datetime] = self.df[Columns.Datetime]
+        if include_extra_cols:
+            expected["extra_col"] = self.df["extra_col"]
 
         pd.testing.assert_frame_equal(actual, expected)
 
@@ -132,7 +145,7 @@ def test_to_external_empty(self) -> None:
         item_id_map = IdMap(np.array(["i1", "i2"]))
         interactions = Interactions(self.df.iloc[:0])
 
-        actual = interactions.to_external(user_id_map, item_id_map)
+        actual = interactions.to_external(user_id_map, item_id_map, include_extra_cols=False)
         expected = pd.DataFrame(
             [],
             columns=Columns.Interactions,