wehub-resource-sync
414 lines
15 KiB
Python
414 lines
15 KiB
Python
# coding: utf-8
|
|
import filecmp
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import lightgbm as lgb
|
|
|
|
from .utils import np_assert_array_equal
|
|
|
|
pl = pytest.importorskip("polars")
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------- #
|
|
# UTILITIES #
|
|
# ----------------------------------------------------------------------------------------------- #
|
|
|
|
_INTEGER_TYPES = [pl.Int8, pl.Int16, pl.Int32, pl.Int64, pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64]
|
|
_FLOAT_TYPES = [pl.Float32, pl.Float64]
|
|
|
|
|
|
def generate_simple_polars_frame() -> pl.DataFrame:
|
|
values = [1, 2, 3, 4, 5]
|
|
bool_values = [True, True, False, False, True]
|
|
columns = {f"col_{i}": pl.Series(values, dtype=dtype) for i, dtype in enumerate(_INTEGER_TYPES + _FLOAT_TYPES)}
|
|
columns[f"col_{len(columns)}"] = pl.Series(bool_values, dtype=pl.Boolean)
|
|
return pl.DataFrame(columns)
|
|
|
|
|
|
def generate_nullable_polars_frame(dtype: Any) -> pl.DataFrame:
|
|
return pl.DataFrame(
|
|
{
|
|
"col_0": pl.Series([1, None, 3, 4, 5], dtype=dtype),
|
|
"col_1": pl.Series([None, 2, 3, 4, 5], dtype=dtype),
|
|
"col_2": pl.Series([1, 2, 3, 4, None], dtype=dtype),
|
|
"col_3": pl.Series([None, None, None, None, None], dtype=dtype),
|
|
}
|
|
)
|
|
|
|
|
|
def generate_dummy_polars_frame() -> pl.DataFrame:
|
|
return pl.DataFrame(
|
|
{
|
|
"a": pl.Series([1, 2, 3, 4, 5], dtype=pl.UInt8),
|
|
"b": pl.Series([0.5, 0.6, 0.1, 0.8, 1.5], dtype=pl.Float32),
|
|
}
|
|
)
|
|
|
|
|
|
def generate_random_polars_frame(
|
|
num_columns: int,
|
|
num_datapoints: int,
|
|
seed: int,
|
|
generate_nulls: bool = True,
|
|
values: Optional[np.ndarray] = None,
|
|
) -> pl.DataFrame:
|
|
return pl.DataFrame(
|
|
{
|
|
f"col_{i}": generate_random_polars_series(
|
|
num_datapoints, seed + i, generate_nulls=generate_nulls, values=values
|
|
)
|
|
for i in range(num_columns)
|
|
}
|
|
)
|
|
|
|
|
|
def generate_random_polars_series(
|
|
num_datapoints: int,
|
|
seed: int,
|
|
generate_nulls: bool = True,
|
|
values: Optional[np.ndarray] = None,
|
|
) -> pl.Series:
|
|
generator = np.random.default_rng(seed)
|
|
data = (
|
|
generator.standard_normal(num_datapoints).astype(np.float32)
|
|
if values is None
|
|
else generator.choice(values, size=num_datapoints, replace=True)
|
|
)
|
|
series = pl.Series("col", data, dtype=pl.Float32)
|
|
if generate_nulls:
|
|
indices = generator.choice(len(data), size=num_datapoints // 10)
|
|
series = series.scatter(indices, None)
|
|
return series
|
|
|
|
|
|
def dummy_dataset_params() -> Dict[str, Any]:
|
|
return {
|
|
"min_data_in_bin": 1,
|
|
"min_data_in_leaf": 1,
|
|
}
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------- #
|
|
# UNIT TESTS #
|
|
# ----------------------------------------------------------------------------------------------- #
|
|
|
|
# ------------------------------------------- DATASET ------------------------------------------- #
|
|
|
|
|
|
def assert_datasets_equal(tmp_path: Path, lhs: lgb.Dataset, rhs: lgb.Dataset):
|
|
lhs._dump_text(tmp_path / "polars.txt")
|
|
rhs._dump_text(tmp_path / "pandas.txt")
|
|
assert filecmp.cmp(tmp_path / "polars.txt", tmp_path / "pandas.txt")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("polars_frame_fn", "dataset_params"),
|
|
[ # Use lambda functions here to minimize memory consumption
|
|
(generate_simple_polars_frame, dummy_dataset_params()),
|
|
(generate_dummy_polars_frame, dummy_dataset_params()),
|
|
(lambda: generate_nullable_polars_frame(pl.Float32), dummy_dataset_params()),
|
|
(lambda: generate_nullable_polars_frame(pl.Int32), dummy_dataset_params()),
|
|
(lambda: generate_random_polars_frame(3, 1000, 42), {}),
|
|
(lambda: generate_random_polars_frame(100, 10000, 43), {}),
|
|
],
|
|
)
|
|
def test_dataset_construct_fuzzy(tmp_path, polars_frame_fn, dataset_params):
|
|
polars_frame = polars_frame_fn()
|
|
|
|
polars_dataset = lgb.Dataset(polars_frame, params=dataset_params)
|
|
polars_dataset.construct()
|
|
|
|
pandas_dataset = lgb.Dataset(polars_frame.to_pandas(), params=dataset_params)
|
|
pandas_dataset.construct()
|
|
|
|
assert_datasets_equal(tmp_path, polars_dataset, pandas_dataset)
|
|
|
|
|
|
def test_dataset_construct_fuzzy_boolean(tmp_path):
|
|
boolean_data = generate_random_polars_frame(10, 10000, 42, generate_nulls=False, values=np.array([True, False]))
|
|
float_data = boolean_data.cast(pl.Float32)
|
|
|
|
polars_dataset = lgb.Dataset(boolean_data)
|
|
polars_dataset.construct()
|
|
|
|
pandas_dataset = lgb.Dataset(float_data.to_pandas())
|
|
pandas_dataset.construct()
|
|
|
|
assert_datasets_equal(tmp_path, polars_dataset, pandas_dataset)
|
|
|
|
|
|
# -------------------------------------------- FIELDS ------------------------------------------- #
|
|
|
|
|
|
def test_dataset_construct_fields_fuzzy():
|
|
polars_frame = generate_random_polars_frame(3, 1000, 42)
|
|
polars_labels = generate_random_polars_series(1000, 42, generate_nulls=False)
|
|
polars_weights = generate_random_polars_series(1000, 42, generate_nulls=False)
|
|
polars_groups = pl.Series("group", [300, 400, 50, 250], dtype=pl.Int32)
|
|
|
|
polars_dataset = lgb.Dataset(polars_frame, label=polars_labels, weight=polars_weights, group=polars_groups)
|
|
polars_dataset.construct()
|
|
|
|
pandas_dataset = lgb.Dataset(
|
|
polars_frame.to_pandas(),
|
|
label=polars_labels.to_numpy(),
|
|
weight=polars_weights.to_numpy(),
|
|
group=polars_groups.to_numpy(),
|
|
)
|
|
pandas_dataset.construct()
|
|
|
|
# Check for equality
|
|
for field in ("label", "weight", "group"):
|
|
np_assert_array_equal(polars_dataset.get_field(field), pandas_dataset.get_field(field), strict=True)
|
|
np_assert_array_equal(polars_dataset.get_label(), pandas_dataset.get_label(), strict=True)
|
|
np_assert_array_equal(polars_dataset.get_weight(), pandas_dataset.get_weight(), strict=True)
|
|
|
|
|
|
# -------------------------------------------- LABELS ------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.parametrize("polars_type", _INTEGER_TYPES + _FLOAT_TYPES)
|
|
def test_dataset_construct_labels(polars_type):
|
|
data = generate_dummy_polars_frame()
|
|
labels = pl.Series("label", [0, 1, 0, 0, 1], dtype=polars_type)
|
|
dataset = lgb.Dataset(data, label=labels, params=dummy_dataset_params())
|
|
dataset.construct()
|
|
|
|
expected = np.array([0, 1, 0, 0, 1], dtype=np.float32)
|
|
np_assert_array_equal(expected, dataset.get_label(), strict=True)
|
|
|
|
|
|
def test_dataset_construct_labels_boolean():
|
|
data = generate_dummy_polars_frame()
|
|
labels = pl.Series("label", [False, True, False, False, True], dtype=pl.Boolean)
|
|
dataset = lgb.Dataset(data, label=labels, params=dummy_dataset_params())
|
|
dataset.construct()
|
|
|
|
expected = np.array([0, 1, 0, 0, 1], dtype=np.float32)
|
|
np_assert_array_equal(expected, dataset.get_label(), strict=True)
|
|
|
|
|
|
# ------------------------------------------- WEIGHTS ------------------------------------------- #
|
|
|
|
|
|
def test_dataset_construct_weights_none():
|
|
data = generate_dummy_polars_frame()
|
|
weight = pl.Series("weight", [1, 1, 1, 1, 1], dtype=pl.Float32)
|
|
dataset = lgb.Dataset(data, weight=weight, params=dummy_dataset_params())
|
|
dataset.construct()
|
|
assert dataset.get_weight() is None
|
|
assert dataset.get_field("weight") is None
|
|
|
|
|
|
@pytest.mark.parametrize("polars_type", _FLOAT_TYPES)
|
|
def test_dataset_construct_weights(polars_type):
|
|
data = generate_dummy_polars_frame()
|
|
weights = pl.Series("weight", [3, 0.7, 1.5, 0.5, 0.1], dtype=polars_type)
|
|
dataset = lgb.Dataset(data, weight=weights, params=dummy_dataset_params())
|
|
dataset.construct()
|
|
|
|
expected = np.array([3, 0.7, 1.5, 0.5, 0.1], dtype=np.float32)
|
|
np_assert_array_equal(expected, dataset.get_weight(), strict=True)
|
|
|
|
|
|
# -------------------------------------------- GROUPS ------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.parametrize("polars_type", _INTEGER_TYPES)
|
|
def test_dataset_construct_groups(polars_type):
|
|
data = generate_dummy_polars_frame()
|
|
groups = pl.Series("group", [2, 3], dtype=polars_type)
|
|
dataset = lgb.Dataset(data, group=groups, params=dummy_dataset_params())
|
|
dataset.construct()
|
|
|
|
expected = np.array([0, 2, 5], dtype=np.int32)
|
|
np_assert_array_equal(expected, dataset.get_field("group"), strict=True)
|
|
|
|
|
|
# ----------------------------------------- INIT SCORES ----------------------------------------- #
|
|
|
|
|
|
@pytest.mark.parametrize("polars_type", _INTEGER_TYPES + _FLOAT_TYPES)
|
|
def test_dataset_construct_init_scores_array(polars_type):
|
|
data = generate_dummy_polars_frame()
|
|
init_scores = pl.Series("init_score", [0, 1, 2, 3, 3], dtype=polars_type)
|
|
dataset = lgb.Dataset(data, init_score=init_scores, params=dummy_dataset_params())
|
|
dataset.construct()
|
|
|
|
expected = np.array([0, 1, 2, 3, 3], dtype=np.float64)
|
|
np_assert_array_equal(expected, dataset.get_init_score(), strict=True)
|
|
|
|
|
|
def test_dataset_construct_init_scores_table():
|
|
data = generate_dummy_polars_frame()
|
|
init_scores = pl.DataFrame(
|
|
{
|
|
"a": generate_random_polars_series(5, seed=1, generate_nulls=False),
|
|
"b": generate_random_polars_series(5, seed=2, generate_nulls=False),
|
|
"c": generate_random_polars_series(5, seed=3, generate_nulls=False),
|
|
}
|
|
)
|
|
dataset = lgb.Dataset(data, init_score=init_scores, params=dummy_dataset_params())
|
|
dataset.construct()
|
|
|
|
actual = dataset.get_init_score()
|
|
expected = init_scores.to_numpy().astype(np.float64)
|
|
np_assert_array_equal(expected, actual, strict=True)
|
|
|
|
|
|
# ------------------------------------------ PREDICTION ----------------------------------------- #
|
|
|
|
|
|
def assert_equal_predict_polars_pandas(booster: lgb.Booster, data: pl.DataFrame):
|
|
pandas_data = data.to_pandas()
|
|
|
|
p_polars = booster.predict(data)
|
|
p_pandas = booster.predict(pandas_data)
|
|
np_assert_array_equal(p_polars, p_pandas, strict=True)
|
|
|
|
p_raw_polars = booster.predict(data, raw_score=True)
|
|
p_raw_pandas = booster.predict(pandas_data, raw_score=True)
|
|
np_assert_array_equal(p_raw_polars, p_raw_pandas, strict=True)
|
|
|
|
p_leaf_polars = booster.predict(data, pred_leaf=True)
|
|
p_leaf_pandas = booster.predict(pandas_data, pred_leaf=True)
|
|
np_assert_array_equal(p_leaf_polars, p_leaf_pandas, strict=True)
|
|
|
|
p_pred_contrib_polars = booster.predict(data, pred_contrib=True)
|
|
p_pred_contrib_pandas = booster.predict(pandas_data, pred_contrib=True)
|
|
np_assert_array_equal(p_pred_contrib_polars, p_pred_contrib_pandas, strict=True)
|
|
|
|
p_first_iter_polars = booster.predict(data, start_iteration=0, num_iteration=1, raw_score=True)
|
|
p_first_iter_pandas = booster.predict(pandas_data, start_iteration=0, num_iteration=1, raw_score=True)
|
|
np_assert_array_equal(p_first_iter_polars, p_first_iter_pandas, strict=True)
|
|
|
|
|
|
def test_predict_regression():
|
|
data_float = generate_random_polars_frame(10, 10000, 42)
|
|
data_bool = generate_random_polars_frame(1, 10000, 42, generate_nulls=False, values=np.array([True, False]))
|
|
data = data_float.with_columns(data_bool["col_0"].alias("col_bool"))
|
|
|
|
dataset = lgb.Dataset(
|
|
data,
|
|
label=generate_random_polars_series(10000, 43, generate_nulls=False),
|
|
params=dummy_dataset_params(),
|
|
)
|
|
booster = lgb.train(
|
|
{"objective": "regression", "num_leaves": 7},
|
|
dataset,
|
|
num_boost_round=5,
|
|
)
|
|
assert_equal_predict_polars_pandas(booster, data)
|
|
|
|
|
|
def test_predict_binary_classification():
|
|
data = generate_random_polars_frame(10, 10000, 42)
|
|
dataset = lgb.Dataset(
|
|
data,
|
|
label=generate_random_polars_series(10000, 43, generate_nulls=False, values=np.arange(2)),
|
|
params=dummy_dataset_params(),
|
|
)
|
|
booster = lgb.train(
|
|
{"objective": "binary", "num_leaves": 7},
|
|
dataset,
|
|
num_boost_round=5,
|
|
)
|
|
assert_equal_predict_polars_pandas(booster, data)
|
|
|
|
|
|
def test_predict_multiclass_classification():
|
|
data = generate_random_polars_frame(10, 10000, 42)
|
|
dataset = lgb.Dataset(
|
|
data,
|
|
label=generate_random_polars_series(10000, 43, generate_nulls=False, values=np.arange(5)),
|
|
params=dummy_dataset_params(),
|
|
)
|
|
booster = lgb.train(
|
|
{"objective": "multiclass", "num_leaves": 7, "num_class": 5},
|
|
dataset,
|
|
num_boost_round=5,
|
|
)
|
|
assert_equal_predict_polars_pandas(booster, data)
|
|
|
|
|
|
def test_predict_ranking():
|
|
data = generate_random_polars_frame(10, 10000, 42)
|
|
dataset = lgb.Dataset(
|
|
data,
|
|
label=generate_random_polars_series(10000, 43, generate_nulls=False, values=np.arange(4)),
|
|
group=np.array([1000, 2000, 3000, 4000]),
|
|
params=dummy_dataset_params(),
|
|
)
|
|
booster = lgb.train(
|
|
{"objective": "lambdarank", "num_leaves": 7},
|
|
dataset,
|
|
num_boost_round=5,
|
|
)
|
|
assert_equal_predict_polars_pandas(booster, data)
|
|
|
|
|
|
def test_polars_feature_name_auto():
|
|
data = generate_dummy_polars_frame()
|
|
dataset = lgb.Dataset(
|
|
data,
|
|
label=pl.Series("label", [0, 1, 0, 0, 1]),
|
|
params=dummy_dataset_params(),
|
|
categorical_feature=["a"],
|
|
)
|
|
booster = lgb.train({"num_leaves": 7}, dataset, num_boost_round=5)
|
|
assert booster.feature_name() == ["a", "b"]
|
|
|
|
|
|
def test_polars_feature_name_manual():
|
|
data = generate_dummy_polars_frame()
|
|
dataset = lgb.Dataset(
|
|
data,
|
|
label=pl.Series("label", [0, 1, 0, 0, 1]),
|
|
params=dummy_dataset_params(),
|
|
feature_name=["c", "d"],
|
|
categorical_feature=["c"],
|
|
)
|
|
booster = lgb.train({"num_leaves": 7}, dataset, num_boost_round=5)
|
|
assert booster.feature_name() == ["c", "d"]
|
|
|
|
|
|
def test_get_data_polars_frame():
|
|
from polars.testing import assert_frame_equal # noqa: PLC0415
|
|
|
|
original_frame = generate_simple_polars_frame()
|
|
dataset = lgb.Dataset(original_frame, free_raw_data=False)
|
|
dataset.construct()
|
|
|
|
returned_data = dataset.get_data()
|
|
assert isinstance(returned_data, pl.DataFrame)
|
|
assert returned_data.schema == original_frame.schema
|
|
assert returned_data.shape == original_frame.shape
|
|
assert_frame_equal(returned_data, original_frame)
|
|
|
|
|
|
def test_get_data_polars_frame_subset(rng):
|
|
from polars.testing import assert_frame_equal # noqa: PLC0415
|
|
|
|
original_frame = generate_random_polars_frame(num_columns=3, num_datapoints=1000, seed=42)
|
|
dataset = lgb.Dataset(original_frame, free_raw_data=False)
|
|
dataset.construct()
|
|
|
|
subset_size = 100
|
|
used_indices = rng.choice(a=original_frame.shape[0], size=subset_size, replace=False)
|
|
used_indices = sorted(used_indices)
|
|
|
|
subset_dataset = dataset.subset(used_indices).construct()
|
|
expected_subset = original_frame[used_indices]
|
|
subset_data = subset_dataset.get_data()
|
|
|
|
assert isinstance(subset_data, pl.DataFrame)
|
|
assert subset_data.schema == expected_subset.schema
|
|
assert subset_data.shape == expected_subset.shape
|
|
assert len(subset_data) == len(used_indices)
|
|
assert subset_data.shape == (subset_size, 3)
|
|
assert_frame_equal(subset_data, expected_subset)
|