[python-package] Add categorical support for Polars/Arrow via Narwhals

[maxzw]

Add categorical support for Polars/Arrow via Narwhals

This PR is a follow up of #7275 (cc @borchero), contributes to #6204 and replaces #7264.

Motivation

Now that DataFrame inputs are routed through narwhals, it seemed like a small step to also pick up categorical columns from polars the same way we already do for pandas. Previously, pl.Categorical/pl.Enum columns were silently ignored, which meant they'd either fail downstream or get mis-encoded.

What Changed?

The pandas-specific _data_from_pandas helper has been replaced with _data_from_narwhals, which detects and encodes categoricals via narwhals. The pandas fast path (_pandas_to_numpy) is preserved for the no-copy case. The pandas_categorical attribute name is kept as-is to stay compatible with saved models (an issue was opened to rename this in the future: #7361).

Added generic dtype validation in the Narwhals dataframe path with _check_for_bad_narwhals_dtypes so LightGBM now explicitly accepts only integer, float, bool and categorical columns across supported backends. I consider this in scope because the old pandas path had implicit dtype validation via pandas_to_numpy conversions, and the new narwhals path needs an explicit equivalent to preserve the same user-facing error behavior.

Note: this validation also includes a pandas-like fallback check to avoid false-negatives from current Narwhals dtype-inference gaps. Issues are known at Narwhals at the time of writing and referenced in the comments. Can be simplified once resolved.

New tests added for polars and arrow based on #7359 and #7376. Currently, only pandas, polars and pyarrow are covered by categorical tests, so those are the only backends I'd claim as supported for now. Other narwhals backends (e.g. modin, cudf, dask) should benefit too in principle - I haven't tested that.

[jameslamb]

replaces #7264

@Ic3fr0g you are the author of #7264 and also ❤️'d this so ... should we close #7264?

[Ic3fr0g]

Absolutely! Done now!

[jameslamb]

hey thanks for keeping this updated @maxzw . I'm going to cancel the latest CI run...I'm about to start the v4.7.0 release and I want to be sure the remaining jobs get priority in scheduling. I'll restart it for you here when that's done. Sorry for the brief disruption.

[maxzw]

Hi, quick status update on this PR.

I've spent more time investigating how the Polars global string cache issue I described above can be handled uniformly across dataframe backends, and how ordered vs unordered categoricals interact with the C++ bin mapper.

The current state:

I've split out a final prerequisite PR (#7376) that adds pandas tests documenting the expected categorical encoding behavior (test_pandas_categorical_encoding_registered_but_unobserved + test_pandas_categorical_with_missing_values). I'd like to get that merged first to establish a regression baseline for pandas. This PR then adds the equivalent polars/arrow tests alongside the implementation.

Once #7376 is in I'll resolve conflicts and mark it ready for proper review. Apologies for continuing to push commits on an open PR, I understand that makes it hard to review a moving target.

[maxzw]

@jameslamb with and #7376 merged and the merge conflicts resolved, I've converged on the current implementation and am stepping back to allow time for a proper review.

Categorical support for Polars and Arrow is (to the best of my knowledge) fully aligned with pandas, except for Polars' unordered categorical (pl.Categorical) where I found one edge-case. I'd like to point your attention to this one divergence; it's a deliberate consequence of Polars' type system rather than a bug and I'd appreciate input on how to handle it.

Why we use .unique() for Polars' unordered categoricals

Polars' recent categorical refactor removed per-column category ownership in favor of a single global string cache. That refactor introduced a category leakage problem that I documented above. Since then, Polars has even deprecated cat.get_categories() entirely, advising users to use .unique() instead. Note that narwhals' cat.get_categories() also exhibits the leakage since it delegates to Polars' implementation.

For Polars' unordered categoricals we use col.unique().drop_nulls().sort().to_list() which gives us per-column observed values only, avoids cross-column pollution from the global cache, and uses the non-deprecated API.

Side effect: divergence from pandas on unobserved categories

Because .unique() returns only values actually present, declared-but-unobserved categories are not included in the encoding. This creates a divergence with pandas, where the category set is attached to the dtype and includes all declared categories regardless of whether they appear in the data. However, this divergence is not a trade-off of the .unique() approach specifically. Pandas allows users to declare categories beyond what's observed (pd.Categorical(["a"], categories=["a, "b])); Polars' pl.Categorical has no equivalent mechanism, so there is no way to match this behavior regardless of implementation choice.

In practice the divergence only surfaces when a user slices a DataFrame into train/valid splits, which I explicitly added as a unit test in #7376. For example:

# Full dataset with all categories observed
full_df_pd = pd.DataFrame({"col": pd.Categorical(["a", "b", "c", "d"])})
full_df_pl = pl.DataFrame({"col": pl.Series(["a", "b", "c", "d"], dtype=pl.Categorical)})

# Slice to simulate a train split (only "a" and "c" observed)
train_df_pd = full_df_pd.iloc[[0, 2, 2]]  # categories still ["a", "b", "c", "d"]
train_df_pl = full_df_pl[[0, 2, 2]]       # .unique() gives ["a", "c"] only

lgb.Dataset(train_df_pd, label=[0, 1, 0]).construct().pandas_categorical[0]  # ['a', 'b', 'c', 'd'], full set preserved through slice
lgb.Dataset(train_df_pl, label=[0, 1, 0]).construct().pandas_categorical[0]  # ['a', 'c'], observed values only

During inference, an unobserved category "b" will be encoded as 1 (pandas) vs -1/missing (polars) on the Python side. However, the C++ bin mapping is the same for both backends since it's built from codes observed during training (see the aforementioned unit test), so the trained model and predictions should be identical.

Options for handling this divergence

1. Accept and document (currently implemented). Since the C++ result is identical this is mostly a Python-side metadata difference. However, I don't know how much users rely on the .pandas_categorical attribute behaving the same for all backends.

2. Use .unique() for all backends. Deliberately ignore pandas' declared-but-unobserved categories and use observed values. This simplifies the code (single path for all backends) and mirrors the C++ bin mapper which already only uses categories observed during training. However, it intentionally breaks existing pandas behavior for users who rely on .pandas_categorical.

3. Add an explicit category list parameter. Allow users to declare the full category universe via some Dataset parameter. This would make pandas_categorical consistent across backends for users who need it, but muddies the API with Polars-specific gap-filling behavior. I would advice against this, but it is an option.

4. Something else I'm missing?

I lean towards option 1 or 2 since the .pandas_categorical attribute is not mentioned anywhere in the documentation, but I'm obviously biased. What are your thoughts?

[maxzw]

Just to document; In the future, nw.factorize might be a very nice abstraction for getting categorical mappings across all Narwhals-supported backends

[borchero]

Sorry for the late review 😅 I have left a few preliminary comments