Merge branch 'feature/21.4_update' into 'release/v0.21.4'

0.21.4 features & fixes

See merge request ai-lab-pmo/mltools/recsys/RePlay!348

Author: OnlyDeniko

projects/pyproject.toml.template

@@ -127,7 +127,7 @@ priority = "explicit"
 
 [tool.poetry-dynamic-versioning]
 enable = true
-format-jinja = """0.21.3{{ env['PACKAGE_SUFFIX'] }}"""
+format-jinja = """0.21.4{{ env['PACKAGE_SUFFIX'] }}"""
 vcs = "git"
 
 [tool.ruff]

replay/data/nn/parquet/parquet_dataset.py

@@ -48,8 +48,9 @@ class ParquetDataset(IterableDataset):
 
     *   ``ParquetDataset`` supports only numeric values (boolean/integer/float),
         therefore, the data paths passed as arguments must contain encoded data.
-    *   For optimal performance, set the ``OMP_NUM_THREADS`` and ``ARROW_IO_THREADS`` to match
-        the number of available CPU cores.
+    *   For achieving maximum performance, set the ARROW_NUM_THREADS and ARROW_IO_THREADS to match
+        the number of available CPU cores and set the OMP_NUM_THREADS and MKL_NUM_THREADS to 1, also set
+        *torch.set_num_threads(1)* and *torch.set_num_interop_threads(1)*.
 
     """
 

replay/data/nn/parquet/parquet_module.py

@@ -30,8 +30,9 @@ class ParquetModule(L.LightningDataModule):
 
     *   ``ParquetModule`` supports only numeric values (boolean/integer/float),
         therefore, the data paths passed as arguments must contain encoded data.
-    *   For optimal performance, set the OMP_NUM_THREADS and ARROW_IO_THREADS to match
-        the number of available CPU cores.
+    *   For achieving maximum performance, set the ARROW_NUM_THREADS and ARROW_IO_THREADS to match
+        the number of available CPU cores and set the OMP_NUM_THREADS and MKL_NUM_THREADS to 1, also set
+        *torch.set_num_threads(1)* and *torch.set_num_interop_threads(1)*.
     *   It's possible to use all train/validate/test/predict splits, then paths to splits should be passed
         as corresponding arguments of ``ParquetModule``.
         Alternatively, all the paths to the splits may be not specified

replay/nn/lightning/callback/metrics_callback.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Optional
 
 import lightning
 import torch
@@ -62,29 +62,22 @@ def __init__(
         self._ground_truth_column = ground_truth_column
         self._train_column = train_column
 
-    def _get_dataloaders_size(self, dataloaders: Optional[Any]) -> list[int]:
-        if isinstance(dataloaders, list):
-            return [len(dataloader) for dataloader in dataloaders]
-        return [len(dataloaders)]
-
     def on_validation_epoch_start(
         self,
         trainer: lightning.Trainer,
         pl_module: LightningModule,  # noqa: ARG002
     ) -> None:
-        self._dataloaders_size = self._get_dataloaders_size(trainer.val_dataloaders)
-        self._metrics_builders = [
-            TorchMetricsBuilder(self._metrics, self._ks, self._item_count) for _ in self._dataloaders_size
-        ]
-        for builder in self._metrics_builders:
-            builder.reset()
+        self._epoch_start(dataloaders_size=trainer.num_val_batches)
 
     def on_test_epoch_start(
         self,
         trainer: lightning.Trainer,
         pl_module: LightningModule,  # noqa: ARG002
     ) -> None:
-        self._dataloaders_size = self._get_dataloaders_size(trainer.test_dataloaders)
+        self._epoch_start(dataloaders_size=trainer.num_test_batches)
+
+    def _epoch_start(self, dataloaders_size: list[int]) -> None:
+        self._dataloaders_size = dataloaders_size
         self._metrics_builders = [
             TorchMetricsBuilder(self._metrics, self._ks, self._item_count) for _ in self._dataloaders_size
         ]
@@ -146,7 +139,6 @@ def _batch_end(
         self._metrics_builders[dataloader_idx].add_prediction(
             sampled_items, batch[self._ground_truth_column], batch.get(self._train_column)
         )
-
         if batch_idx + 1 == self._dataloaders_size[dataloader_idx]:
             pl_module.log_dict(
                 self._metrics_builders[dataloader_idx].get_metrics(),

replay/preprocessing/label_encoder.py

@@ -630,13 +630,22 @@ def partial_fit(self, df: DataFrameLike) -> "SequenceEncodingRule":
         return self
 
     def _transform_spark(self, df: SparkDataFrame, default_value: Optional[int]) -> SparkDataFrame:
+        fake_column_name = "__fake_column__"
+        df = df.withColumn(fake_column_name, sf.rand(0).cast("string"))
+
         other_columns = [col for col in df.columns if col != self._col]
 
         mapping_on_spark = get_spark_session().createDataFrame(
             data=list(self.get_mapping().items()), schema=[self._col, self._target_col]
         )
+
+        array_size_column_name = "__array_size__"
+        df_with_sz = df.withColumn(array_size_column_name, sf.size(self._col))
+        empty_arrays = df_with_sz.filter(sf.col(array_size_column_name) == 0).drop(array_size_column_name)
+        non_empty_arrays = df_with_sz.filter(sf.col(array_size_column_name) != 0).drop(array_size_column_name)
+
         encoded_df = (
-            df.select(*other_columns, sf.posexplode(self._col))
+            non_empty_arrays.select(*other_columns, sf.posexplode(self._col))
             .withColumnRenamed("col", self._col)
             .join(mapping_on_spark, on=self._col, how="left")
         )
@@ -663,7 +672,9 @@ def _transform_spark(self, df: SparkDataFrame, default_value: Optional[int]) ->
                     LabelEncoderTransformWarning,
                 )
 
-        return result
+        # for correct concatenation, it is necessary that the column schemes match
+        empty_arrays = empty_arrays.withColumn(self._col, sf.col(self._col).cast(result.schema[self._col].dataType))
+        return result.unionByName(empty_arrays).drop(fake_column_name)
 
     def _transform_pandas(self, df: PandasDataFrame, default_value: Optional[int]) -> PandasDataFrame:
         mapping = self.get_mapping()

tests/experimental/preprocessing/test_sequence_generator.py

@@ -10,12 +10,17 @@
     columns,
     columns_target,
     columns_target_list_len,
+    data_target,
+    data_target_ordered,
+    data_target_ordered_list_len,
     schema_target,
     schema_target_list_len,
     simple_dataframe,
     simple_dataframe_additional,
+    simple_dataframe_additional_data,
     simple_dataframe_additional_pandas,
     simple_dataframe_array,
+    simple_dataframe_array_data,
     simple_dataframe_array_pandas,
     simple_dataframe_pandas,
     simple_dataframe_target,

tests/preprocessing/conftest.py

@@ -785,31 +785,48 @@ def generate_random_string(length=10):
 
 
 @pytest.fixture(scope="module")
-def simple_dataframe_array(spark):
-    columns_array = ["user_id", "item_id", "timestamp"]
-    data_array = [
-        (1, [2, 1, 0], 19842),
+def simple_dataframe_array_data():
+    columns = ["user_id", "item_id", "timestamp"]
+    data = [
+        (1, [], 19842),
         (1, [4, 1], 19844),
         (1, [3, 1, 0], 19843),
         (1, [5, 1], 19845),
-        (1, [6, 1, 0], 19846),
+        (1, [], 19846),
         (1, [7, 1], 19847),
         (2, [1, 0, 1], 19841),
-        (2, [2, 0], 19842),
+        (2, [], 19842),
         (2, [3, 0, 1], 19843),
         (2, [4, 0], 19844),
         (3, [10, 0], 19844),
         (4, [11, 0, 1], 19843),
         (4, [12, 0], 19845),
-        (1, [1, 0], 19841),
+        (1, [], 19841),
     ]
-    return spark.createDataFrame(data_array, schema=columns_array)
+    return data, columns
+
+
+@pytest.fixture(scope="module")
+def simple_dataframe_array(simple_dataframe_array_data, spark):
+    data, columns = simple_dataframe_array_data
+    return spark.createDataFrame(data, schema=columns)
+
+
+@pytest.fixture(scope="module")
+def simple_dataframe_array_pandas(simple_dataframe_array_data):
+    data, columns = simple_dataframe_array_data
+    return pd.DataFrame(data, columns=columns)
+
+
+@pytest.fixture(scope="module")
+def simple_dataframe_array_polars(simple_dataframe_array_pandas):
+    return pl.from_pandas(simple_dataframe_array_pandas)
 
 
 @pytest.fixture(scope="module")
-def simple_dataframe_additional(spark):
-    columns_additional = ["user_id", "item_id", "timestamp", "other_column"]
-    data_additional = [
+def simple_dataframe_additional_data():
+    columns = ["user_id", "item_id", "timestamp", "other_column"]
+    data = [
         (1, 2, 19842, 0),
         (1, 4, 19844, 1),
         (1, 3, 19843, 0),
@@ -825,12 +842,24 @@ def simple_dataframe_additional(spark):
         (4, 12, 19845, 1),
         (1, 1, 19841, 1),
     ]
-    return spark.createDataFrame(data_additional, schema=columns_additional)
+    return data, columns
 
 
 @pytest.fixture(scope="module")
-def simple_dataframe_target(spark, schema_target):
-    data_target = [
+def simple_dataframe_additional(simple_dataframe_additional_data, spark):
+    data, columns = simple_dataframe_additional_data
+    return spark.createDataFrame(data, schema=columns)
+
+
+@pytest.fixture(scope="module")
+def simple_dataframe_additional_pandas(simple_dataframe_additional_data):
+    data, columns = simple_dataframe_additional_data
+    return pd.DataFrame(data, columns=columns)
+
+
+@pytest.fixture(scope="module")
+def data_target():
+    return [
         (1, 4, 19844, [2], [19842]),
         (1, 3, 19843, [2, 4], [19842, 19844]),
         (1, 5, 19845, [2, 4, 3], [19842, 19844, 19843]),
@@ -842,12 +871,21 @@ def simple_dataframe_target(spark, schema_target):
         (2, 4, 19844, [1, 2, 3], [19841, 19842, 19843]),
         (4, 12, 19845, [11], [19843]),
     ]
+
+
+@pytest.fixture(scope="module")
+def simple_dataframe_target(spark, data_target, schema_target):
     return spark.createDataFrame(data_target, schema=schema_target)
 
 
 @pytest.fixture(scope="module")
-def simple_dataframe_target_ordered(spark, schema_target):
-    data_target_ordered = [
+def simple_dataframe_target_pandas(data_target, columns_target):
+    return pd.DataFrame(data_target, columns=columns_target)
+
+
+@pytest.fixture(scope="module")
+def data_target_ordered():
+    return [
         (1, 2, 19842, [1], [19841]),
         (1, 3, 19843, [1, 2], [19841, 19842]),
         (1, 4, 19844, [1, 2, 3], [19841, 19842, 19843]),
@@ -859,12 +897,21 @@ def simple_dataframe_target_ordered(spark, schema_target):
         (2, 4, 19844, [1, 2, 3], [19841, 19842, 19843]),
         (4, 12, 19845, [11], [19843]),
     ]
+
+
+@pytest.fixture(scope="module")
+def simple_dataframe_target_ordered(spark, data_target_ordered, schema_target):
     return spark.createDataFrame(data_target_ordered, schema=schema_target)
 
 
 @pytest.fixture(scope="module")
-def simple_dataframe_target_ordered_list_len(spark, schema_target_list_len):
-    data_target_ordered_list_len = [
+def simple_dataframe_target_ordered_pandas(data_target_ordered, columns_target):
+    return pd.DataFrame(data_target_ordered, columns=columns_target)
+
+
+@pytest.fixture(scope="module")
+def data_target_ordered_list_len():
+    return [
         (1, 2, 19842, [1], [19841], 1),
         (1, 3, 19843, [1, 2], [19841, 19842], 2),
         (1, 4, 19844, [1, 2, 3], [19841, 19842, 19843], 3),
@@ -876,9 +923,18 @@ def simple_dataframe_target_ordered_list_len(spark, schema_target_list_len):
         (2, 4, 19844, [1, 2, 3], [19841, 19842, 19843], 3),
         (4, 12, 19845, [11], [19843], 1),
     ]
+
+
+@pytest.fixture(scope="module")
+def simple_dataframe_target_ordered_list_len(spark, data_target_ordered_list_len, schema_target_list_len):
     return spark.createDataFrame(data_target_ordered_list_len, schema=schema_target_list_len)
 
 
+@pytest.fixture(scope="module")
+def simple_dataframe_target_ordered_list_len_pandas(data_target_ordered_list_len, columns_target_list_len):
+    return pd.DataFrame(data_target_ordered_list_len, columns=columns_target_list_len)
+
+
 @pytest.fixture(scope="module")
 def simple_dataframe_pandas(columns):
     data = [
@@ -910,106 +966,6 @@ def dataframe_not_implemented(simple_dataframe_pandas):
     return simple_dataframe_pandas.to_numpy()
 
 
-@pytest.fixture(scope="module")
-def simple_dataframe_array_pandas():
-    columns_array = ["user_id", "item_id", "timestamp"]
-    data_array = [
-        (1, [2, 1, 0], 19842),
-        (1, [4, 1], 19844),
-        (1, [3, 1, 0], 19843),
-        (1, [5, 1], 19845),
-        (1, [6, 1, 0], 19846),
-        (1, [7, 1], 19847),
-        (2, [1, 0, 1], 19841),
-        (2, [2, 0], 19842),
-        (2, [3, 0, 1], 19843),
-        (2, [4, 0], 19844),
-        (3, [10, 0], 19844),
-        (4, [11, 0, 1], 19843),
-        (4, [12, 0], 19845),
-        (1, [1, 0], 19841),
-    ]
-    return pd.DataFrame(data_array, columns=columns_array)
-
-
-@pytest.fixture(scope="module")
-def simple_dataframe_array_polars(simple_dataframe_array_pandas):
-    return pl.from_pandas(simple_dataframe_array_pandas)
-
-
-@pytest.fixture(scope="module")
-def simple_dataframe_additional_pandas():
-    columns_additional = ["user_id", "item_id", "timestamp", "other_column"]
-    data_additional = [
-        (1, 2, 19842, 0),
-        (1, 4, 19844, 1),
-        (1, 3, 19843, 0),
-        (1, 5, 19845, 1),
-        (1, 6, 19846, 0),
-        (1, 7, 19847, 1),
-        (2, 1, 19841, 0),
-        (2, 2, 19842, 0),
-        (2, 3, 19843, 0),
-        (2, 4, 19844, 1),
-        (3, 10, 19844, 0),
-        (4, 11, 19843, 1),
-        (4, 12, 19845, 1),
-        (1, 1, 19841, 1),
-    ]
-    return pd.DataFrame(data_additional, columns=columns_additional)
-
-
-@pytest.fixture(scope="module")
-def simple_dataframe_target_pandas(columns_target):
-    data_target = [
-        (1, 4, 19844, [2], [19842]),
-        (1, 3, 19843, [2, 4], [19842, 19844]),
-        (1, 5, 19845, [2, 4, 3], [19842, 19844, 19843]),
-        (1, 6, 19846, [2, 4, 3, 5], [19842, 19844, 19843, 19845]),
-        (1, 7, 19847, [2, 4, 3, 5, 6], [19842, 19844, 19843, 19845, 19846]),
-        (1, 1, 19841, [4, 3, 5, 6, 7], [19844, 19843, 19845, 19846, 19847]),
-        (2, 2, 19842, [1], [19841]),
-        (2, 3, 19843, [1, 2], [19841, 19842]),
-        (2, 4, 19844, [1, 2, 3], [19841, 19842, 19843]),
-        (4, 12, 19845, [11], [19843]),
-    ]
-    return pd.DataFrame(data_target, columns=columns_target)
-
-
-@pytest.fixture(scope="module")
-def simple_dataframe_target_ordered_pandas(columns_target):
-    data_target_ordered = [
-        (1, 2, 19842, [1], [19841]),
-        (1, 3, 19843, [1, 2], [19841, 19842]),
-        (1, 4, 19844, [1, 2, 3], [19841, 19842, 19843]),
-        (1, 5, 19845, [1, 2, 3, 4], [19841, 19842, 19843, 19844]),
-        (1, 6, 19846, [1, 2, 3, 4, 5], [19841, 19842, 19843, 19844, 19845]),
-        (1, 7, 19847, [2, 3, 4, 5, 6], [19842, 19843, 19844, 19845, 19846]),
-        (2, 2, 19842, [1], [19841]),
-        (2, 3, 19843, [1, 2], [19841, 19842]),
-        (2, 4, 19844, [1, 2, 3], [19841, 19842, 19843]),
-        (4, 12, 19845, [11], [19843]),
-    ]
-    return pd.DataFrame(data_target_ordered, columns=columns_target)
-
-
-@pytest.fixture(scope="module")
-def simple_dataframe_target_ordered_list_len_pandas(columns_target_list_len):
-    data_target_ordered_list_len = [
-        (1, 2, 19842, [1], [19841], 1),
-        (1, 3, 19843, [1, 2], [19841, 19842], 2),
-        (1, 4, 19844, [1, 2, 3], [19841, 19842, 19843], 3),
-        (1, 5, 19845, [1, 2, 3, 4], [19841, 19842, 19843, 19844], 4),
-        (1, 6, 19846, [1, 2, 3, 4, 5], [19841, 19842, 19843, 19844, 19845], 5),
-        (1, 7, 19847, [2, 3, 4, 5, 6], [19842, 19843, 19844, 19845, 19846], 5),
-        (2, 2, 19842, [1], [19841], 1),
-        (2, 3, 19843, [1, 2], [19841, 19842], 2),
-        (2, 4, 19844, [1, 2, 3], [19841, 19842, 19843], 3),
-        (4, 12, 19845, [11], [19843], 1),
-    ]
-    return pd.DataFrame(data_target_ordered_list_len, columns=columns_target_list_len)
-
-
 @pytest.fixture(scope="module")
 def dataframe_sessionizer(spark):
     columns = ["user_id", "item_id", "timestamp"]