Count distinct values for all columns of a numeric array at once in modality detection

changelog/1255.changed.md

@@ -0,0 +1 @@
+Faster modality detection on wide inputs: a numeric array's distinct values are counted for all columns at once instead of per column, a numeric column stored as `object` is recognized with `pd.api.types.infer_dtype` instead of a value-by-value walk, and the nullable-dtype coercion is decided once per dtype. The inferred modalities are unchanged.

src/tabpfn/preprocessing/clean.py

@@ -226,12 +226,16 @@ def coerce_nullable_dtypes_to_numpy(X: pd.DataFrame) -> pd.DataFrame:
 
     ``category``/``string``/``object`` columns are left untouched.
     """
-    cols = [
-        col
-        for col, dtype in X.dtypes.items()
+    dtypes = X.dtypes
+    # Decided once per distinct dtype rather than once per column: a wide frame has
+    # thousands of columns and a handful of dtypes.
+    dtypes_to_cast = {
+        dtype
+        for dtype in dtypes.unique()
         if pd.api.types.is_bool_dtype(dtype)
         or (pd.api.types.is_extension_array_dtype(dtype) and dtype.kind in "iuf")
-    ]
+    }
+    cols = [col for col, dtype in dtypes.items() if dtype in dtypes_to_cast]
     return _cast_columns(X, cols, "float64")
 
 

src/tabpfn/preprocessing/modality_detection.py

@@ -7,8 +7,8 @@
 import math
 import warnings
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
 
+import numpy as np
 import pandas as pd
 
 from tabpfn.errors import TabPFNUserError
@@ -21,9 +21,6 @@
     build_input_feature_names,
 )
 
-if TYPE_CHECKING:
-    import numpy as np
-
 _EARLY_EXIT_PREFIX_ROWS = 1024
 
 #: Cap on how many column names the likely-text warning lists, so a wide frame of
@@ -68,18 +65,35 @@ def detect_feature_modalities(
     features: list[Feature] = []
     big_enough_n_to_infer_cat = len(X) > min_samples_for_inference
     unique_feature_names = build_input_feature_names(feature_names, X.shape[1])
+    provided = set(provided_categorical_indices or ())
+    decided_at = _decided_at(
+        max_unique_for_category=max_unique_for_category,
+        min_unique_for_numerical=min_unique_for_numerical,
+        min_cardinality_for_text=min_cardinality_for_text,
+    )
+    # A numeric array needs no per-column parsing: every column is numeric, so only
+    # the distinct-value count decides, and that is counted for all columns at once.
+    n_unique_per_column = _numeric_n_unique_per_column(X, decided_at=decided_at)
     for i, index in enumerate(range(X.shape[1])):
         feature_name = unique_feature_names[i]
-        X_slice: np.ndarray = X[:, index]
-        reported_categorical = index in (provided_categorical_indices or ())
-        feat_modality = _detect_feature_modality(
-            s=pd.Series(X_slice, name=feature_name),
-            reported_categorical=reported_categorical,
-            max_unique_for_category=max_unique_for_category,
-            min_unique_for_numerical=min_unique_for_numerical,
-            min_cardinality_for_text=min_cardinality_for_text,
-            big_enough_n_to_infer_cat=big_enough_n_to_infer_cat,
-        )
+        reported_categorical = index in provided
+        if n_unique_per_column is not None:
+            feat_modality = _numeric_modality(
+                n_unique=int(n_unique_per_column[index]),
+                reported_categorical=reported_categorical,
+                max_unique_for_category=max_unique_for_category,
+                min_unique_for_numerical=min_unique_for_numerical,
+                big_enough_n_to_infer_cat=big_enough_n_to_infer_cat,
+            )
+        else:
+            feat_modality = _detect_feature_modality(
+                s=pd.Series(X[:, index], name=feature_name),
+                reported_categorical=reported_categorical,
+                max_unique_for_category=max_unique_for_category,
+                min_unique_for_numerical=min_unique_for_numerical,
+                min_cardinality_for_text=min_cardinality_for_text,
+                big_enough_n_to_infer_cat=big_enough_n_to_infer_cat,
+            )
         features.append(Feature(name=feature_name, modality=feat_modality))
     feature_schema = FeatureSchema(features=features)
     _warn_on_text(feature_schema)
@@ -142,17 +156,10 @@ def _detect_feature_modality(
         "Categorical dtype must be converted before modality detection; "
         "preserve its intent in provided_categorical_indices."
     )
-    # Early exit: once a prefix already clears every threshold below, the full
-    # count would land in the same bucket, so skip scanning the rest.
-    # min_cardinality_for_text is included since it can exceed the other two.
-    decided_at = (
-        max(
-            max_unique_for_category,
-            min_unique_for_numerical,
-            min_cardinality_for_text,
-            1,
-        )
-        + 1
+    decided_at = _decided_at(
+        max_unique_for_category=max_unique_for_category,
+        min_unique_for_numerical=min_unique_for_numerical,
+        min_cardinality_for_text=min_cardinality_for_text,
     )
     n_unique = 0
     if len(s) > _EARLY_EXIT_PREFIX_ROWS:
@@ -169,15 +176,13 @@ def _detect_feature_modality(
         return FeatureModality.CONSTANT
 
     if _is_numeric_pandas_series(s):
-        if _detect_numeric_as_categorical(
+        return _numeric_modality(
             n_unique=n_unique,
             reported_categorical=reported_categorical,
             max_unique_for_category=max_unique_for_category,
             min_unique_for_numerical=min_unique_for_numerical,
             big_enough_n_to_infer_cat=big_enough_n_to_infer_cat,
-        ):
-            return FeatureModality.CATEGORICAL
-        return FeatureModality.NUMERICAL
+        )
 
     # A pandas `category` column never arrives here as such: `X` is a numpy array
     # by now, and its intent travels in `provided_categorical_indices` instead.
@@ -192,9 +197,115 @@ def _detect_feature_modality(
     )
 
 
+def _decided_at(
+    *,
+    max_unique_for_category: int,
+    min_unique_for_numerical: int,
+    min_cardinality_for_text: int,
+) -> int:
+    """The distinct-value count at which every threshold below is cleared.
+
+    Once a prefix of a column already holds this many distinct values, the full count
+    would land in the same bucket, so the rest of the column need not be scanned.
+    `min_cardinality_for_text` is included since it can exceed the other two.
+    """
+    return (
+        max(
+            max_unique_for_category,
+            min_unique_for_numerical,
+            min_cardinality_for_text,
+            1,
+        )
+        + 1
+    )
+
+
+def _numeric_modality(
+    *,
+    n_unique: int,
+    reported_categorical: bool,
+    max_unique_for_category: int,
+    min_unique_for_numerical: int,
+    big_enough_n_to_infer_cat: bool,
+) -> FeatureModality:
+    """The modality of a numeric column with `n_unique` distinct values (NaN counted).
+
+    A constant (or all-missing) column is `CONSTANT` unless declared categorical, so
+    that it still routes through the ordinal encoder instead of crashing as a
+    constant numeric column when predict sees an unseen value.
+    """
+    if n_unique <= 1 and not reported_categorical:
+        return FeatureModality.CONSTANT
+    if _detect_numeric_as_categorical(
+        n_unique=n_unique,
+        reported_categorical=reported_categorical,
+        max_unique_for_category=max_unique_for_category,
+        min_unique_for_numerical=min_unique_for_numerical,
+        big_enough_n_to_infer_cat=big_enough_n_to_infer_cat,
+    ):
+        return FeatureModality.CATEGORICAL
+    return FeatureModality.NUMERICAL
+
+
+def _numeric_n_unique_per_column(
+    X: np.ndarray, *, decided_at: int
+) -> np.ndarray | None:
+    """Distinct values per column of a numeric or bool array, NaN counted as a value.
+
+    `None` for anything else (an object array is parsed column by column). Mirrors
+    the per-column early exit: a column whose first `_EARLY_EXIT_PREFIX_ROWS` rows
+    already hold `decided_at` distinct values keeps that prefix count, which lands
+    in the same bucket as the full count; only the other columns are counted in
+    full.
+    """
+    if not isinstance(X, np.ndarray) or X.ndim != 2 or X.dtype.kind not in "biuf":
+        return None
+    n_rows, n_columns = X.shape
+    if n_rows == 0:
+        return np.zeros(n_columns, dtype=np.int64)
+    if n_rows <= _EARLY_EXIT_PREFIX_ROWS:
+        return _count_distinct_per_column(X)
+    n_unique = _count_distinct_per_column(X[:_EARLY_EXIT_PREFIX_ROWS])
+    undecided = np.flatnonzero(n_unique < decided_at)
+    if len(undecided):
+        n_unique[undecided] = _count_distinct_per_column(X[:, undecided])
+    return n_unique
+
+
+def _count_distinct_per_column(X: np.ndarray) -> np.ndarray:
+    """`pd.Series(column).nunique(dropna=False)` for every column of a numeric array.
+
+    Sorting puts equal values next to each other and NaN last, so the count is one
+    plus the number of adjacent unequal pairs, with NaN counted once when present.
+    `-0.0` equals `0.0` and `inf` equals `inf` here as under `nunique`.
+    """
+    values = np.sort(X, axis=0)
+    if values.dtype.kind == "f":
+        missing = np.isnan(values)
+        differs = (values[1:] != values[:-1]) & ~missing[1:]
+        return (
+            differs.sum(axis=0)
+            + (~missing).any(axis=0).astype(np.int64)
+            + missing.any(axis=0).astype(np.int64)
+        )
+    return (values[1:] != values[:-1]).sum(axis=0) + 1
+
+
+#: `pd.api.types.infer_dtype` kinds whose every non-missing value is a number. A
+#: `string` or `mixed` column is not settled by them: a spelled-out number counts too.
+_INFERRED_NUMERIC_KINDS = frozenset(
+    {"integer", "floating", "mixed-integer-float", "boolean", "decimal", "empty"}
+)
+
+
 def _is_numeric_pandas_series(s: pd.Series) -> bool:
     if pd.api.types.is_numeric_dtype(s.dtype):
         return True
+    # A numeric column stored as object is the common case: a frame with one
+    # non-numeric column arrives as a single object array. `infer_dtype` settles it in
+    # C rather than a Python-level walk over every value.
+    if pd.api.types.infer_dtype(s, skipna=True) in _INFERRED_NUMERIC_KINDS:
+        return True
     if PANDAS_BELOW_3:
         return all(_is_numeric_or_missing_for_old_pandas(value) for value in s)
     # The generator above stops at the first non-numeric value; `pd.to_numeric`

tests/test_preprocessing/test_data_cleaning.py

@@ -21,6 +21,7 @@
 from tabpfn.preprocessing.clean import (
     _is_single_float_block,
     clean_data_transform,
+    coerce_nullable_dtypes_to_numpy,
     fix_dtypes,
     process_text_na_dataframe,
 )
@@ -1344,3 +1345,30 @@ def test__fit_predict__unicode_array__is_accepted_like_an_object_array(
     np.testing.assert_array_equal(
         fitted_on_frame.predict(X), fitted_on_frame.predict(X.astype(object))
     )
+
+
+def test__coerce_nullable_dtypes_to_numpy__selects_by_dtype() -> None:
+    """Bool and nullable numeric columns become float64, all others keep their dtype."""
+    X = pd.DataFrame(
+        {
+            "bool": [True, False, True],
+            "boolean": pd.array([True, None, False], dtype="boolean"),
+            "int64_na": pd.array([1, None, 3], dtype="Int64"),
+            "float64_na": pd.array([1.5, None, 3.5], dtype="Float64"),
+            "uint8_na": pd.array([1, 2, None], dtype="UInt8"),
+            "int": [1, 2, 3],
+            "float": [1.5, 2.5, 3.5],
+            "cat": pd.Categorical(["a", "b", "a"]),
+            "string": pd.array(["a", None, "b"], dtype="string"),
+            "obj": ["a", 1, None],
+            "bool_again": [False, False, True],
+        }
+    )
+    out = coerce_nullable_dtypes_to_numpy(X)
+    cast = ["bool", "boolean", "int64_na", "float64_na", "uint8_na", "bool_again"]
+    assert all(out[c].dtype == np.float64 for c in cast)
+    for c in X.columns:
+        if c not in cast:
+            assert out[c].dtype == X[c].dtype, c
+    assert list(out.columns) == list(X.columns)
+    assert out["boolean"].isna().tolist() == [False, True, False]