From 9f65344efa916de4df2155e6e6ff78e661d3edec Mon Sep 17 00:00:00 2001 From: Jingang Qu Date: Fri, 25 Sep 2026 04:17:19 -0700 Subject: [PATCH 1/2] Add `low_cardinality` to `infer_stypes` Type numerical columns with two or three distinct values as categorical on request, and use it in the Kumo TabArena adapter. Signed-off-by: Jingang Qu --- benchmark/tabular/model.py | 28 ++++++++++++++++- sdm/stype.py | 63 ++++++++++++++++++++++++++++++++++++-- test/test_stype.py | 34 ++++++++++++++++++++ 3 files changed, 122 insertions(+), 3 deletions(-) diff --git a/benchmark/tabular/model.py b/benchmark/tabular/model.py index 5a5caead5..2e02e4f94 100644 --- a/benchmark/tabular/model.py +++ b/benchmark/tabular/model.py @@ -88,7 +88,7 @@ def _fit( self.random_seed ) - X = self.preprocess(X, y=y) + X = self.preprocess(X, y=y, is_train=True) self.stypes = sdm.infer_stypes(X) x_context = sdm.TableTensor.from_pandas( df=X, @@ -349,6 +349,32 @@ def __setstate__(self, state: dict[str, Any]) -> None: device=self._device, ) + def _preprocess( + self, + X: pd.DataFrame, + is_train: bool = False, + **kwargs: Any, + ) -> pd.DataFrame: + X = super()._preprocess(X, **kwargs) + # AutoGluon's feature generator hands binary columns over as integers, + # so low-cardinality numerical columns are typed categorical. The + # recipe runs without AlignCategories (see `_create_recipe`), so their + # codes would follow the order of appearance: pin the categories seen + # in training in value order instead; other values become missing. + if is_train: + stypes = sdm.infer_stypes(X, low_cardinality="infer") + self._low_cardinality_dtypes = { + column: pd.CategoricalDtype( + categories=np.sort(X[column].dropna().unique()) + ) + for column, stype in stypes.items() + if stype == sdm.Stype.categorical + and X[column].dtype.kind in "iuf" + } + if self._low_cardinality_dtypes: + X = X.astype(self._low_cardinality_dtypes, copy=False) + return X + def _create_recipe(self) -> sdm.Recipe: recipe = super()._create_recipe() # TabArena aligns features with its fitted generator and targets with diff --git a/sdm/stype.py b/sdm/stype.py index c9d2aabd2..eb67962f2 100644 --- a/sdm/stype.py +++ b/sdm/stype.py @@ -58,6 +58,7 @@ def infer_stypes( *, text: Literal["off", "infer", "drop"] = "off", id: Literal["off", "infer", "drop"] = "off", + low_cardinality: Literal["off", "infer"] = "off", unsupported: Literal["error", "warn", "drop"] = "error", ) -> dict[str, StypeLike]: r"""Infer semantic types from raw data statistics. @@ -80,6 +81,10 @@ def infer_stypes( :attr:`~Stype.id` if its name contains ``"id"`` as a whole word (*e.g.*, ``"user_id"``, ``"userId"``, ``"id"``, but not ``"solid"`` or ``"covid"``). + * Integer, floating-point, and decimal columns are inferred as + :attr:`~Stype.categorical` if the table has more than 150 rows and the + column contains two or three distinct values, counting missing values as + one. Args: table: A :class:`pandas.DataFrame`, :class:`pyarrow.Table`, or @@ -93,6 +98,10 @@ def infer_stypes( ``"off"`` disables :attr:`~Stype.id` column detection. ``"infer"`` includes inferred :attr:`~Stype.id` columns. ``"drop"`` omits inferred :attr:`~Stype.id` columns. + low_cardinality: The detection policy for low-cardinality integer, + floating-point, and decimal columns. + ``"off"`` keeps them :attr:`~Stype.numerical`. + ``"infer"`` infers them as :attr:`~Stype.categorical`. unsupported: How to handle unsupported dtypes. ``"error"`` raises a :class:`TypeError`. ``"warn"`` emits a warning and omits the column. @@ -103,7 +112,13 @@ def infer_stypes( """ overrides = overrides or {} - fn: Callable[[str, object, Policy, Policy], Stype | None] | None = None + fn: ( + Callable[ + [str, object, Policy, Policy, Literal["off", "infer"]], + Stype | None, + ] + | None + ) = None columns: Iterable[tuple[Hashable, object]] | None = None if isinstance(table, pa.Table): fn = _infer_arrow_stype @@ -136,7 +151,7 @@ def infer_stypes( continue try: - stype = fn(name, column, text, id) + stype = fn(name, column, text, id, low_cardinality) except TypeError: if unsupported == "error": raise @@ -169,6 +184,7 @@ def _infer_arrow_stype( array: object, text: Policy, id: Policy, + low_cardinality: Literal["off", "infer"], ) -> Stype | None: assert isinstance(array, pa.Array | pa.ChunkedArray) dtype = array.type @@ -189,6 +205,8 @@ def _infer_arrow_stype( or pa.types.is_floating(dtype) or pa.types.is_decimal(dtype) ): + if low_cardinality != "off" and _is_arrow_low_cardinality(array): + return Stype.categorical return Stype.numerical if pa.types.is_boolean(dtype) or pa.types.is_dictionary(dtype): @@ -210,6 +228,7 @@ def _infer_pandas_stype( ser: object, text: Policy, id: Policy, + low_cardinality: Literal["off", "infer"], ) -> Stype | None: import pandas as pd from pandas.api.types import ( @@ -237,6 +256,8 @@ def _infer_pandas_stype( return None if id == "drop" else Stype.id if is_integer_dtype(dtype) or is_float_dtype(dtype): + if low_cardinality != "off" and _is_series_low_cardinality(ser): + return Stype.categorical return Stype.numerical if is_bool_dtype(dtype) or isinstance(dtype, pd.CategoricalDtype): @@ -258,6 +279,7 @@ def _infer_cudf_stype( ser: object, text: Policy, id: Policy, + low_cardinality: Literal["off", "infer"], ) -> Stype | None: import cudf from cudf.api.types import ( @@ -284,6 +306,8 @@ def _infer_cudf_stype( or is_float_dtype(dtype) or is_decimal_dtype(dtype) ): + if low_cardinality != "off" and _is_series_low_cardinality(ser): + return Stype.categorical return Stype.numerical if is_bool_dtype(dtype) or isinstance(dtype, cudf.CategoricalDtype): @@ -343,3 +367,38 @@ def _is_cudf_text(ser: cudf.Series) -> bool: unique = ser.dropna().unique() avg_words = unique.str.token_count().mean() return avg_words >= _TEXT_MIN_AVERAGE_WORD_COUNT + + +_LOW_CARDINALITY_MIN_ROWS = 151 +_LOW_CARDINALITY_MAX_UNIQUE_VALUES = 3 +_LOW_CARDINALITY_PREFIX_ROWS = 1024 + + +def _is_arrow_low_cardinality(array: pa.Array | pa.ChunkedArray) -> bool: + if len(array) < _LOW_CARDINALITY_MIN_ROWS: + return False + + # A prefix holds a subset of the distinct values, so most columns are + # ruled out without a full pass. + options = pc.CountOptions(mode="all") + prefix = array.slice(0, _LOW_CARDINALITY_PREFIX_ROWS) + num_unique = pc.call_function("count_distinct", [prefix], options).as_py() + if num_unique > _LOW_CARDINALITY_MAX_UNIQUE_VALUES: + return False + + num_unique = pc.call_function("count_distinct", [array], options).as_py() + return 1 < num_unique <= _LOW_CARDINALITY_MAX_UNIQUE_VALUES + + +def _is_series_low_cardinality(ser: pd.Series | cudf.Series) -> bool: + if len(ser) < _LOW_CARDINALITY_MIN_ROWS: + return False + + # A prefix holds a subset of the distinct values, so most columns are + # ruled out without a full pass. + prefix = ser.iloc[:_LOW_CARDINALITY_PREFIX_ROWS] + if prefix.nunique(dropna=False) > _LOW_CARDINALITY_MAX_UNIQUE_VALUES: + return False + + num_unique = ser.nunique(dropna=False) + return 1 < num_unique <= _LOW_CARDINALITY_MAX_UNIQUE_VALUES diff --git a/test/test_stype.py b/test/test_stype.py index f6534fc04..dd81ae89e 100644 --- a/test/test_stype.py +++ b/test/test_stype.py @@ -145,6 +145,40 @@ def test_id_detection() -> None: } +@pytest.mark.parametrize("backend", _BACKENDS) +def test_low_cardinality_detection(backend: str) -> None: + def make_table(num_rows: int) -> pa.Table | pd.DataFrame | cudf.DataFrame: + data = { + "binary": [i % 2 for i in range(num_rows)], + "ternary": [(0.5, 1.5, None)[i % 3] for i in range(num_rows)], + "count": [(0, 1, 2, None)[i % 4] for i in range(num_rows)], + # Two more values in the last rows only. + "late": [ + i % 2 if i < num_rows - 2 else i for i in range(num_rows) + ], + "constant": [1] * num_rows, + } + if backend == "pandas": + return pd.DataFrame(data) + if backend == "arrow": + return pa.table(data) + cudf = pytest.importorskip("cudf") + return cudf.DataFrame(data) + + expected = dict.fromkeys( + ["binary", "ternary", "count", "late", "constant"], + Stype.numerical, + ) + assert infer_stypes(make_table(2048)) == expected + assert infer_stypes(make_table(150), low_cardinality="infer") == expected + for num_rows in (151, 2048): + assert infer_stypes(make_table(num_rows), low_cardinality="infer") == { + **expected, + "binary": Stype.categorical, + "ternary": Stype.categorical, + } + + def test_overrides() -> None: table = pa.table( { From df413d382c706349c6fc0347d3fb160e9430664f Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Fri, 25 Sep 2026 23:18:23 +0000 Subject: [PATCH 2/2] mark as experimental --- benchmark/tabular/model.py | 2 +- sdm/stype.py | 12 ++---------- test/test_stype.py | 5 +++-- 3 files changed, 6 insertions(+), 13 deletions(-) diff --git a/benchmark/tabular/model.py b/benchmark/tabular/model.py index 2e02e4f94..9b7180cc3 100644 --- a/benchmark/tabular/model.py +++ b/benchmark/tabular/model.py @@ -362,7 +362,7 @@ def _preprocess( # codes would follow the order of appearance: pin the categories seen # in training in value order instead; other values become missing. if is_train: - stypes = sdm.infer_stypes(X, low_cardinality="infer") + stypes = sdm.infer_stypes(X, _low_cardinality="infer") self._low_cardinality_dtypes = { column: pd.CategoricalDtype( categories=np.sort(X[column].dropna().unique()) diff --git a/sdm/stype.py b/sdm/stype.py index eb67962f2..5f6e5130d 100644 --- a/sdm/stype.py +++ b/sdm/stype.py @@ -58,7 +58,7 @@ def infer_stypes( *, text: Literal["off", "infer", "drop"] = "off", id: Literal["off", "infer", "drop"] = "off", - low_cardinality: Literal["off", "infer"] = "off", + _low_cardinality: Literal["off", "infer"] = "off", unsupported: Literal["error", "warn", "drop"] = "error", ) -> dict[str, StypeLike]: r"""Infer semantic types from raw data statistics. @@ -81,10 +81,6 @@ def infer_stypes( :attr:`~Stype.id` if its name contains ``"id"`` as a whole word (*e.g.*, ``"user_id"``, ``"userId"``, ``"id"``, but not ``"solid"`` or ``"covid"``). - * Integer, floating-point, and decimal columns are inferred as - :attr:`~Stype.categorical` if the table has more than 150 rows and the - column contains two or three distinct values, counting missing values as - one. Args: table: A :class:`pandas.DataFrame`, :class:`pyarrow.Table`, or @@ -98,10 +94,6 @@ def infer_stypes( ``"off"`` disables :attr:`~Stype.id` column detection. ``"infer"`` includes inferred :attr:`~Stype.id` columns. ``"drop"`` omits inferred :attr:`~Stype.id` columns. - low_cardinality: The detection policy for low-cardinality integer, - floating-point, and decimal columns. - ``"off"`` keeps them :attr:`~Stype.numerical`. - ``"infer"`` infers them as :attr:`~Stype.categorical`. unsupported: How to handle unsupported dtypes. ``"error"`` raises a :class:`TypeError`. ``"warn"`` emits a warning and omits the column. @@ -151,7 +143,7 @@ def infer_stypes( continue try: - stype = fn(name, column, text, id, low_cardinality) + stype = fn(name, column, text, id, _low_cardinality) except TypeError: if unsupported == "error": raise diff --git a/test/test_stype.py b/test/test_stype.py index dd81ae89e..f38f8ac25 100644 --- a/test/test_stype.py +++ b/test/test_stype.py @@ -170,9 +170,10 @@ def make_table(num_rows: int) -> pa.Table | pd.DataFrame | cudf.DataFrame: Stype.numerical, ) assert infer_stypes(make_table(2048)) == expected - assert infer_stypes(make_table(150), low_cardinality="infer") == expected + assert infer_stypes(make_table(150), _low_cardinality="infer") == expected for num_rows in (151, 2048): - assert infer_stypes(make_table(num_rows), low_cardinality="infer") == { + table = make_table(num_rows) + assert infer_stypes(table, _low_cardinality="infer") == { **expected, "binary": Stype.categorical, "ternary": Stype.categorical,