Skip to content

[python-package] Improve pandas categorical test coverage #7376

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jameslamb merged 11 commits into from
Jul 27, 2026
+82 −0
Merged

[python-package] Improve pandas categorical test coverage #7376

Changes from 10 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
1f86647
Add categorical tests
maxzw Jul 22, 2026
f4b833c
Merge branch 'main' into tests/pandas-categorical-null-and-registred-…
maxzw Jul 22, 2026
650fc50
Combine None/NaN tests in single test run
maxzw Jul 23, 2026
009e40f
Improve strictness for test_pandas_categorical_encoding_registered_bu…
maxzw Jul 23, 2026
dd1351d
Align ordered / unordered
maxzw Jul 23, 2026
d32c33f
Use slicing instead of defining full set upfront (cannot be done in p…
maxzw Jul 24, 2026
aa169e2
Remove unnecessary .reset_index
maxzw Jul 24, 2026
f4971cb
Fix identical codes
maxzw Jul 25, 2026
4ea329d
Remove 'pandas' prefix
maxzw Jul 25, 2026
aa9a446
Merge branch 'main' into tests/pandas-categorical-null-and-registred-…
maxzw Jul 25, 2026
78d2b3c
Merge main
maxzw Jul 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 106 additions & 8 deletions tests/python_package_test/test_pandas.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
# coding: utf-8
from typing import Any, Dict
from zoneinfo import ZoneInfo

import numpy as np
Expand All @@ -11,6 +12,22 @@
pd = pytest.importorskip("pandas")


# ----------------------------------------------------------------------------------------------- #
# UTILITIES #
# ----------------------------------------------------------------------------------------------- #


def dummy_dataset_params() -> Dict[str, Any]:
return {
"min_data_in_bin": 1,
"min_data_in_leaf": 1,
}


# ----------------------------------------------------------------------------------------------- #
# UNIT TESTS #
# ----------------------------------------------------------------------------------------------- #

# ------------------------------------------- CATEGORICAL ----------------------------------------- #


Expand All @@ -32,7 +49,7 @@ def test_pandas_categorical_encoding(tmp_path):
)
y = [0, 1, 0, 1, 0]

ds = lgb.Dataset(df, label=y, params={"min_data_in_bin": 1})
ds = lgb.Dataset(df, label=y, params=dummy_dataset_params())
ds.construct()

assert ds.num_data() == 5
Expand All @@ -55,7 +72,7 @@ def test_pandas_categorical_encoding(tmp_path):
"num_col": [1.0, 2.0, 3.0, 4.0, 5.0],
}
)
ref_ds = lgb.Dataset(ref_df, label=y, categorical_feature=[0, 1], params={"min_data_in_bin": 1})
ref_ds = lgb.Dataset(ref_df, label=y, categorical_feature=[0, 1], params=dummy_dataset_params())
ref_ds.construct()

assert_datasets_equal(tmp_path, ds, ref_ds)
Expand All @@ -66,12 +83,11 @@ def test_pandas_categorical_encoding_unseen_category(tmp_path):
train_values = ["a", "b", "c", "a", "b"]
valid_values = ["a", "c", "d", "d", "a"] # "d" is unseen in training data

params = {"min_data_in_bin": 1, "min_data_in_leaf": 1}
train_df = pd.DataFrame({"cat_col": pd.Categorical(train_values), "num_col": [1.0, 2.0, 3.0, 4.0, 5.0]})
valid_df = pd.DataFrame({"cat_col": pd.Categorical(valid_values), "num_col": [6.0, 7.0, 8.0, 9.0, 10.0]})

train_ds = lgb.Dataset(train_df, label=[0, 1, 0, 1, 0], params=params)
valid_ds = lgb.Dataset(valid_df, label=[1, 0, 1, 0, 1], reference=train_ds, params=params)
train_ds = lgb.Dataset(train_df, label=[0, 1, 0, 1, 0], params=dummy_dataset_params())
valid_ds = lgb.Dataset(valid_df, label=[1, 0, 1, 0, 1], reference=train_ds, params=dummy_dataset_params())
train_ds.construct()
valid_ds.construct()

Expand All @@ -82,12 +98,94 @@ def test_pandas_categorical_encoding_unseen_category(tmp_path):
"num_col": [6.0, 7.0, 8.0, 9.0, 10.0],
}
)
ref_valid_ds = lgb.Dataset(ref_valid_df, label=[1, 0, 1, 0, 1], reference=train_ds, params=params)
ref_valid_ds = lgb.Dataset(ref_valid_df, label=[1, 0, 1, 0, 1], reference=train_ds, params=dummy_dataset_params())
ref_valid_ds.construct()

assert_datasets_equal(tmp_path, valid_ds, ref_valid_ds)


def test_categorical_encoding_registered_but_unobserved(tmp_path):
# Define full DataFrame with all categories observed
full_df = pd.DataFrame(
{
"unordered_col": pd.Categorical(["a", "b", "c", "d"]),
"ordered_col": pd.Categorical(["e", "f", "g", "h"], ordered=True),
}
)

# Slice train from full_df so all categories are preserved despite not all being observed
train_df = full_df.iloc[[0, 2, 2]] # ["a", "c", "c"] and ["e", "g", "g"]
valid_df = pd.DataFrame(
{
"unordered_col": pd.Categorical(["a", "b", "d"]),
"ordered_col": pd.Categorical(["h", "e", "f"], ordered=True),
}
)

