_

Author: rcap107

skrub/_session_encoder.py

@@ -11,7 +11,11 @@
 The name of the session column is "{timestamp}_{suffix}", where "timestamp" is the name
 of the timestamp column, and "suffix" is a string that can be set via the "suffix"
 parameter (default is "session_id"). The session column contains a unique identifier for
-each session, which is a combination of the "split_by" column(s) and a session number
+each session, which is a combination of the "split_by" column(s) and a session number.
+
+Null values in "timestamp" or "split_by" columns are assigned a session ID of -1,
+and are ignored when computing time intervals.
+
 """
 
 import numbers
@@ -35,14 +39,14 @@
 
 
 @dispatch
-def _add_session_column(
+def _get_session_column(
     X, split_by_columns, timestamp_column, session_gap, session_id_column
 ):
     raise_dispatch_unregistered_type(X, kind="Dataframe")
 
 
-@_add_session_column.specialize("pandas")
-def _add_session_column_pandas(
+@_get_session_column.specialize("pandas")
+def _get_session_column_pandas(
     X, split_by_columns, timestamp_column, session_gap, session_id_column
 ):
     # Adding a row order column to sort lines back
@@ -102,8 +106,8 @@ def _add_session_column_pandas(
     return X_with_session_id.sort_values(by=row_order_col)[session_id_column]
 
 
-@_add_session_column.specialize("polars")
-def _add_session_column_polars(
+@_get_session_column.specialize("polars")
+def _get_session_column_polars(
     X, split_by_columns, timestamp_column, session_gap, session_id_column
 ):
     # Adding a row order column to sort lines back
@@ -566,7 +570,7 @@ def transform(self, X, y=None):
                 X, **{self.session_id_column_: np.array([], dtype=np.int64)}
             )
 
-        session_id = _add_session_column(
+        session_id = _get_session_column(
             X,
             self._split_by_columns,
             self.timestamp_col,

skrub/tests/test_session_encoder.py

@@ -6,7 +6,7 @@
 from packaging.version import parse
 
 from .. import _dataframe as sbd
-from .._session_encoder import SessionEncoder, _add_session_column
+from .._session_encoder import SessionEncoder, _get_session_column
 
 
 @pytest.fixture
@@ -542,7 +542,7 @@ def test_get_feature_names(df_module):
 
 
 @pytest.mark.skipif(parse(pd.__version__).major >= 3, reason="Test only for pandas < 3")
-def test_add_session_column_old_pandas(df_module):
+def test_get_session_column_old_pandas(df_module):
     """Old versions of pandas have a different branch that needs to be covered"""
     df = df_module.make_dataframe(
         {
@@ -555,7 +555,7 @@ def test_add_session_column_old_pandas(df_module):
         }
     )
     session_id = sbd.to_list(
-        _add_session_column(
+        _get_session_column(
             df, [], "timestamp", 30 * 60, session_id_column="timestamp_session_id"
         ),
     )
@@ -592,7 +592,7 @@ def test_proper_suffix(timestamp, suffix, df_module):
         result = SessionEncoder(
             timestamp_col=timestamp, split_by="user_id", suffix=suffix
         ).fit_transform(df)
-        # _add_session_column now returns the full dataframe with session_id added
+        # SessionEncoder now returns the full dataframe with session_id added
         expected_name = f"{timestamp}_{suffix}"
         assert expected_name in sbd.column_names(result)
 
@@ -625,7 +625,7 @@ def test_preserves_input_order(df_module):
 
 def test_error_dispatch():
     with pytest.raises(TypeError, match="Expecting a Pandas or Polars Dataframe"):
-        _add_session_column(
+        _get_session_column(
             np.array([1]),
             split_by_columns=[],
             timestamp_column="timestamp",