train_ds = lgb.Dataset(train_df, label=[0, 1, 0], params=dummy_dataset_params())
valid_ds = lgb.Dataset(valid_df, label=[0, 1, 0], reference=train_ds, params=dummy_dataset_params())
train_ds.construct()
valid_ds.construct()

assert train_ds.pandas_categorical[0] == ["a", "b", "c", "d"]
assert train_ds.pandas_categorical[1] == ["e", "f", "g", "h"]
assert train_ds.params["categorical_column"] == [0] # only unordered column is treated as categorical

# Python-side encoding: both ordered and unordered columns use all registered categories to encode
valid_df_encoded = lgb.basic._data_from_pandas(
data=valid_df,
feature_name="auto",
categorical_feature="auto",
pandas_categorical=train_ds.pandas_categorical,
)[0]
assert valid_df_encoded[:, 0].tolist() == [0.0, 1.0, 3.0] # a -> 0, b -> 1, d -> 3
assert valid_df_encoded[:, 1].tolist() == [3.0, 0.0, 1.0] # h -> 3, e -> 0, f -> 1

# C++ binning
# - Unordered columns: only codes observed during training are binned. Unseen codes are treated as missing.
# - Ordered columns: treats as continuous. Unseen values interpolate (e<f<g) or clip (h clipped to g).
ref_valid_df = pd.DataFrame(
{
"unordered_col": pd.Categorical(["a", None, None], categories=["a", "b", "c", "d"]),
"ordered_col": pd.Categorical(["g", "e", "g"], categories=["e", "f", "g", "h"], ordered=True),
}
)
ref_valid_ds = lgb.Dataset(ref_valid_df, label=[0, 1, 0], reference=train_ds, params=dummy_dataset_params())
ref_valid_ds.construct()

assert_datasets_equal(tmp_path, valid_ds, ref_valid_ds)

@jameslamb jameslamb Jul 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This tests that these Datasets are equal but they could also be equal in the case of a bug where these inputs were all handled incorrectly. For example, if they were all converted to floats and treated as continuous variables.

I think this test should be strengthened with some assertions that the variables were handled in the expected way.

Like:

  • both columns were detected as categorical
  • the encodings in Dataset are as expected (I think you could see this in a model file if you added an lgb.train() to this or output of Dataset._dump_text(), don't recall exactly and can't spend the time to look right now)

@maxzw maxzw Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've rewritten the test a little bit to improve the strictness.

  • Aligned the (relative) values and categories between ordered / unordered to emphasize the differences in how they are handled.
  • Check more metadata (e.g. pandas_categorical and params["categorical_column"])
  • Inspect the resulting encodings from _data_from_pandas directly (so python-side, but I've also kept the C++ side verification with a comment explaining why the output is expected)

@jameslamb jameslamb Jul 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep that's good for this PR, much stricter, thank you.



def test_categorical_with_missing_values(tmp_path):
categories = ["a", "b"]
values_none = ["a", "b", None, "a", None]
values_nan = ["b", "a", np.nan, "b", np.nan]

df = pd.DataFrame(
{
"cat_none": pd.Categorical(values_none, categories=categories),
"cat_nan": pd.Categorical(values_nan, categories=categories),
"num": [1.0, 2.0, 3.0, 4.0, 5.0],
}
)
y = [0, 1, 0, 1, 0]

ds = lgb.Dataset(df, label=y, params=dummy_dataset_params())
ds.construct()
assert ds.pandas_categorical == [categories, categories]

ref_df = pd.DataFrame(
{
"cat_none": [0.0, 1.0, np.nan, 0.0, np.nan],
"cat_nan": [1.0, 0.0, np.nan, 1.0, np.nan],
"num": [1.0, 2.0, 3.0, 4.0, 5.0],
}
)
ref_ds = lgb.Dataset(ref_df, label=y, categorical_feature=[0, 1], params=dummy_dataset_params())
ref_ds.construct()
assert_datasets_equal(tmp_path, ds, ref_ds)


def test_pandas_dataset_construction_with_high_cardinality_categorical_succeeds(rng):
X = pd.DataFrame({"x1": rng.integers(low=0, high=5_000, size=(10_000,))})
y = rng.uniform(size=(10_000,))
Expand Down Expand Up @@ -158,7 +256,7 @@ def test_pandas_supported_dtypes(tmp_path, dtype, values):
df = pd.DataFrame({"test_col": pd.Series(values, dtype=dtype), "num_col": [4.0, 5.0, 6.0]})
y = [0, 1, 0]

ds = lgb.Dataset(df, label=y, params={"min_data_in_bin": 1})
ds = lgb.Dataset(df, label=y, params=dummy_dataset_params())
ds.construct()

assert ds.num_data() == 3
Expand All @@ -168,7 +266,7 @@ def test_pandas_supported_dtypes(tmp_path, dtype, values):

# Verify values are preserved
ref_df = pd.DataFrame({"test_col": values, "num_col": [4.0, 5.0, 6.0]})
ref_ds = lgb.Dataset(ref_df, label=y, params={"min_data_in_bin": 1})
ref_ds = lgb.Dataset(ref_df, label=y, params=dummy_dataset_params())
ref_ds.construct()

assert_datasets_equal(tmp_path, ds, ref_ds)
Expand Down
Loading